from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import pandas as pd


@dataclass(frozen=True)
class LeadPeriods:
    train_start: pd.Timestamp
    train_end: pd.Timestamp
    valid_start: pd.Timestamp
    valid_end: pd.Timestamp
    clim_start: pd.Timestamp
    clim_end: pd.Timestamp


def make_time_axis(year_start: int = 1982, n_months: int = 480) -> pd.DatetimeIndex:
    return pd.date_range(f'{year_start}-01-01', periods=int(n_months), freq='MS')


def _shift_month(year: int, month: int, lead: int) -> pd.Timestamp:
    return pd.Timestamp(year=int(year), month=int(month), day=1) + pd.DateOffset(months=int(lead))


def lead_periods(cfg: dict, lead: int) -> LeadPeriods:
    tr0, tr1 = cfg['time']['train_init_years']
    va0, va1 = cfg['time']['valid_init_years']
    cl0, cl1 = cfg['time']['clim_init_years']
    return LeadPeriods(
        train_start=_shift_month(tr0, 1, lead),
        train_end=_shift_month(tr1, 12, lead),
        valid_start=_shift_month(va0, 1, lead),
        valid_end=_shift_month(va1, 12, lead),
        clim_start=_shift_month(cl0, 1, lead),
        clim_end=_shift_month(cl1, 12, lead),
    )


def between(time: pd.DatetimeIndex, start: pd.Timestamp, end: pd.Timestamp) -> np.ndarray:
    return (time >= start) & (time <= end)


def split_masks(cfg: dict, time: pd.DatetimeIndex, lead: int) -> dict[str, np.ndarray]:
    p = lead_periods(cfg, lead)
    train = between(time, p.train_start, p.train_end)
    valid = between(time, p.valid_start, p.valid_end)
    clim = between(time, p.clim_start, p.clim_end)
    full = train | valid
    return {'train': train, 'valid': valid, 'clim': clim, 'full': full}