import numpy as np
import xarray as xr
import metpy.calc as mpcalc


def smthClmDay(clmDay, nHarm):  # from athul
    # nharm is the number of harmonic modes that needs to be removed. So we retain that in smthClm to later 
    # subtract from raw.
    print(clmDay)

    cf = np.fft.rfft(clmDay, axis=0)  # ([nday/2]+1)x[nlat]x[mlon]<==n:even
    # ([nday+1/2])x[nlat]x[mlon]<==n:odd
    # cf(:,0:nharm-1)
    cf[nHarm, :, :] = 0.5 * cf[nHarm, :, :]  # mini-taper
    cf[nHarm + 1:, :, :] = 0.0  # set all higher coef. to zero

    clmDaySmth = np.fft.irfft(cf, axis=0)  # reverse FT to real reconstruct data

    clmDaySmth = xr.DataArray(clmDaySmth, dims=('hourofyear', 'lat', 'lon'),
                              coords={'hourofyear': clmDay.hourofyear, 'lat': clmDay.lat, 'lon': clmDay.lon})
    return (clmDaySmth)

# also weird dims. le_ti_level_0
def compute_climatology(cfg):
    strt = '1979-01-01'
    last = '2019-12-31'

    cv_database_dir = cfg.cv_data_dir

    f = xr.open_mfdataset(cv_database_dir + "*.nc").sel(
        # plev=slice(plev, plev), 
        # lat=slice(latS,latN), 
        # lon=slice(lonW,lonE),
        time=slice(strt, last))
    # smoothing
    cv = mpcalc.smooth_n_point(f.cv, n=9, passes=2).metpy.dequantify()
    # generate hourofyear climatology, save it to file for next time

    # Calculate a hour-of-year climatology to be subtracted later.
    cv = cv.assign_coords(hourofyear=cv.time.dt.strftime("%m-%d %H"))
    print('Computing climatology...')
    cv_clim_da = cv.groupby('hourofyear').mean('time')

    # print('Smoothing climatology...')
    # clm = smthClmDay(clm, 3)  # calculate climatology with just the 3 harmonic modes retained.
    print('Climatology computed ')

    cv_clim_da.attrs['units'] = 's**-1'
    cv_clim_da.attrs['standard_name'] = 'Curvature Vorticity'
    cv_clim_da.attrs['long_name'] = 'Curvature Vorticity climatology (hour of year)'

    cv_clim = cv_clim_da.to_dataset(name="cv")

    # compute anomalies
    print("Compute anomalies of " + str(cfg.cv_percentile) + " percentile (total)...")
    full_cv_anomalies = cv.groupby('hourofyear') - cv_clim.cv

    # quantile applied over time,lat,lon -> so one scalar of cfg.cv_percentile% quantile per pressure level
    quantiles_full = full_cv_anomalies.quantile(float(cfg.cv_percentile) / 100.0, dim=['lat', 'lon', 'time'])
    cv_clim['cva_quantile_full'] = quantiles_full
    print(quantiles_full.data)

    print("Compute anomalies of " + str(cfg.cv_percentile) + " percentile (hourofyear)...")
    # quantile for each hourofyear, bit more complicated. bascially for each Jan 1 00UTC the x% quantile etc.
    full_cv_anomalies = full_cv_anomalies.assign_coords(hourofyear=full_cv_anomalies.time.dt.strftime("%m-%d %H"))

    full_cv_anomalies = full_cv_anomalies.swap_dims({"time": "hourofyear"})
    full_cv_anomalies_stacked = full_cv_anomalies.stack(le_ti=("plev", "hourofyear"))

    # group by pressure level and hourofyear -> quantile over lat+lon+date(hoy)
    quantiles_hoy_stacked = full_cv_anomalies_stacked.groupby('le_ti').quantile(q=float(cfg.cv_percentile) / 100.0,
                                                                                dim=...)
    quantiles_hoy = quantiles_hoy_stacked.unstack(dim='le_ti')

    cv_clim['cva_quantile_hoy'] = quantiles_hoy
    quantiles_hoy = cv_clim.cva_quantile_hoy
    # somehow xarray cant reconstruct stacked groupby results, fix manually.
    quantiles_hoy = quantiles_hoy.swap_dims({"le_ti_level_0": "plev"})
    quantiles_hoy = quantiles_hoy.swap_dims({"le_ti_level_1": "hourofyear"})
    quantiles_hoy = quantiles_hoy.rename({"le_ti_level_0": "plev"})
    quantiles_hoy = quantiles_hoy.rename({"le_ti_level_1": "hourofyear"})
    print(quantiles_hoy)

    cv_clim['cva_quantile_hoy'] = quantiles_hoy
    cv_clim = cv_clim.drop_dims(["le_ti_level_0", "le_ti_level_1"])
    print("Done!")
    print("Generated dataset:")
    print(cv_clim)

    return cv_clim