import xarray as xr
import pandas as pd
import numpy as np
from typing import Literal

_MONTH_INT_2_ACRONYM = {
    1: "J", 2: "F", 3: "M", 4: "A", 5: "M", 6: "J",
    7: "J", 8: "A", 9: "S", 10: "O", 11: "N", 12: "D"
}

def get_season_str(month, nmonths):
    """
     create coord with the aggregated months
     i.e nmonths = 3 -> 3: MAM, 6: JJA, 9: SON, 12: DJF
    """
    return ''.join(
        _MONTH_INT_2_ACRONYM[(month + i - 1) % 12 + 1] 
        for i in range(nmonths)
    )
    

## TODO: masked vars return NAs after droping time nas
def monthly2seasonal(
    ds: xr.Dataset, 
    nmonths: int = 3,
    month_reference: Literal["first", "last"] = "first"
    ) -> xr.Dataset:
    
    # Sort by time and create a complete monthly time series
    ds = ds.sortby("time")
    start, end = ds.time.to_index()[[0, -1]]
    new_time_index = pd.date_range(start=f"{start.year}-01-01", 
                                    end=f"{end.year}-12-31", 
                                    freq="MS")
    
    # reindex with the monthly all-year around time series and compute means
    # shifts the time to be the first month of the season
    # drops the NAs at the end/start of the time series
    ds = (ds.reindex(time=new_time_index)
            .rolling(time=nmonths).mean().shift(time= -nmonths+1 if month_reference == "first" else 0)
            .sel(time=slice(start, end))
            .dropna(dim="time")
            )
    
    ds = ds.assign_coords(
        season=("time", np.vectorize(get_season_str)(ds.time.dt.month.values, nmonths))
        )
    ds.attrs["frequency"] = f"{nmonths}M"
    
    return ds