from __future__ import annotations
import numpy as np
from .io_utils import ALL_VARS, TARGET_TO_OBS


def monthly_climatology(ds, months: np.ndarray, ref_mask: np.ndarray, variables=None, precip_floor: float = 1.0) -> dict[str, np.ndarray]:
    variables = variables or ALL_VARS
    out: dict[str, np.ndarray] = {}
    for v in variables:
        if v not in ds:
            continue
        arr = ds[v].values.astype(float)  # time, lat, lon
        clim = np.full((12, arr.shape[1], arr.shape[2]), np.nan, dtype=np.float32)
        for m in range(1, 13):
            idx = ref_mask & (months == m)
            if np.any(idx):
                clim[m-1] = np.nanmean(arr[idx], axis=0).astype(np.float32)
        if v in {'pr','pr_o'}:
            clim = np.where(np.isfinite(clim) & (clim > 0), clim, precip_floor).astype(np.float32)
        out[v] = clim
    return out


def var_transform(ds, clim: dict[str, np.ndarray], var: str, tidx: np.ndarray, months: np.ndarray, precip_floor: float = 1.0) -> np.ndarray:
    arr = ds[var].values[tidx].astype(float)
    out = np.full_like(arr, np.nan, dtype=np.float32)
    tidx_arr = np.asarray(tidx)
    for m in range(1, 13):
        k = months[tidx_arr] == m
        if not np.any(k):
            continue
        c = clim[var][m-1][None, :, :].astype(float)
        if var == 'pr':
            c = np.where(np.isfinite(c) & (c > 0), c, precip_floor)
            out[k] = (arr[k] / c).astype(np.float32)
        else:
            out[k] = (arr[k] - c).astype(np.float32)
    return out


def target_transform(ds, clim: dict[str, np.ndarray], target: str, tidx: np.ndarray, months: np.ndarray, precip_floor: float = 1.0) -> np.ndarray:
    obs = TARGET_TO_OBS[target]
    arr = ds[obs].values[tidx].astype(float)
    out = np.full_like(arr, np.nan, dtype=np.float32)
    tidx_arr = np.asarray(tidx)
    for m in range(1, 13):
        k = months[tidx_arr] == m
        if not np.any(k):
            continue
        c = clim[obs][m-1][None, :, :].astype(float)
        if target == 'pr':
            c = np.where(np.isfinite(c) & (c > 0), c, precip_floor)
            out[k] = (arr[k] / c).astype(np.float32)
        else:
            out[k] = (arr[k] - c).astype(np.float32)
    return out


def inverse_target(pred_tr: np.ndarray, clim: dict[str, np.ndarray], target: str, tidx: np.ndarray, months: np.ndarray, clip_pr: bool = True) -> np.ndarray:
    obs = TARGET_TO_OBS[target]
    out = np.full_like(pred_tr, np.nan, dtype=np.float32)
    tidx_arr = np.asarray(tidx)
    for m in range(1, 13):
        k = months[tidx_arr] == m
        if not np.any(k):
            continue
        c = clim[obs][m-1][None, :, :].astype(float)
        if target == 'pr':
            x = pred_tr[k] * c
            if clip_pr:
                x = np.where(np.isfinite(x), np.maximum(x, 0.0), np.nan)
            out[k] = x.astype(np.float32)
        else:
            out[k] = (pred_tr[k] + c).astype(np.float32)
    return out


def obs_physical(ds, target: str, tidx: np.ndarray) -> np.ndarray:
    return ds[TARGET_TO_OBS[target]].values[tidx].astype(np.float32)


def raw_physical(ds, target: str, tidx: np.ndarray) -> np.ndarray:
    return ds[target].values[tidx].astype(np.float32)