"""Functions to fetch and process data from the ECMWF WeatherBench dataset."""
import numpy as np
import pandas as pd
import xarray as xr
from dateutil.relativedelta import relativedelta
from nuthatch import cache
from nuthatch.processors import timeseries

from sheerwater.reanalysis import era5_rolled
from sheerwater.utils import dask_remote, get_grid, get_variable, lon_base_change, regrid, roll_and_agg, shift_by_days
from sheerwater.interfaces import forecast as sheerwater_forecast, spatial


@dask_remote
@timeseries(timeseries=['start_date', 'model_issuance_date'])
@cache(cache=False,
       cache_args=['variable', 'forecast_type', 'run_type', 'time_group', 'grid'],
       backend_kwargs={'chunking': {"lat": 121, "lon": 240, "lead_time": 46,
                                    "start_date": 29, "start_year": 29,
                                    "model_issuance_date": 1}})
def ifs_extended_range_raw(start_time, end_time, variable, forecast_type,  # noqa ARG001
                           run_type='average', time_group='weekly', grid="global1_5"):
    """Fetches IFS extended range forecast data from the WeatherBench2 dataset.

    Args:
        start_time (str): The start date to fetch data for.
        end_time (str): The end date to fetch.
        variable (str): The weather variable to fetch.
        forecast_type (str): The type of forecast to fetch. One of "forecast" or "reforecast".
        run_type (str): The type of run to fetch. One of:
            - average: to download the averaged of the perturbed runs
            - perturbed: to download all perturbed runs
            - [int 0-50]: to download a specific  perturbed run
        time_group (str): The time grouping to use. One of: "daily", "weekly", "biweekly"
        grid (str): The grid resolution to fetch the data at. One of:
            - global1_5: 1.5 degree global grid
    """
    if grid != 'global1_5':
        raise NotImplementedError("Only global 1.5 degree grid is implemented.")

    forecast_str = "-reforecast" if forecast_type == "reforecast" else ""
    run_str = "_ens_mean" if run_type == "average" else ""
    avg_str = "-avg" if time_group == "daily" else "_avg"  # different naming for daily
    time_str = "" if time_group == "daily" else "-weekly" if time_group == "weekly" else "-biweekly"
    file_str = f'ifs-ext{forecast_str}-full-single-level{time_str}{avg_str}{run_str}.zarr'
    filepath = f'gs://weatherbench2/datasets/ifs_extended_range/{time_group}/{file_str}'

    # Pull the google dataset
    ds = xr.open_zarr(filepath, decode_timedelta=True)

    # Select the right variable
    var = get_variable(variable, 'ecmwf_ifs_er')
    ds = ds[var].to_dataset()
    ds = ds.rename_vars(name_dict={var: variable})

    # Convert local dataset naming and units
    ds = ds.rename({'latitude': 'lat', 'longitude': 'lon', 'prediction_timedelta': 'lead_time'})
    if run_type != 'average':
        ds = ds.rename({'number': 'member'})
    if forecast_type == 'reforecast':
        ds = ds.rename({'hindcast_year': 'start_year'})
        ds = ds.rename({'forecast_time': 'model_issuance_date'})
        ds = ds.drop('time')
    else:
        ds = ds.rename({'time': 'start_date'})

    ds = ds.drop_vars('valid_time')

    # If a specific run, select
    if isinstance(run_type, int):
        ds = ds.sel(member=run_type)
    return ds


@dask_remote
@timeseries(timeseries=['start_date', 'model_issuance_date'])
@spatial()
@cache(cache_args=['variable', 'forecast_type', 'run_type', 'time_group', 'grid'],
       cache_disable_if={'grid': 'global1_5'},
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1,
                        "start_date": 1000,
                        "model_issuance_date": 1000, "start_year": 1,
                        "member": 1},
           'chunk_by_arg': {
               'grid': {
                   # A note: a setting where time is in groups of 200 works better for regridding tasks,
                   # but is less good for storage.
                   'global0_25': {"lat": 721, "lon": 1440, 'model_issuance_date': 30, "start_date": 30}
               },
           }
})
def ifs_extended_range(start_time, end_time, variable, forecast_type,
                       run_type='average', time_group='weekly',
                       grid="global1_5", mask=None, region='global'):  # noqa: ARG001
    """Fetches IFS extended range forecast and reforecast data from the WeatherBench2 dataset.

    Args:
        start_time (str): The start date to fetch data for.
        end_time (str): The end date to fetch.
        variable (str): The weather variable to fetch.
        forecast_type (str): The type of forecast to fetch. One of "forecast" or "reforecast".
        run_type (str): The type of run to fetch. One of:
            - average: to download the averaged of the perturbed runs
            - perturbed: to download all perturbed runs
            - [int 0-50]: to download a specific  perturbed run
        time_group (str): The time grouping to use. One of: "daily", "weekly", "biweekly"
        grid (str): The grid resolution to fetch the data at. One of:
            - global1_5: 1.5 degree global grid
        mask: Spatial mask to apply.
        region: Region to fetch data for.
    """
    """IRI ECMWF average forecast with regridding."""
    ds = ifs_extended_range_raw(start_time, end_time, variable, forecast_type,
                                run_type, time_group=time_group, grid='global1_5')
    # Convert to base180 longitude
    ds = lon_base_change(ds, to_base="base180")

    if variable in ['tmp2m', 'tmax2m', 'tmin2m']:
        ds[variable] = ds[variable] - 273.15
        ds.attrs.update(units='C')
    elif variable == 'precip':
        ds[variable] = ds[variable] * 1000.0
        ds.attrs.update(units='mm')
        ds = np.maximum(ds, 0)
    elif variable == 'ssrd':
        ds.attrs.update(units='Joules/m^2')
        ds = np.maximum(ds, 0)
    if grid == 'global1_5':
        return ds

    _, _, size, _ = get_grid(grid)
    if size < 1.5:
        raise NotImplementedError("Unable to regrid ECMWF smaller than 1.5x1.5")
    if forecast_type == 'reforecast':
        raise NotImplementedError("Regridding reforecast data should be done with extreme care. It's big.")

    # Need all lats / lons in a single chunk to be reasonable
    chunks = {'lat': 121, 'lon': 240, 'lead_time': 1}
    if run_type == 'perturbed':
        chunks['member'] = 1
        chunks['start_date'] = 200
    else:
        chunks['start_date'] = 200
    ds = ds.chunk(chunks)
    # Need all lats / lons in a single chunk for the output to be reasonable
    ds = regrid(ds, grid, base='base180', method='conservative',
                output_chunks={"lat": 721, "lon": 1440}, region=region)
    return ds


@dask_remote
@timeseries(timeseries='model_issuance_date')
@spatial()
@cache(cache_args=['variable', 'lead', 'run_type', 'time_group', 'grid'],
       backend_kwargs={'chunking': {"lat": 121, "lon": 240, "lead_time": 1, "model_issuance_date": 200, "member": 50}})
def ifs_er_reforecast_lead_bias(start_time, end_time, variable, lead=0, run_type='average',
                                time_group='weekly', grid="global1_5", mask=None, region='global'):
    """Computes the bias of ECMWF reforecasts for a specific lead."""
    # Fetch the reforecast data; get's the past 20 years associated with each start date
    ds_deb = ifs_extended_range(start_time, end_time, variable, forecast_type="reforecast",
                                run_type=run_type, time_group=time_group, grid=grid, mask=mask, region=region)

    # Get the appropriate lead
    n_leads = len(ds_deb.lead_time)
    if lead >= n_leads:
        # Lead does not exist
        return None
    ds_deb = ds_deb.sel(lead_time=np.timedelta64(lead, 'D'))

    first_date = ds_deb.model_issuance_date.min().values
    last_date = ds_deb.model_issuance_date.max().values
    # We need ERA5 data from the start_time to 20 years before the first date
    new_start = (pd.Timestamp(first_date) - relativedelta(years=20)).strftime("%Y-%m-%d")
    # We need ERA5 data from the end time to the end time plus last lead in days
    new_end = (pd.Timestamp(last_date + ds_deb.lead_time.values) - relativedelta(years=1)).strftime("%Y-%m-%d")

    # Get the pre-aggregated ERA5 data
    agg = {'daily': 1, 'weekly': 7, 'biweekly': 14}[time_group]
    ds_truth = era5_rolled(new_start, new_end, variable, agg_days=agg, grid=grid, mask=mask, region=region)

    def get_bias(ds_sub):
        """Get the 20-year estimated bias of the reforecast data."""
        # The the corresponding forecast dates for the reforecast data
        dates = [np.datetime64(pd.Timestamp(ds_sub['model_issuance_date'].values[0] + relativedelta(years=x)))
                 for x in ds_sub.start_year]

        # Adjust each forecast date by the lead time (0, 1, 2, ... days)
        lead_td = ds_deb.lead_time.values
        lead_dates = [x + lead_td for x in dates]

        # Select the subset of the ground truth matching this lead
        ds_truth_lead = ds_truth.sel(time=lead_dates)

        # # Assign the time coordinate to match the forecast dataframe for alignment and subtract
        ds_truth_lead = ds_truth_lead.assign_coords(time=ds_sub.start_year.values).rename(time='start_year')
        bias = (ds_truth_lead - ds_sub).mean(dim='start_year')
        return bias

    bias = ds_deb.groupby('model_issuance_date').map(get_bias)
    return bias


@dask_remote
@timeseries(timeseries='model_issuance_date')
@spatial()
@cache(cache_args=['variable', 'run_type', 'time_group', 'grid'],
       backend_kwargs={'chunking': {"lat": 121, "lon": 240, "lead_time": 1, "model_issuance_date": 1000, "member": 1}})
def ifs_er_reforecast_bias(start_time, end_time, variable, run_type='average',
                            time_group='weekly', grid="global1_5", mask=None,
                            region='global'):
    """Computes the bias of ECMWF reforecasts for all leads."""
    # Fetch the reforecast data to calculate how many leads we need
    if time_group == 'weekly':
        leads = [0, 7, 14, 21, 28, 35]
    elif time_group == 'biweekly':
        leads = [0, 7, 14, 21, 28]
    elif time_group == 'daily':
        leads = list(range(46))
    else:
        raise NotImplementedError(f"Time group {time_group} not implemented for ECMWF reforecasts.")

    # Accumulate all the per lead biases
    biases = []
    for i in leads:
        biases.append(ifs_er_reforecast_lead_bias(
            start_time, end_time, variable, lead=i,
            run_type=run_type, time_group=time_group, grid=grid,
            mask=mask, region=region))
    # Concatenate leads and unstack
    ds_biases = xr.concat(biases, dim='lead_time')
    return ds_biases


@dask_remote
@timeseries(timeseries='start_date')
@spatial()
@cache(cache_args=['variable', 'margin_in_days', 'run_type', 'time_group', 'grid'],
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1,
                        "start_date": 1000,
                        "model_issuance_date": 1000, "start_year": 1,
                        "member": 1},
           'chunk_by_arg': {
               'grid': {
                   # A note: a setting where time is in groups of 200 works better for regridding tasks,
                   # but is less good for storage.
                   'global0_25': {"lat": 721, "lon": 1440, 'model_issuance_date': 30, "start_date": 30}
               },
           }
})
def ifs_extended_range_debiased(start_time, end_time, variable, margin_in_days=6,
                                run_type='average', time_group='weekly', grid="global1_5", mask=None, region='global'):
    """Computes the debiased ECMWF forecasts."""
    # Get bias data from reforecast; for now, debias with deterministic bias
    ds_b = ifs_er_reforecast_bias(start_time, end_time, variable,
                                  run_type='average', time_group=time_group, grid=grid, mask=mask, region=region)

    # Get forecast data
    ds_f = ifs_extended_range(start_time, end_time, variable, forecast_type='forecast',
                              run_type=run_type, time_group=time_group, grid=grid, mask=mask, region=region)
    if time_group == 'weekly':
        leads = [np.timedelta64(x, 'D') for x in [0, 7, 14, 21, 28, 35]]
    elif time_group == 'biweekly':
        leads = [np.timedelta64(x, 'D') for x in [0, 7, 14, 21, 28]]
    elif time_group == 'daily':
        leads = [np.timedelta64(x, 'D') for x in range(46)]
    else:
        raise NotImplementedError(f"Time group {time_group} not implemented for ECMWF reforecasts.")
    ds_f = ds_f.sel(lead_time=leads)

    def bias_correct(ds_sub, mid=6):
        date = ds_sub.start_date.values
        dt = np.timedelta64(mid, 'D')
        nbhd = (ds_b.model_issuance_date.values >= date - dt) & \
            (ds_b.model_issuance_date.values <= date + dt)
        if nbhd.sum() == 0:  # No data to debias
            raise ValueError('No debiasing data found.')

        nbhd_df = ds_b.isel(model_issuance_date=nbhd).mean(dim='model_issuance_date')
        dsp = ds_sub + nbhd_df
        return dsp

    ds = ds_f.groupby('start_date').map(bias_correct, mid=margin_in_days)
    # Should not be below zero after bias correction
    if variable in ['precip', 'ssrd']:
        ds = np.maximum(ds, 0)
    return ds


@dask_remote
@timeseries(timeseries='start_date')
@spatial()
@cache(cache_args=['variable', 'margin_in_days', 'run_type', 'time_group', 'grid'],
       cache_disable_if={'grid': 'global1_5'},
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1,
                        "start_date": 1000,
                        "model_issuance_date": 1000, "start_year": 1,
                        "member": 1},
           'chunk_by_arg': {
               'grid': {
                   # A note: a setting where time is in groups of 200 works better for regridding tasks,
                   # but is less good for storage.
                   'global0_25': {"lat": 721, "lon": 1440, 'model_issuance_date': 30, "start_date": 30}
               },
           }
})
def ifs_extended_range_debiased_regrid(start_time, end_time, variable,
                                       margin_in_days=6, run_type='average',
                                       time_group='weekly', grid="global1_5",
                                       mask=None, region='global'):
    """Computes the debiased ECMWF forecasts."""
    ds = ifs_extended_range_debiased(start_time, end_time, variable,
                                     margin_in_days=margin_in_days,
                                     run_type=run_type, time_group=time_group,
                                     grid='global1_5', mask=mask, region=region)
    if grid == 'global1_5':
        return ds

    # Need all lats / lons in a single chunk to be reasonable
    chunks = {'lat': 121, 'lon': 240, 'lead_time': 1}
    if run_type == 'perturbed':
        chunks['member'] = 1
        chunks['start_date'] = 200
    else:
        chunks['start_date'] = 200
    ds = ds.chunk(chunks)
    # Need all lats / lons in a single chunk for the output to be reasonable
    ds = regrid(ds, grid, base='base180', method='conservative',
                output_chunks={"lat": 721, "lon": 1440}, region=region)
    return ds


@dask_remote
@timeseries(timeseries='start_date')
@spatial()
@cache(cache_args=['variable', 'prob_type', 'agg_days', 'grid', 'debiased'],
       cache_disable_if={'agg_days': [1, 7, 14]},
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1,
                        "start_date": 1000, "member": 1},
           'chunk_by_arg': {
               'grid': {
                   'global0_25': {"lat": 721, "lon": 1440, "start_date": 30}
               },
           }
})
def ifs_extended_range_rolled(start_time, end_time, variable,
                              prob_type='deterministic', agg_days=7,
                              grid="global1_5", debiased=True, mask=None, region='global'):
    """Standard format forecast data for aggregated ECMWF forecasts."""
    run_type = 'perturbed' if prob_type == 'probabilistic' else 'average'
    if debiased:
        fn = ifs_extended_range_debiased_regrid
        kwargs = {'margin_in_days': 6}
    else:
        fn = ifs_extended_range
        kwargs = {'forecast_type': 'forecast'}

    if agg_days not in [1, 7, 14]:  # not one of the precomputed time groups
        ds = fn(start_time, end_time, variable,
                run_type=run_type, time_group='daily', grid=grid, mask=mask, region=region, **kwargs)
        # Get aggregated variable
        ds[variable] = roll_and_agg(ds[variable], agg=agg_days, agg_col='lead_time', agg_fn='mean')
    else:
        time_group = {1: 'daily', 7: 'weekly', 14: 'biweekly'}[agg_days]
        ds = fn(start_time, end_time, variable,
                run_type=run_type, time_group=time_group, grid=grid, mask=mask, region=region, **kwargs)
    return ds


@dask_remote
def _ecmwf_ifs_er_unified(start_time, end_time, variable, agg_days, prob_type='deterministic',
                          grid="global1_5", mask='lsm', region="global", debiased=True):  # noqa: ARG001
    """Unified API accessor for ECMWF raw and debiased forecasts."""
    # The earliest and latest forecast dates for the set of all leads
    # ECMWF extended range forecasts have 46 days, so we shift to include all forecasters who could
    # overlaps with the start and end period
    forecast_start = shift_by_days(start_time, -46) if start_time is not None else None
    forecast_end = shift_by_days(end_time, 46) if end_time is not None else None

    ds = ifs_extended_range_rolled(forecast_start, forecast_end, variable, prob_type=prob_type,
                                   agg_days=agg_days, grid=grid, debiased=debiased, mask=mask, region=region)

    # Assign probability label
    prob_label = prob_type if prob_type == 'deterministic' else 'ensemble'
    ds = ds.assign_attrs(prob_type=prob_label)
    if 'spatial_ref' in ds.variables:
        ds = ds.drop_vars('spatial_ref')

    # TODO: remove this once we update ECMWF caches
    if variable == 'precip' and agg_days in [7, 14] and \
            not (grid == 'global1_5' and not debiased):
        # The dataframes for un-debiased, global1_5 are already divided by days
        print("Warning: Dividing precip by days to get daily values. Do you still want to do this?")
        ds['precip'] /= agg_days

    # Rename to standard naming
    ds = ds.rename({'start_date': 'init_time', 'lead_time': 'prediction_timedelta'})
    return ds


@dask_remote
@sheerwater_forecast()
@cache(cache=False,
       cache_args=['variable', 'agg_days', 'prob_type', 'grid', 'mask', 'region'],
       backend_kwargs={'chunking': {'lat': 300, 'lon': 300, 'time': 365, 'lead_time': 1, 'member': 1}})
def ecmwf_ifs_er(start_time=None, end_time=None, variable="precip", agg_days=1, prob_type='deterministic',
                 grid='global1_5', mask='lsm', region="global"):
    """Standard format forecast data for ECMWF forecasts."""
    return _ecmwf_ifs_er_unified(start_time=start_time, end_time=end_time, variable=variable,
                                 agg_days=agg_days, prob_type=prob_type,
                                 grid=grid, mask=mask, region=region, debiased=False)


@dask_remote
@sheerwater_forecast()
@cache(cache=False, cache_args=['variable', 'agg_days', 'prob_type', 'grid', 'mask', 'region'],
       backend_kwargs={'chunking': {'lat': 300, 'lon': 300, 'time': 365, 'lead_time': 1, 'member': 1}})
def ecmwf_ifs_er_debiased(start_time=None, end_time=None, variable="precip", agg_days=1, prob_type='deterministic',
                          grid='global1_5', mask='lsm', region="global"):
    """Standard format forecast data for ECMWF forecasts."""
    return _ecmwf_ifs_er_unified(start_time=start_time, end_time=end_time, variable=variable,
                                 agg_days=agg_days, prob_type=prob_type,
                                 grid=grid, mask=mask, region=region, debiased=True)
