
import tinysunset as ts
import xarray as xr
import pandas as pd
import numpy as np
    
    
def preprocess_era5(ds):
    
    ds = ts.loaders.common_preprocess(ds)
        
    # floor time values to the nearest hour
    ds['time'] = ds['time'].dt.floor('H')
    
    ds['time'] = pd.to_datetime(ds['time'].dt.strftime('%Y-%m-01'))
    
    return ds


def load_esarchive_era5(
    vars, months, years,
    ref_grid: xr.Dataset = None,
    seasonal_mean: bool = False
    ):

    # --------------------------
    #       Hcst files load
    # --------------------------

    _ARCHIVE = "/esarchive/recon/ecmwf/era5/monthly_mean/"
    _var_prefix = {
        "tas": "_f1h-r1440x721cds/", "prlr": "_f1h-r1440x721cds/", 
        "g500": "_f1h-r1440x721cds/", "tasmin": "_f1h-r1440x721cds/", 
        "tasmax": "_f1h-r1440x721cds/"
        }

    # Compute lead_months and format dates as strings
    # example: sd_month = 12, target_months = [1,2,3] -> lead_months = [1,2,3]
    _months = [str(m).zfill(2) for m in months]
    
    exps = []
    for var in vars:

        path = _ARCHIVE + var + _var_prefix[var]
        
        year_months = [f"{year}{month}" for year in years for month in _months]
        
        files = [path + f"{var}_{year_month}.nc" for year_month in year_months]
        
        exp = xr.open_mfdataset(
            files,
            drop_variables=['time_bnds'],
            combine='nested',
            combine_attrs='override',
            concat_dim='time',
            join="override",
            preprocess=preprocess_era5
            )
        
        exps.append(exp)

    exps = xr.merge(exps)

    if seasonal_mean:
        exps = exps.rolling(time=3, center=False).mean()
    
    exps = ts.loaders.fix_coords(exps).load()
    exps = ts.loaders.to_monthly(exps)
    
    if ref_grid is not None:
        exps = exps.interp_like(ref_grid, method="linear")

    # combine all the xarray objects into one
    return exps

def load_esarchive_era5_v2(
    vars, months, years,
    ref_grid: xr.Dataset = None,
    seasonal_mean: bool = False
    ):

    _ARCHIVE = "/esarchive/recon/ecmwf/era5/monthly_mean/"
    _var_prefix = {
        "tas": "_f1h-r1440x721cds/", "prlr": "_f1h-r1440x721cds/", 
        "g500": "_f1h-r1440x721cds/", "tasmin": "_f1h-r1440x721cds/", 
        "tasmax": "_f1h-r1440x721cds/"
        }

    # Compute lead_months and format dates as strings
    # example: sd_month = 12, target_months = [1,2,3] -> lead_months = [1,2,3]
    _months = [str(m).zfill(2) for m in months]

    exps = []
    for var in vars:

        path = _ARCHIVE + var + _var_prefix[var]
        
        year_months = [f"{year}{month}" for year in years for month in _months]
        
        files = [path + f"{var}_{year_month}.nc" for year_month in year_months]
        
        exp = xr.open_mfdataset(
            files,
            drop_variables=['time_bnds'],
            combine='nested',
            combine_attrs='override',
            concat_dim='time',
            join="override",
            preprocess=preprocess_era5
            )
        
        exps.append(exp)
        
    exps = xr.merge(exps)

    exps = ts.loaders.fix_coords(exps).load()

    if ref_grid is not None:
        exps = exps.interp_like(ref_grid, method="linear")
        
    if seasonal_mean:
        exps = ts.tools.monthly2seasonal(exps, nmonths=len(months))
    

    exps = ts.loaders.to_monthly(exps)

    # combine all the xarray objects into one
    return exps
