from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import numpy as np
import pandas as pd


@dataclass(frozen=True)
class YearMonth:
    year: int
    month: int

    @classmethod
    def parse(cls, s: str) -> 'YearMonth':
        parts = str(s).replace('_', '-').split('-')
        if len(parts) != 2:
            raise ValueError(f'Expected YYYY-MM, got {s!r}')
        return cls(int(parts[0]), int(parts[1]))

    def to_timestamp(self) -> pd.Timestamp:
        return pd.Timestamp(self.year, self.month, 1)


def monthly_time_axis(start_year: int, n_months: int) -> pd.DatetimeIndex:
    return pd.date_range(f'{start_year}-01-01', periods=int(n_months), freq='MS')


def to_datetime_index(values, start_year: int = 1982) -> pd.DatetimeIndex:
    """Return a robust monthly DatetimeIndex.

    xarray/pandas/netCDF combinations can return DatetimeIndex, ndarray of
    datetime64, cftime objects, or plain numeric coordinates.  This helper keeps
    all pipeline stages independent of whether the input object exposes
    `.to_numpy()`.
    """
    try:
        dt = pd.to_datetime(values)
        if len(dt) > 0 and not pd.isna(dt[0]):
            return pd.DatetimeIndex(dt)
    except Exception:
        pass
    return pd.date_range(f'{int(start_year)}-01-01', periods=len(values), freq='MS')


def reference_mask_by_ym(time, start_ym: str, end_ym: str) -> np.ndarray:
    t0 = YearMonth.parse(start_ym).to_timestamp()
    t1 = YearMonth.parse(end_ym).to_timestamp()
    ti = pd.DatetimeIndex(pd.to_datetime(time))
    return np.asarray((ti >= t0) & (ti <= t1), dtype=bool)


def lead_to_target_month(init_year: int, init_month: int, lead: int) -> Tuple[int, int]:
    base = pd.Timestamp(int(init_year), int(init_month), 1)
    target = base + pd.DateOffset(months=int(lead))
    return int(target.year), int(target.month)


def days_in_month(year: int, month: int) -> int:
    return int(pd.Period(f'{int(year):04d}-{int(month):02d}').days_in_month)