from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

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

MODEL_VARIABLES = [
    "pr", "t2", "tx", "tm", "h500", "h850", "olr", "q2", "rh2", "slp", "u200", "u850", "w700"
]
OBS_VARIABLES = ["pr_o", "t2_o", "tx_o", "tm_o"]
ALL_DATA_VARIABLES = OBS_VARIABLES + MODEL_VARIABLES
TARGET_TO_MODEL = {"pr_o": "pr", "t2_o": "t2", "tx_o": "tx", "tm_o": "tm"}
TARGET_SHORT = {"pr_o": "pr", "t2_o": "t2", "tx_o": "tx", "tm_o": "tm"}


@dataclass
class GridData:
    ds: xr.Dataset
    mask: xr.Dataset
    years: np.ndarray
    months: np.ndarray
    time: pd.DatetimeIndex
    valid_mask: np.ndarray  # lat, lon bool
    reg: np.ndarray  # lat, lon float/int with nan outside regions
    lon2d: np.ndarray
    lat2d: np.ndarray


def _replace_missing(ds: xr.Dataset, missing_value: float) -> xr.Dataset:
    out = ds.copy()
    for name in list(out.data_vars):
        arr = out[name]
        fill_values = [missing_value]
        for att in ["missing_value", "_FillValue"]:
            if att in arr.attrs:
                try:
                    fill_values.append(float(arr.attrs[att]))
                except Exception:
                    pass
        mask = xr.zeros_like(arr, dtype=bool)
        for mv in set(fill_values):
            mask = mask | np.isclose(arr, mv)
        out[name] = arr.where(~mask)
    return out


def open_lead_dataset(data_dir: str | Path, file_pattern: str, source_model: str, lead: int, missing_value: float) -> xr.Dataset:
    fname = file_pattern.format(lead=lead, source_model=source_model)
    path = Path(data_dir) / fname
    if not path.exists():
        raise FileNotFoundError(f"Input file not found: {path}")
    ds = xr.open_dataset(path, decode_times=False)
    # Ensure standard dimension order for data variables.
    for name in list(ds.data_vars):
        if set(["mon", "lat", "lon"]).issubset(ds[name].dims):
            ds[name] = ds[name].transpose("mon", "lat", "lon")
    return _replace_missing(ds, missing_value)


def open_mask(mask_file: str | Path, missing_value: float) -> xr.Dataset:
    path = Path(mask_file)
    if not path.exists():
        raise FileNotFoundError(f"Mask file not found: {path}")
    mask = xr.open_dataset(path, decode_times=False)
    for name in list(mask.data_vars):
        if set(["lat", "lon"]).issubset(mask[name].dims):
            mask[name] = mask[name].transpose("lat", "lon")
    return _replace_missing(mask, missing_value)


def make_time_axis(year_start: int, ntime: int) -> Tuple[pd.DatetimeIndex, np.ndarray, np.ndarray]:
    # User convention: mon=1 means YYYY-01, not one month after YYYY-01.
    time = pd.date_range(f"{year_start}-01-01", periods=ntime, freq="MS")
    years = time.year.to_numpy()
    months = time.month.to_numpy()
    return time, years, months


def load_grid_data(
    data_dir: str | Path,
    mask_file: str | Path,
    file_pattern: str,
    source_model: str,
    lead: int,
    year_start: int,
    missing_value: float,
) -> GridData:
    ds = open_lead_dataset(data_dir, file_pattern, source_model, lead, missing_value)
    mask = open_mask(mask_file, missing_value)

    # Basic grid compatibility checks.
    if ds.sizes["lat"] != mask.sizes["lat"] or ds.sizes["lon"] != mask.sizes["lon"]:
        raise ValueError("Lead file and mask file have different lat/lon sizes.")
    if not np.allclose(ds["lat"].values, mask["lat"].values) or not np.allclose(ds["lon"].values, mask["lon"].values):
        raise ValueError("Lead file and mask file have different lat/lon coordinates.")

    ntime = ds.sizes["mon"]
    time, years, months = make_time_axis(year_start, ntime)
    reg = mask["reg"].values.astype(float)
    valid_mask = np.isfinite(reg) & (reg > 0)
    lon2d, lat2d = np.meshgrid(ds["lon"].values, ds["lat"].values)
    return GridData(ds=ds, mask=mask, years=years, months=months, time=time, valid_mask=valid_mask, reg=reg, lon2d=lon2d, lat2d=lat2d)


def year_range_mask(years: np.ndarray, year_pair: Iterable[int]) -> np.ndarray:
    y1, y2 = list(year_pair)
    return (years >= y1) & (years <= y2)


def calendar_month_mask(months: np.ndarray, month: int) -> np.ndarray:
    return months == month


def compute_train_climatology(
    ds: xr.Dataset,
    years: np.ndarray,
    train_years: Iterable[int],
    variables: List[str],
) -> Dict[str, np.ndarray]:
    """Return climatology[var][month_index, lat, lon] computed from training years only."""
    train_mask = year_range_mask(years, train_years)
    clim: Dict[str, np.ndarray] = {}
    for var in variables:
        arr = ds[var].values  # mon, lat, lon
        out = np.full((12, arr.shape[1], arr.shape[2]), np.nan, dtype=np.float32)
        for m in range(1, 13):
            tmask = train_mask & (((np.arange(arr.shape[0]) % 12) + 1) == m)
            with np.errstate(all="ignore"):
                out[m - 1] = np.nanmean(arr[tmask, :, :], axis=0)
        clim[var] = out
    return clim


def anomaly_for_var(ds: xr.Dataset, clim: Dict[str, np.ndarray], var: str, time_indices: np.ndarray, month: int) -> np.ndarray:
    arr = ds[var].values[time_indices, :, :]  # time, lat, lon
    return arr - clim[var][month - 1][None, :, :]


def raw_for_var(ds: xr.Dataset, var: str, time_indices: np.ndarray) -> np.ndarray:
    return ds[var].values[time_indices, :, :]


def target_obs_climatology(clim: Dict[str, np.ndarray], target: str, month: int) -> np.ndarray:
    return clim[target][month - 1]