"""Interface for graphcast forecasts."""
import numpy as np
import pandas as pd
import xarray as xr
from nuthatch import cache
from nuthatch.processors import timeseries

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


@dask_remote
@timeseries()
@spatial()
@cache(cache_args=['variable', 'init_hour', 'grid'],
       backend_kwargs={'chunking': {"lat": 721, "lon": 1440, "lead_time": 1, "time": 30}})
def graphcast_daily(start_time, end_time, variable, init_hour=0, grid='global0_25', mask=None, region='global'):  # noqa: ARG001
    """A daily Graphcast forecast."""
    if init_hour != 0:
        raise ValueError("Only 0 init hour supported.")

    # Read the three years for gcloud
    ds1 = xr.open_zarr(
        'gs://weathernext/59572747_4_0/zarr/99140631_1_2020_to_2021/forecasts_10d/date_range_2019-12-01_2021-01-22_6_hours.zarr',
        chunks='auto',
        decode_timedelta=True)
    ds1 = ds1.rename({'prediction_timedelta': 'lead_time',
                     '2m_temperature': 'tmp2m', 'total_precipitation_6hr': 'precip'})
    ds1 = ds1[[variable]]
    ds1 = ds1.sel({'time': slice(pd.to_datetime("2019-12-01"), pd.to_datetime("2020-11-30"))})

    ds2 = xr.open_zarr(
        'gs://weathernext/59572747_4_0/zarr/99140631_2_2021_to_2022/forecasts_10d/date_range_2020-12-01_2022-01-22_6_hours.zarr',
        decode_timedelta=True,
        chunks='auto')
    ds2 = ds2.rename({'prediction_timedelta': 'lead_time',
                     '2m_temperature': 'tmp2m', 'total_precipitation_6hr': 'precip'})
    ds2 = ds2[[variable]]
    ds2 = ds2.sel({'time': slice(pd.to_datetime("2020-12-01"), pd.to_datetime("2021-12-31"))})

    ds3 = xr.open_zarr(
        'gs://weathernext/59572747_4_0/zarr/132880704_2022_to_2023/1/forecasts_10d/date_range_2022-01-01_2023-01-01_12_hours.zarr',
        decode_timedelta=True,
        chunks='auto')
    ds3 = ds3.rename({'prediction_timedelta': 'lead_time',
                     '2m_temperature': 'tmp2m', 'total_precipitation_6hr': 'precip'})
    ds3 = ds3[[variable]]
    ds3 = ds3.sel({'time': slice(pd.to_datetime("2022-01-01"), pd.to_datetime("2023-01-01"))})

    # concat them together
    ds = xr.concat([ds1, ds2, ds3], dim='time')

    # Select only the midnight initialization times
    ds = ds.where(ds.time.dt.hour == init_hour, drop=True)

    # Convert units
    K_const = 273.15
    if variable == 'tmp2m':
        ds[variable] = ds[variable] - K_const
        ds.attrs.update(units='C')
        ds = ds.resample(lead_time='1D').mean(dim='lead_time')
    elif variable == 'precip':
        ds[variable] = ds[variable] * 1000.0
        ds.attrs.update(units='mm')
        # Convert from 6hrly to daily precip, robust to different numbers of 6hrly samples in a day
        ds = ds.resample(lead_time='1D').mean(dim='lead_time') * 4.0
        # Can't have precip less than zero (there are some very small negative values)
        ds = np.maximum(ds, 0)
    else:
        raise ValueError(f"Variable {variable} not implemented.")

    # Shift the lead back 6 hours + init time to align with the start date convention
    ds['lead_time'] = ds['lead_time'] - np.timedelta64(6+init_hour, 'h')

    # Convert lat/lon
    ds = lon_base_change(ds)

    return ds


@dask_remote
@timeseries()
@spatial()
@cache(cache_args=['variable', 'init_hour', 'grid'],
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1, "time": 1000},
           'chunk_by_arg': {
               'grid': {
                   'global0_25': {"lat": 721, "lon": 1440, 'lead_time': 1, 'time': 30}
               },
           }
})
def graphcast_daily_wb(start_time, end_time, variable, init_hour=0, grid='global0_25', mask=None, region='global'):  # noqa: ARG001
    """A daily Graphcast forecast."""
    if init_hour != 0:
        raise ValueError("Only 0 init hour supported.")

    # Read the three years for gcloud
    if grid == 'global0_25':
        filename1 = 'gs://weatherbench2/datasets/graphcast/2018/date_range_2017-11-16_2019-02-01_12_hours.zarr'
        filename2 = 'gs://weatherbench2/datasets/graphcast/2020/date_range_2019-11-16_2021-02-01_12_hours.zarr'
    elif grid == 'global1_5':
        # Global 1.5 degree grid is fractionally off from the standard 1.5 grid
        filename1 = 'gs://weatherbench2/datasets/graphcast/2018/date_range_2017-11-16_2019-02-01_12_hours-240x121_equiangular_with_poles_conservative.zarr'
        filename2 = 'gs://weatherbench2/datasets/graphcast/2020/date_range_2019-11-16_2021-02-01_12_hours-240x121_equiangular_with_poles_conservative.zarr'
    else:
        raise ValueError(f"Grid {grid} not implemented.")

    ds1 = xr.open_zarr(filename1,
                       chunks='auto',
                       decode_timedelta=True)
    ds1 = ds1.rename({'prediction_timedelta': 'lead_time',
                     '2m_temperature': 'tmp2m', 'total_precipitation_6hr': 'precip'})
    if 'latitude' in ds1.dims:
        ds1 = ds1.rename({'latitude': 'lat', 'longitude': 'lon'})
    ds1 = ds1[[variable]]

    ds2 = xr.open_zarr(filename2,
                       decode_timedelta=True,
                       chunks='auto')
    ds2 = ds2.rename({'prediction_timedelta': 'lead_time',
                     '2m_temperature': 'tmp2m', 'total_precipitation_6hr': 'precip'})
    if 'latitude' in ds2.dims:
        ds2 = ds2.rename({'latitude': 'lat', 'longitude': 'lon'})
    ds2 = ds2[[variable]]

    # concat them together
    ds = xr.concat([ds1, ds2], dim='time')
    ds = ds.sortby('time')

    # Select only the midnight initialization times
    ds = ds.where(ds.time.dt.hour == init_hour, drop=True)

    # Round the lats to two decimal places
    ds['lat'] = np.round(ds.lat, decimals=2)
    ds['lon'] = np.round(ds.lon, decimals=2)

    # Convert units
    K_const = 273.15
    if variable == 'tmp2m':
        ds[variable] = ds[variable] - K_const
        ds.attrs.update(units='C')
        ds = ds.resample(lead_time='1D').mean(dim='lead_time')
    elif variable == 'precip':
        ds[variable] = ds[variable] * 1000.0
        ds.attrs.update(units='mm')
        # Convert from 6hrly to daily precip, robust to different numbers of 6hrly samples in a day
        ds = ds.resample(lead_time='1D').mean(dim='lead_time') * 4.0
        # Can't have precip less than zero (there are some very small negative values)
        ds = np.maximum(ds, 0)
    else:
        raise ValueError(f"Variable {variable} not implemented.")

    # Shift the lead back 6 hours + init time to align with the start date convention
    ds['lead_time'] = ds['lead_time'] - np.timedelta64(6+init_hour, 'h')

    # Convert lat/lon
    ds = lon_base_change(ds)

    return ds


@dask_remote
@timeseries()
@spatial()
@cache(cache_args=['variable', 'init_hour', 'grid'],
       cache_disable_if={'grid': 'global0_25'},
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 1, "time": 1000},
           'chunk_by_arg': {
               'grid': {
                   'global0_25': {"lat": 721, "lon": 1440, 'lead_time': 1, 'time': 30}
               },
           }
})
def graphcast_daily_regrid(start_time, end_time, variable, init_hour=0,
                           grid='global0_25', mask=None,
                           region='global'):  # noqa: ARG001
    """Regrid for the original Weathernext datasource."""
    ds = graphcast_daily(start_time, end_time, variable, init_hour=init_hour,
                         grid='global0_25', mask=mask, region=region)
    if grid == 'global0_25':
        return ds

    # Regrid onto appropriate grid
    ds = regrid(ds, grid, base='base180', method='conservative', output_chunks={"lat": 121, "lon": 240}, region=region)

    return ds


@dask_remote
@timeseries()
@spatial()
@cache(cache_args=['variable', 'agg_days', 'grid'],
       cache_disable_if={'agg_days': 1},
       backend_kwargs={
           'chunking': {"lat": 121, "lon": 240, "lead_time": 10, "time": 100},
           'chunk_by_arg': {
               'grid': {
                   'global0_25': {"lat": 721, "lon": 1440, 'lead_time': 1, 'time': 30}
               },
           }
})
def graphcast_wb_rolled(start_time, end_time, variable, agg_days, grid='global0_25', mask=None, region='global'):
    """A rolled and aggregated Graphcast forecast."""
    # Grab the init 0 forecast; don't need to regrid
    ds = graphcast_daily_wb(start_time, end_time, variable, init_hour=0, grid=grid, mask=mask, region=region)
    ds = roll_and_agg(ds, agg=agg_days, agg_col="lead_time", agg_fn="mean")
    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 graphcast(start_time=None, end_time=None, variable="precip", agg_days=1, prob_type='deterministic',
              grid='global1_5', mask='lsm', region="global"):  # noqa: ARG001
    """Final Graphcast interface."""
    if prob_type != 'deterministic':
        raise NotImplementedError("Only deterministic forecast implemented for graphcast")

    # Get the data with the right days
    forecast_start = shift_by_days(start_time, -15) if start_time is not None else None
    forecast_end = shift_by_days(end_time, 15) if end_time is not None else None

    # Get the data with the right days
    ds = graphcast_wb_rolled(forecast_start, forecast_end, variable,
                             agg_days=agg_days, grid=grid, mask=mask,
                             region=region)
    ds = ds.assign_attrs(prob_type="deterministic")

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

    return ds
