
import numpy as np
import xarray as xr
import pandas as pd
import tinysunset as ts

def preprocess_seas5(ds, leadtimes: list, members: list):
    ds = ds.sel(dict(ensemble=members, time=leadtimes))
    
    try:
        ds = ds.drop("realization")
    except:
        pass

    # Rename longitude to lon and latitude to lat if needed
    if 'longitude' in ds.coords:
        ds = ds.rename({'longitude': 'lon'})
    if 'latitude' in ds.coords:
        ds = ds.rename({'latitude': 'lat'})
    
    ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
    ds = ds.sortby('lon')

    ds = ts.loaders.common_preprocess(ds)
    return ds

def load_seas5(
    vars, sd_month, target_months, hcst_years,
    ref_grid: xr.Dataset = None,
    seasonal_mean: bool = False,
    n_members = 25
    ):
    
    _ARCHIVE = "/esarchive/exp/ecmwf/system51c3s/monthly_mean/"
    _var_prefix = {"tas": "_f6h/", "prlr": "_s0-24h/", "g500": "_f12h/", "tasmin": "_f24h/", "tasmax": "_f24h/"}

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

        # ex. /esarhive/exp/ecmwf/system5c3s/monthly_mean/tas_f6h/tas_20190101.nc"
        exp_files = [
            _ARCHIVE + var + _var_prefix[var] + var + "_" + str(year) + sd_month + '01.nc' 
            for year in hcst_years
            ]
        
        print("Loading " + var + " files")
        exp = xr.open_mfdataset(exp_files,
                                drop_variables=['time_bnds','time'],
                                combine='nested',
                                concat_dim='syear',
                                # preprocess is preprocess setting leadtimes=lead_months
                                preprocess= lambda x: preprocess_seas5(
                                    x, leadtimes=lead_months, members=slice(n_members)
                                    )
                                )
        # Adds dates to the sdate dimension the month is not relevant
        exp['syear'] = pd.date_range(
            str(list(hcst_years)[0])+"-01-01", freq="YS", periods=len(hcst_years)
            )
        exps.append(exp)

    exps = xr.merge(exps)

    if seasonal_mean:
        exps = exps.mean(dim="time")
        exps = exps.expand_dims("time")
        target_months = [sd_month]

    # reindex syear and time into a single time dimension
    exps = exps.rename({"time":"tt"})
    exps = exps.stack(time=["syear","tt"])
    
    dates_list = []
    for year in hcst_years:
        for month in target_months:
            # create a NumPy datetime64 object for the first day of the month
            date = np.datetime64('{0}-{1:02}-01'.format(year, month))
            # add the datetime64 object to the list
            dates_list.append(date)

    exps = exps.assign_coords({"time": dates_list})
    
    exps = ts.loaders.fix_coords(exps)
    if ref_grid is not None:
        exps = exps.interp_like(ref_grid, method="linear")

    exps = exps.load()
    exps = ts.loaders.to_monthly(exps)
    
    return exps