from __future__ import annotations

from typing import Dict, List, Tuple

import numpy as np
import xarray as xr

from .data_utils import TARGET_TO_MODEL, anomaly_for_var


def shift2d(arr: np.ndarray, direction: str) -> np.ndarray:
    """Shift a 2D or 3D time-lat-lon array by one grid cell, filling new edge with NaN."""
    out = np.full_like(arr, np.nan, dtype=float)
    if arr.ndim == 2:
        if direction == "west":
            out[:, 1:] = arr[:, :-1]
        elif direction == "east":
            out[:, :-1] = arr[:, 1:]
        elif direction == "south":
            out[1:, :] = arr[:-1, :]
        elif direction == "north":
            out[:-1, :] = arr[1:, :]
        else:
            raise ValueError(direction)
    elif arr.ndim == 3:
        if direction == "west":
            out[:, :, 1:] = arr[:, :, :-1]
        elif direction == "east":
            out[:, :, :-1] = arr[:, :, 1:]
        elif direction == "south":
            out[:, 1:, :] = arr[:, :-1, :]
        elif direction == "north":
            out[:, :-1, :] = arr[:, 1:, :]
        else:
            raise ValueError(direction)
    else:
        raise ValueError("arr must be 2D or 3D")
    return out


def predictor_variable_names(target: str, predictor_set: str) -> List[str]:
    direct = TARGET_TO_MODEL[target]
    if predictor_set == "P2":
        return [direct]
    if predictor_set == "P4":
        if target == "pr_o":
            return ["pr", "olr", "q2", "rh2", "slp", "u850", "w700", "h500", "h850"]
        if target == "t2_o":
            return ["t2", "tx", "tm", "h500", "h850", "slp", "q2", "rh2", "olr"]
        return [direct]
    if predictor_set == "P3":
        return ["pr", "t2", "tx", "tm", "h500", "h850", "olr", "q2", "rh2", "slp", "u200", "u850", "w700"]
    if predictor_set == "P_fortran22":
        # Handled separately because it includes derived and neighbour features.
        return []
    raise ValueError(f"Unknown predictor_set: {predictor_set}")


def build_feature_cube(
    ds: xr.Dataset,
    clim: Dict[str, np.ndarray],
    target: str,
    predictor_set: str,
    time_indices: np.ndarray,
    month: int,
    lat2d: np.ndarray,
    lon2d: np.ndarray,
    reg: np.ndarray,
    spatial_mode: str,
    include_lat_lon: bool = True,
    include_region_feature: bool = True,
) -> Tuple[np.ndarray, List[str]]:
    """Build feature cube with shape (time, lat, lon, n_features)."""
    features: List[np.ndarray] = []
    names: List[str] = []

    if predictor_set == "P_fortran22":
        direct = TARGET_TO_MODEL[target]
        direct_anom = anomaly_for_var(ds, clim, direct, time_indices, month)
        olr = anomaly_for_var(ds, clim, "olr", time_indices, month)
        q2 = anomaly_for_var(ds, clim, "q2", time_indices, month)
        rh2 = anomaly_for_var(ds, clim, "rh2", time_indices, month)
        w700 = anomaly_for_var(ds, clim, "w700", time_indices, month)
        u200 = anomaly_for_var(ds, clim, "u200", time_indices, month)
        u850 = anomaly_for_var(ds, clim, "u850", time_indices, month)
        shear = u200 - u850
        base_items = [
            (direct_anom, direct),
            (olr, "olr"),
            (q2, "q2"),
            (rh2, "rh2"),
            (w700, "w700"),
            (shear, "u200_minus_u850"),
        ]
        for arr, nm in base_items:
            features.append(arr)
            names.append(nm)
        for arr, base_nm in [(direct_anom, direct), (olr, "olr"), (w700, "w700"), (q2, "q2")]:
            for direction in ["west", "east", "south", "north"]:
                features.append(shift2d(arr, direction))
                names.append(f"{base_nm}_{direction}")
    else:
        for var in predictor_variable_names(target, predictor_set):
            features.append(anomaly_for_var(ds, clim, var, time_indices, month))
            names.append(var)

    nt = len(time_indices)
    if include_lat_lon:
        features.append(np.broadcast_to(lat2d[None, :, :], (nt, lat2d.shape[0], lat2d.shape[1])).astype(float))
        names.append("lat")
        features.append(np.broadcast_to(lon2d[None, :, :], (nt, lon2d.shape[0], lon2d.shape[1])).astype(float))
        names.append("lon")

    if include_region_feature and spatial_mode == "national":
        features.append(np.broadcast_to(reg[None, :, :], (nt, reg.shape[0], reg.shape[1])).astype(float))
        names.append("reg")

    cube = np.stack(features, axis=-1).astype(np.float32)
    return cube, names


def select_samples(
    feature_cube: np.ndarray,
    target_anom: np.ndarray,
    spatial_mask: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Flatten time-lat-lon samples and drop missing features/target.

    Returns X, y, t_indices_flat, lat_indices_flat, lon_indices_flat relative to the supplied cube.
    """
    nt, nlat, nlon, nf = feature_cube.shape
    spatial_flat = spatial_mask.reshape(-1)
    X_all = feature_cube.reshape(nt, nlat * nlon, nf)
    y_all = target_anom.reshape(nt, nlat * nlon)
    keep_space = np.where(spatial_flat)[0]
    Xs = X_all[:, keep_space, :].reshape(-1, nf)
    ys = y_all[:, keep_space].reshape(-1)
    tt = np.repeat(np.arange(nt), len(keep_space))
    flat_space = np.tile(keep_space, nt)
    lat_idx, lon_idx = np.unravel_index(flat_space, (nlat, nlon))
    valid = np.isfinite(ys) & np.all(np.isfinite(Xs), axis=1)
    return Xs[valid], ys[valid], tt[valid], lat_idx[valid], lon_idx[valid]


def predict_to_grid(
    model,
    feature_cube: np.ndarray,
    spatial_mask: np.ndarray,
) -> np.ndarray:
    """Predict to a time-lat-lon grid, leaving invalid samples as NaN."""
    nt, nlat, nlon, nf = feature_cube.shape
    pred = np.full((nt, nlat, nlon), np.nan, dtype=np.float32)
    spatial_flat = spatial_mask.reshape(-1)
    keep_space = np.where(spatial_flat)[0]
    X_all = feature_cube.reshape(nt, nlat * nlon, nf)
    Xs = X_all[:, keep_space, :].reshape(-1, nf)
    valid = np.all(np.isfinite(Xs), axis=1)
    if np.any(valid):
        yhat = np.full(Xs.shape[0], np.nan, dtype=np.float32)
        yhat[valid] = model.predict(Xs[valid]).astype(np.float32)
        pred_flat = pred.reshape(nt, nlat * nlon)
        pred_flat[:, keep_space] = yhat.reshape(nt, len(keep_space))
    return pred