import collections.abc
import os
import numpy as np
import pandas as pd
import datetime as dt
import xarray as xr
import pyproj
import parse
from toolz import merge_with,cons
from scipy.optimize import curve_fit

import trosat.sunpos as sp
import pyrnet.logger as pyrlogger
import pyrnet.data as pyrdata
import pyrnet.pyrnet as pyrnet
import pyrnet.utils as pyrutils
pyrconfig = pyrdata.get_config()

def lennard_jones_exp6(x,e,s,a=0.2):
    xm = -s*10 # x at maximum potential peak 
    xc = x-xm # center at maximum potential peak
    f =  1 - 1e-2*(e/(a-6))*( 6*np.exp(a*(1-(xc/(10*s)))) - a*(10*s/xc)**6 )
    f = f/(1+e*1e-2) # normalize at peak value
    return f

def fit_trans(ds,values,label='',plot=False):
    # find Fmax somewhere in the initial -8 to 20m
    dstart,dend = -8,20
    idxs = np.searchsorted(ds.dist.values,dstart)
    idxe = np.searchsorted(ds.dist.values,dend)
    Fmax = np.nanmax(values[idxs:idxe])

    # define index for Fmax/e  Fmax/e interval
    idx0 = np.searchsorted(ds.dist.values,0)

    offset = np.argmax(values[idxs:idxe])+idxs-idx0
    dval = np.diff(values[idx0-offset:np.searchsorted(ds.dist.values,150)-offset])
    dval = np.convolve(dval,np.ones(100),'same')

    idxe = idx0 + np.argmin(dval)
    idx = np.argwhere(values<values[idxe]).ravel()
    idxs = int(np.max(idx[idx<idx0]))
    
    # center at Fmax
    #offset = np.argmax(values[idxs:idxe])+idxs-idx0
    ndist = ds.dist.values[idxs-offset:idxe-offset]
    nvalues = values[idxs:idxe]/Fmax
    nvalues = np.roll(nvalues,-offset)
    values = np.roll(values,-offset)
    
    # constrain to 1 after 1km (weight 0.5 versus observed data)
    lvalues = np.full(int(0.1*(idxe-idxs)),1./Fmax)
    ldist = ds.dist.values[idxs:idxs+len(lvalues)] - dstart + 10000

    # prepare fitting data
    fdist = np.concatenate((ndist,ldist),axis=0)
    fvalues = np.concatenate((nvalues,lvalues),axis=0)

    # get optimized epsilon and sigma
    fitfunc = lambda x,y,z: lennard_jones_exp6(x,y,z,a=0.2)
    try:
        popt,pcov = curve_fit(fitfunc,fdist,fvalues,p0=(10,10),bounds=((1e-6,1e-6),(1e3,1e3)))
    except:
        return np.full(4,np.nan)
    epsilon,sigma = popt

    # calculate x_e from sigma and fixed alpha
    alpha=0.2
    xe = -10*(sigma/alpha)*np.log((6-alpha)/(6*np.exp(1)))
    
    # plot data and fit
    idxxe = np.searchsorted(ds.dist.values,xe)
    if idxxe>=ds.dist.size:
        idxxe=ds.dist.size-1
    if plot:
        print()
        print(label,':')
        print(f"epsilon={epsilon:.2f}; sigma={sigma:.2f}; alpha={alpha:.3f}",f"x_e={xe:.1f}m; f(x_e)={lfc(xe,*popt)*Fmax:.3f}, values(x_e)={values[idxxe]:.3f}")
        fig,ax = plt.subplots(1,1,figsize=(10,5))
        ax.set_title(label)
        ax.plot(ds.dist,values,c='C0', ls=':')
        ax.plot(ndist,values[idxs-offset:idxe-offset],c='C0',label='values')
        ax.plot(ds.dist,lennard_jones_exp6(sig.dist,*popt)*Fmax,c='C1',ls=':')
        ax.plot(ds.dist[idxs-offset:idxe-offset],lennard_jones_exp6(sig.dist[idxs-offset:idxe-offset],*popt)*Fmax,c='C1',label='fitted potential')
        ax.axline((xe,0),(xe,1),c='k',ls='--',label='e-folding')
        ax.grid(True)
        ax.set_ylim((0.9,1.1))
        ax.set_ylim((1,np.max(values[idxs:idxe])+0.05))
        ax.legend(loc='upper right')
    
    return epsilon,sigma,alpha,xe

def get_sfclatlon(szen ,sazi ,lat ,lon ,heights ,GEOD=None):
    """
    Calculate lat/lon coordiantes on surface projected from a point with certain height along sun beam.

    Parameters:
    -----------
    szen: float
        solar zenith angle in degree
    sazi: float
        solar azimuth angle in degree
    lat: float or array of floats
        latitude coordinates at height
    lon: float or array of float
        longitude coordinates at height
    heights: float or array of float
        geometric height

    Returns:
    --------
    sfc_lats, sfc_lon: float or array of float same size as lat/lon
       projected coordiantes on surface
    """
    if GEOD is None:
        GEOD = pyproj.Geod(ellps='sphere')
    sfc_dradius = np.tan(np.deg2rad(szen)) * heights
    sfc_dazi = sazi + 180.
    lons ,lats = np.meshgrid(lon ,lat)
    sfc_lons = np.full((heights.size, lat.size, lon.size) ,np.nan)
    sfc_lats = sfc_lons.copy()

    for i in range(heights.size):
        sfc_lon ,sfc_lat ,_ = GEOD.fwd(
            lons=lons.flatten(),
            lats=lats.flatten(),
            az=np.full(lats.size ,sfc_dazi),
            dist=np.full(lats.size ,sfc_dradius[i]),
            radians=False
        )
        sfc_lons[i ,: ,:] = sfc_lon.reshape((lat.size ,lon.size))
        sfc_lats[i ,: ,:] = sfc_lat.reshape((lat.size ,lon.size))
    return sfc_lats, sfc_lons

def load_l1a(date,station,path):
    date=pd.to_datetime(date)
    pyrconfig = pyrdata.get_config()
    fname_format = pyrconfig["output"]

    fnames = np.sort(os.listdir(path))
    results = False
    for fname in fnames[::-1]:
        res = parse.parse(fname_format, fname)
        if not results:
            results = {k:[v] for k,v in res.named.items()}
            results.update({"fname":[fname]})
        else:
            res = {k:[v] for k,v in res.named.items()}
            res.update({"fname":[fname]})
            results = merge_with(lambda x: x[1]+x[0], results, res)
    c = pd.DataFrame.from_dict(results)
    c = c.query(f'station=={station}').reset_index()

    startdts = pd.to_datetime(c["dt"])
    enddts = startdts + pd.to_timedelta(c["period"])
    # get file index with maintenance interval including date
    mask = date>=startdts
    mask *= date<enddts

    ds_l1b = None
    for i,fname in enumerate(c["fname"][mask.values]):
        print(fname)
        ds = pyrdata.to_l1b(
            os.path.join(path,fname),
            config=dict(l1bfreq="100ms")
        )
        if ds is None:
            continue
        ds = ds.drop_vars([var for var in ds if var not in ["ghi","esd","szen"]])
        udays = np.unique(ds.time.values.astype("datetime64[D]"))
        dslist= []
        for day in udays:
            day = pd.to_datetime(day)
            dslist.append( ds.sel(time=f"{day:%Y-%m-%d}") )

        # if ds_l1b is None:
        #     ds_l1b = pyrdata.merge_l1b(dslist,freq='100ms')
        # else:
        #     ds_l1b = pyrdata.merge_l1b([ds_l1b]+dslist,freq='100ms')

        if ds_l1b is not None:
            dslist = [ds_l1b] + dslist

        freq = '100ms'
        timevar = 'time'
        # sort by first station coordinate
        dslist = pyrdata._sort_by_station(dslist)
        ## Unify datasets
        # reindex timevar:
        dslist = pyrdata._reindex_time(dslist, freq=freq, timevar=timevar)
        # reindex station var:
        dslist = pyrdata._reindex_station(dslist)
        # reindex maintenancetime var:
        dslist = pyrdata._reindex_maintenancetime(dslist)
        
        #####################################################################
        ## Merge datasets
        # merge vars with (time,station) dims
        for i in range(len(dslist)):
            dst = dslist[i].copy()
            dst = dst.drop_vars(
                [var for var in dst if not (timevar in dst[var].dims and "station" in dst[var].dims)]
            )
            if i==0:
                ds_time_station = dst.copy()
            else:
                # handle overlapping values by dropping from the first (override from second)
                for var in dst:
                    overlap = (~np.isnan(dst[var].values))*(~np.isnan(dst[var].values))
                    ds_time_station[var].values = ds_time_station[var].values.astype(float)
                    ds_time_station[var].values[overlap] = np.nan
                ds_time_station = ds_time_station.merge(dst)
                
        # merge vars with (station) dims
        for i in range(len(dslist)):
            dst = dslist[i].copy()
            dst = dst.drop_vars(
                [var for var in dst if not (len(dst[var].dims)==1 and "station" in dst[var].dims)]
            )
            if i==0:
                ds_station = dst.copy()
            else:
                try:
                    # works there is no overlap with non null values ( new stations )
                    ds_station = ds_station.merge(dst, compat='no_conflicts')
                except:
                    # override if station already exists
                    ds_station = ds_station.merge(dst, compat='override')
                #ds_station = xr.concat((ds_station, dst), dim='station')
        
        ds_l1b = xr.merge([ds_time_station,ds_station])
    

    if ds_l1b is None:
        return None,None
    
    tr0 = ds_l1b.ghi.sel(station=station)
    tr0 /= (1367.0*ds_l1b.esd.values**2*np.cos(np.deg2rad(ds_l1b.szen.values[:,0])))
    # filter night
    tr0 = tr0.where(ds_l1b.szen.values[:,0]<80)
                                
    return ds_l1b, tr0

def get_cloudshade(szen, sazi, cloudmask, lat, lon, heights, extents, bins=200, GEOD=None):
    """
    Project a 3D cloud mask (coordianes lat,lon,height) to surface along sun beam.

    Parameters:
    -----------
    szen: float
        solar zenith angle in degree
    sazi: float
        solar azimuth angle in degree
    cloud_mask: array of bool, shape(height.size,lat.size,lon.size)
        Boolean cloud mask
    lat: float or array of floats
        latitude coordinates at height
    lon: float or array of float
        longitude coordinates at height
    heights: float or array of float
        geometric height
    extents: array of float with length 4
        the extends of the projection area (longitude min, longitude max, latitude min, latitude max)
    bins: int, or array of int with lenght 2
        number of bins of the projection area, if scalar the number of bins in both directions is the same,
        else (longitude bins, latitude bins)

    Returns:
    --------
    hshades: array of bools of size (bins[0], bins[1])
        projected shade mask on the surface
    hlat, hlon: array of floats
        the coordinates of hshades
    """
    if GEOD is None:
        GEOD = pyproj.Geod(ellps='sphere')
    if not isinstance(bins,collections.abc.Sequence):
        bins = (int(bins),int(bins))
    bins = np.array(bins).astype(int)
    assert len(bins)==2


    # calculate surface lat lon
    sfc_lats, sfc_lons = get_sfclatlon(
        szen=szen,
        sazi=sazi,
        lat=lat,
        lon=lon,
        heights=heights,
        GEOD=GEOD
    )

    lonbins = np.linspace(extents[0], extents[1], bins[0])
    latbins = np.linspace(extents[2], extents[3], bins[1])

    hlat = latbins[:-1] + np.diff(latbins)
    hlon = lonbins[:-1] + np.diff(lonbins)

    hshade = np.zeros((heights.size, latbins.size - 1, lonbins.size - 1))

    for i in range(heights.size):
        mask = cloudmask[i, :, :].flatten()
        mask = mask.astype(bool)
        X = sfc_lats[i, :, :].flatten()[mask]
        Y = sfc_lons[i, :, :].flatten()[mask]

        hist, _, _ = np.histogram2d(X, Y, bins=[latbins, lonbins])
        hist = hist.astype(bool)

        hshade[i, :, :] = hist

    hshades = np.sum(hshade, axis=0).astype(bool)
    return hshades, hlat, hlon