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

from typing import Literal


def common_preprocess(ds):
    # Remove unnecessary variables
    for var in ['plev', 'time_bnds']:
        ds = ds.drop_vars(var, errors='ignore')
    
    # Rename 'prlr' to 'pr' if present
    if "prlr" in ds:
        ds = ds.rename({'prlr': 'pr'})

    # Unit correction
    unit_correction = {
        "tas": ("degC", 273.15),
        "tasmin": ("degC", 273.15),
        "tasmax": ("degC", 273.15),
        "pr": ("kg m-2 s-1", 0)
    }

    for key, (unit, offset) in unit_correction.items():
        if key in ds and ds[key].attrs.get('units') != unit:
            ds[key] = ds[key] - offset if offset else ds[key]
            ds[key].attrs['units'] = unit

    # Set lat and lon units
    for coord in ['lat', 'lon']:
        ds[coord].attrs['units'] = 'degree'

    return ds

def fix_coords(ds):
    # Rename coordinates if necessary
    rename_dict = {'longitude': 'lon', 'latitude': 'lat', 'X': 'lon', 'Y': 'lat'}
    ds = ds.rename({old: new for old, new in rename_dict.items() if old in ds})

    # Adjust longitudes and sort coordinates
    ds = (ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
            .sortby(['lon', 'lat']))

    # Round times to start of the day
    try:
        ds = ds.assign_coords(time=ds.time.dt.floor("D"))
    except AttributeError:
        pass

    return ds

def to_monthly(ds: xr.Dataset):

    # assign new coords in-place
    ds.coords["year"] = ("time", ds.time.dt.year.data)
    ds.coords["month"] = ("time", ds.time.dt.month.data)
    # reshape the array to (..., "month", "year") in-place
    return ds.set_index(time=("year", "month")).unstack("time", fill_value=np.nan)

def split_months_by_year(
    ds: xr.Dataset,
    split_months: list,
    time_dim: str = "time",
    split_dim: str = "month",
    year_dim: str = "year",
    fill_value: float = np.nan,
    year_reference: Literal["first", "last"] = "first"
) -> xr.Dataset:
    """
    Splits a dataset with monthly time dimension into a dataset with year and split dimensions.
    Handles year shifts for split months spanning across years (e.g., [12, 1, 2]).
    
    Parameters:
    -----------
    ...
    year_reference : {"first", "last"}, optional
        Determines whether the year is chosen based on the first or last month in split_months.
        Default is "first".
    """
    ds = ds.copy()
    time = pd.to_datetime(ds[time_dim].values)
    
    reference_month = split_months[0] if year_reference == "first" else split_months[-1]
    
    if year_reference == "first":
        split_year = time.year - (time.month < reference_month).astype(int)
    else:
        split_year = time.year + (time.month > reference_month).astype(int)
    
    ds = ds.assign_coords({
        year_dim: (time_dim, split_year),
        split_dim: (time_dim, time.month)
    })
    
    return (ds.set_index({time_dim: (year_dim, split_dim)})
              .unstack(time_dim, fill_value=fill_value)
              .sel({split_dim: split_months}))