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

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 monthly2seasonal(ds: xr.Dataset, nmonths: int = 3) -> 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, compute seasonal means, and trim to original time range
    return (ds.reindex(time=new_time_index)
            .rolling(time=nmonths, center=True).mean()
            .sel(time=slice(start, end))
            .dropna(dim="time"))

#def common_preprocess(ds):
#    try:
#        ds = ds.drop('plev')
#    except:
#        pass
#    try:
#        ds = ds.drop('time_bnds')
#    except:
#        pass

#    if 'prlr' in ds:
#        ds = ds.rename({'prlr': 'pr'})

#    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:
#            if ds[key].attrs.get('units') != unit:
#                ds[key] = ds[key] - offset if offset else ds[key]
#                ds[key].attrs['units'] = unit

#    ds.lat.attrs['units'] = 'degree'
#    ds.lon.attrs['units'] = 'degree'

#    return ds
#def fix_coords(dataset):
#    """ Function that sets longitudes in the -180,180 range and sorts
#        latitudes.

#        dataset: xarray object with 2D coords dims (lat,lon) or (latitude,longitude)
#        return: arranged dataset
#    """
#    # Changes the name of the lon lat coords if necessary
#    try:
#        dataset = dataset.rename({'longitude':'lon','latitude':'lat'})
#    except:
#        #log = "no need to rename lon lat dims"
#        pass
#    try:
#        dataset = dataset.rename({'X':'lon','Y':'lat'})
#    except:
#        #log = "no need to rename lon lat dims"
#        pass

#    # Converts longitudes to 0 360 domain
#    dataset = dataset.assign_coords(lon=(((dataset.lon + 180) % 360) - 180))
#    dataset = dataset.sortby('lon')
#    dataset = dataset.sortby('lat')

#    #dataset = dataset.assign_coords(month=("time", dataset.time.dt.month.values), 
#    #                                year=("time", dataset.time.dt.year.values))

#    # TODO: only works from xarray version 2022.09 onwards
#    #dataset = dataset.set_xindex("month")
#    #dataset = dataset.set_xindex("year")
#    # Round times to start of the day
#    try:
#        dataset = dataset.assign_coords(time=dataset.time.dt.floor("D"))
#    except:
#        #log = "no need to reset hour un leadtimes"
#        pass
#    return dataset