from __future__ import annotations

import json
import warnings
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple

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

from src.safe_ops import safe_mean
from sklearn.exceptions import ConvergenceWarning

from .data_utils import ALL_DATA_VARIABLES, TARGET_TO_MODEL, compute_train_climatology, load_grid_data, raw_for_var, year_range_mask
from .features import build_feature_cube, predict_to_grid, select_samples
from .metrics import grid_metrics, metrics_to_dataset, summarize_by_region
from .models import extract_linear_weights, make_model, make_stack_meta_model
from .plotting import plot_all_metric_maps

warnings.filterwarnings("ignore", category=ConvergenceWarning)


def stack_task_label(task: Dict[str, Any]) -> str:
    return (
        f"STACK_{task['stack_scope']}_lead{int(task['lead'])}_"
        f"{task['target']}_{task['spatial_mode']}_{task['meta_method']}"
    )


def _target_units(target: str) -> str:
    if target == "pr_o":
        return "mm/month"
    if target == "t2_o":
        return "C"
    return ""


def _scope_sources(scope: str, cfg: Dict[str, Any]) -> List[str]:
    scope = str(scope)
    if scope == "multi_source":
        return list(cfg.get("source_models", ["RegCM", "clWRF"]))
    if scope == "within_RegCM":
        return ["RegCM"]
    if scope == "within_clWRF":
        return ["clWRF"]
    if scope.startswith("within_"):
        return [scope.replace("within_", "")]
    raise ValueError(f"Unknown stack_scope: {scope}")


def _scope_source_label(scope: str) -> str:
    if scope == "multi_source":
        return "RegCM+clWRF"
    if scope.startswith("within_"):
        return scope.replace("within_", "")
    return scope


def _spatial_masks(valid_mask: np.ndarray, reg: np.ndarray, spatial_mode: str, regions: Iterable[int]) -> List[Tuple[str, np.ndarray]]:
    if spatial_mode == "national":
        return [("ALL", valid_mask)]
    if spatial_mode == "regional":
        return [(f"R{int(r)}", valid_mask & (reg.astype(float) == float(r))) for r in regions]
    raise ValueError(f"Unknown spatial_mode: {spatial_mode}")


def _year_kfolds(years: np.ndarray, train_years: Iterable[int], n_splits: int = 8) -> List[Tuple[np.ndarray, np.ndarray]]:
    train_mask = year_range_mask(years, train_years)
    unique_years = np.unique(years[train_mask])
    if n_splits <= 1:
        raise ValueError("n_splits must be >= 2")
    folds = np.array_split(unique_years, min(n_splits, len(unique_years)))
    out = []
    for val_years in folds:
        val = train_mask & np.isin(years, val_years)
        trn = train_mask & ~np.isin(years, val_years)
        out.append((trn, val))
    return out


def _fit_base_predict_month(gd, clim, cfg, source_model: str, target: str, predictor_set: str, method: str, month: int, train_idx: np.ndarray, pred_idx: np.ndarray, spatial_mask: np.ndarray) -> np.ndarray:
    """Fit one base model on train_idx and predict anomaly on pred_idx for one month and mask."""
    target_train_anom = gd.ds[target].values[train_idx, :, :] - clim[target][month - 1][None, :, :]
    feat_train, _ = build_feature_cube(
        gd.ds, clim, target, predictor_set, train_idx, month, gd.lat2d, gd.lon2d, gd.reg,
        spatial_mode="national",  # ensure same feature definition; region mask controls samples
        include_lat_lon=cfg.get("include_lat_lon", True),
        include_region_feature=cfg.get("include_region_feature", True),
    )
    feat_pred, _ = build_feature_cube(
        gd.ds, clim, target, predictor_set, pred_idx, month, gd.lat2d, gd.lon2d, gd.reg,
        spatial_mode="national",
        include_lat_lon=cfg.get("include_lat_lon", True),
        include_region_feature=cfg.get("include_region_feature", True),
    )
    X_train, y_train, *_ = select_samples(feat_train, target_train_anom, spatial_mask)
    out = np.full((len(pred_idx), gd.ds.sizes["lat"], gd.ds.sizes["lon"]), np.nan, dtype=np.float32)
    if len(y_train) < 20:
        return out
    model = make_model(method, cfg, random_seed=int(cfg.get("random_seed", 42)))
    model.fit(X_train, y_train)
    out = predict_to_grid(model, feat_pred, spatial_mask)
    return out.astype(np.float32)


def _meta_feature_cube(base_pred_list: List[np.ndarray], lat2d: np.ndarray, lon2d: np.ndarray, reg: np.ndarray, include_geo: bool) -> Tuple[np.ndarray, List[str]]:
    arrays = []
    names = []
    for name, arr in base_pred_list:
        arrays.append(arr)
        names.append(name)
    if include_geo:
        nt = arrays[0].shape[0]
        arrays.append(np.broadcast_to(lat2d[None, :, :], arrays[0].shape).astype(np.float32)); names.append("lat")
        arrays.append(np.broadcast_to(lon2d[None, :, :], arrays[0].shape).astype(np.float32)); names.append("lon")
        arrays.append(np.broadcast_to(reg[None, :, :], arrays[0].shape).astype(np.float32)); names.append("reg")
    cube = np.stack(arrays, axis=-1).astype(np.float32)
    return cube, names


def _select_xy(feature_cube: np.ndarray, target_anom: np.ndarray, spatial_mask: np.ndarray):
    return select_samples(feature_cube, target_anom, spatial_mask)


def _predict_meta_to_grid(meta_model, method: str, feature_cube: np.ndarray, spatial_mask: np.ndarray) -> np.ndarray:
    nt, nlat, nlon, nf = feature_cube.shape
    pred = np.full((nt, nlat, nlon), np.nan, dtype=np.float32)
    keep_space = np.where(spatial_mask.reshape(-1))[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)
        if method == "stack_mean":
            yhat[valid] = safe_mean(Xs[valid], axis=1, min_count=1, dtype=np.float32)
        else:
            yhat[valid] = meta_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


def _run_stack_month(gds: Dict[str, Any], clims: Dict[str, Dict[str, np.ndarray]], cfg: Dict[str, Any], task: Dict[str, Any], month: int, train_mask_global: np.ndarray, test_mask_global: np.ndarray, spatial_mask: np.ndarray, meta_weights_accum: List[Dict[str, Any]]):
    st = dict(cfg.get("stacking", {}))
    target = task["target"]
    lead = int(task["lead"])
    meta_method = task["meta_method"]
    scope = task["stack_scope"]
    source_models = _scope_sources(scope, cfg)
    base_methods = list(st.get("base_methods", ["ridge", "xgb_regularized", "mlp_small"]))
    base_predictor_sets = list(st.get("base_predictor_sets", ["P2", "P4"]))
    include_meta_geo = bool(st.get("include_meta_geo_features", False))
    n_splits = int(st.get("n_splits", 8))

    first_gd = gds[source_models[0]]
    train_idx_all = np.where(train_mask_global & (first_gd.months == month))[0]
    test_idx = np.where(test_mask_global & (first_gd.months == month))[0]
    if len(test_idx) == 0:
        shape = (0, first_gd.ds.sizes["lat"], first_gd.ds.sizes["lon"])
        return np.full(shape, np.nan, np.float32), np.full(shape, np.nan, np.float32), np.full(shape, np.nan, np.float32)

    # Target anomaly based on observed climatology from the first source (obs should be common across sources).
    target_train_anom = first_gd.ds[target].values[train_idx_all, :, :] - clims[source_models[0]][target][month - 1][None, :, :]

    base_specs = []
    for src in source_models:
        for ps in base_predictor_sets:
            for meth in base_methods:
                base_specs.append({"source_model": src, "predictor_set": ps, "method": meth, "label": f"{src}_{ps}_{meth}"})

    # OOF base predictions for training period and refit base predictions for test period.
    oof_base = [(spec["label"], np.full((len(train_idx_all), first_gd.ds.sizes["lat"], first_gd.ds.sizes["lon"]), np.nan, dtype=np.float32)) for spec in base_specs]
    test_base = []

    folds = _year_kfolds(first_gd.years, cfg["train_years"], n_splits=n_splits)
    train_pos_lookup = {int(idx): pos for pos, idx in enumerate(train_idx_all)}

    for ib, spec in enumerate(base_specs):
        src = spec["source_model"]
        gd = gds[src]
        clim = clims[src]
        for fold_train_mask, fold_val_mask in folds:
            tr_idx = np.where(fold_train_mask & (gd.months == month))[0]
            va_idx = np.where(fold_val_mask & (gd.months == month))[0]
            if len(va_idx) == 0 or len(tr_idx) == 0:
                continue
            pred_va = _fit_base_predict_month(gd, clim, cfg, src, target, spec["predictor_set"], spec["method"], month, tr_idx, va_idx, spatial_mask)
            for local_va, global_idx in enumerate(va_idx):
                pos = train_pos_lookup.get(int(global_idx))
                if pos is not None:
                    oof_base[ib][1][pos, :, :] = pred_va[local_va, :, :]
        # Refit on all training years and predict final test period.
        pred_test = _fit_base_predict_month(gd, clim, cfg, src, target, spec["predictor_set"], spec["method"], month, train_idx_all, test_idx, spatial_mask)
        test_base.append((spec["label"], pred_test))

    Xmeta_train_cube, meta_names = _meta_feature_cube(oof_base, first_gd.lat2d, first_gd.lon2d, first_gd.reg, include_meta_geo)
    Xmeta_test_cube, _ = _meta_feature_cube(test_base, first_gd.lat2d, first_gd.lon2d, first_gd.reg, include_meta_geo)

    X_meta, y_meta, *_ = _select_xy(Xmeta_train_cube, target_train_anom, spatial_mask)
    if len(y_meta) < int(st.get("min_meta_samples", 30)):
        pred_anom_test = np.full((len(test_idx), first_gd.ds.sizes["lat"], first_gd.ds.sizes["lon"]), np.nan, dtype=np.float32)
    elif meta_method == "stack_mean":
        meta_model = None
        pred_anom_test = _predict_meta_to_grid(meta_model, meta_method, Xmeta_test_cube, spatial_mask)
        meta_weights_accum.append({"month": month, "meta_method": meta_method, "n_samples": len(y_meta), "feature_names": json.dumps(meta_names), "weights_json": json.dumps(extract_linear_weights(meta_model, meta_names))})
    else:
        meta_model = make_stack_meta_model(meta_method, cfg, random_seed=int(cfg.get("random_seed", 42)))
        meta_model.fit(X_meta, y_meta)
        pred_anom_test = _predict_meta_to_grid(meta_model, meta_method, Xmeta_test_cube, spatial_mask)
        meta_weights_accum.append({"month": month, "meta_method": meta_method, "n_samples": len(y_meta), "feature_names": json.dumps(meta_names), "weights_json": json.dumps(extract_linear_weights(meta_model, meta_names))})

    pred_raw = pred_anom_test + clims[source_models[0]][target][month - 1][None, :, :]
    if target == "pr_o" and cfg.get("clip_precipitation_to_zero", True):
        pred_raw = np.where(np.isfinite(pred_raw), np.maximum(pred_raw, 0.0), np.nan).astype(np.float32)

    obs_raw = raw_for_var(first_gd.ds, target, test_idx).astype(np.float32)
    direct_var = TARGET_TO_MODEL[target]
    raw_members = [raw_for_var(gds[src].ds, direct_var, test_idx).astype(np.float32) for src in source_models]
    with np.errstate(invalid="ignore"):
        raw_ref = safe_mean(np.stack(raw_members, axis=0), axis=0, min_count=1, dtype=np.float32)
    return pred_raw.astype(np.float32), obs_raw, raw_ref


def run_stack_task(cfg: Dict[str, Any], task: Dict[str, Any], dirs: Dict[str, Path]) -> Dict[str, Path]:
    label = stack_task_label(task)
    done_file = dirs["done"] / f"{label}.done"
    metrics_file = dirs["metrics"] / f"grid_metrics_{label}.nc"
    summary_file = dirs["metrics"] / f"summary_{label}.csv"
    pred_file = dirs["corrected"] / f"pred_{label}.nc"
    weights_file = dirs["metrics"] / f"stack_weights_{label}.csv"

    if done_file.exists() and not cfg.get("overwrite", False):
        return {"status": "skipped", "label": label, "metrics": metrics_file, "summary": summary_file, "pred": pred_file}

    scope = task["stack_scope"]
    source_models = _scope_sources(scope, cfg)
    gds = {}
    clims = {}
    for src in source_models:
        gd = load_grid_data(
            data_dir=cfg["data_dir"],
            mask_file=cfg["mask_file"],
            file_pattern=cfg.get("file_pattern", "lead{lead}_{source_model}_01_BL_Test_VN.nc"),
            source_model=src,
            lead=int(task["lead"]),
            year_start=int(cfg.get("year_start", 1982)),
            missing_value=float(cfg.get("missing_value", -99.0)),
        )
        gds[src] = gd
        clims[src] = compute_train_climatology(gd.ds, gd.years, cfg["train_years"], list(dict.fromkeys(ALL_DATA_VARIABLES)))

    first_gd = gds[source_models[0]]
    train_mask = year_range_mask(first_gd.years, cfg["train_years"])
    test_mask = year_range_mask(first_gd.years, cfg["test_years"])
    test_global_indices = np.where(test_mask)[0]
    ntest = int(test_mask.sum())
    nlat = first_gd.ds.sizes["lat"]
    nlon = first_gd.ds.sizes["lon"]
    pred = np.full((ntest, nlat, nlon), np.nan, dtype=np.float32)
    obs = np.full_like(pred, np.nan)
    raw = np.full_like(pred, np.nan)
    weights_rows: List[Dict[str, Any]] = []

    for region_label, smask in _spatial_masks(first_gd.valid_mask, first_gd.reg, task["spatial_mode"], cfg.get("valid_regions", [1,2,3,4,5,6,7])):
        if not np.any(smask):
            continue
        for month in range(1, 13):
            test_idx = np.where(test_mask & (first_gd.months == month))[0]
            if len(test_idx) == 0:
                continue
            weight_accum: List[Dict[str, Any]] = []
            pred_m, obs_m, raw_m = _run_stack_month(gds, clims, cfg, task, month, train_mask, test_mask, smask, weight_accum)
            compact_positions = np.searchsorted(test_global_indices, test_idx)
            jj, ii = np.where(smask)
            for kk, pos in enumerate(compact_positions):
                pred[pos, jj, ii] = pred_m[kk, jj, ii]
                obs[pos, jj, ii] = obs_m[kk, jj, ii]
                raw[pos, jj, ii] = raw_m[kk, jj, ii]
            for wr in weight_accum:
                wr.update({"region_unit": region_label, "label": label})
                weights_rows.append(wr)

    invalid = ~first_gd.valid_mask
    pred[:, invalid] = np.nan
    obs[:, invalid] = np.nan
    raw[:, invalid] = np.nan

    metric_values = grid_metrics(pred, obs, rmae_percent=bool(cfg.get("rmae_percent", True)))
    metric_ds = metrics_to_dataset(metric_values, first_gd.ds["lat"].values, first_gd.ds["lon"].values)
    metric_ds.attrs.update({k: str(v) for k, v in task.items()})
    metric_ds.attrs["source_model"] = _scope_source_label(scope)
    metric_ds.attrs["method"] = task["meta_method"]
    metric_ds.to_netcdf(metrics_file)

    st = dict(cfg.get("stacking", {}))
    meta = {
        "source_model": _scope_source_label(scope),
        "lead": int(task["lead"]),
        "target": task["target"],
        "spatial_mode": task["spatial_mode"],
        "predictor_set": "+".join(st.get("base_predictor_sets", ["P2", "P4"])),
        "method": task["meta_method"],
        "stack_scope": scope,
        "base_methods": "+".join(st.get("base_methods", [])),
        "base_source_models": "+".join(source_models),
    }
    rows = summarize_by_region(metric_values, first_gd.reg, cfg.get("valid_regions", [1, 2, 3, 4, 5, 6, 7]), meta)
    pd.DataFrame(rows).to_csv(summary_file, index=False)

    if weights_rows:
        pd.DataFrame(weights_rows).to_csv(weights_file, index=False)

    if cfg.get("save_predictions", True):
        time_test = first_gd.time[test_global_indices]
        ds_pred = xr.Dataset(
            coords={"time": time_test, "lat": first_gd.ds["lat"].values, "lon": first_gd.ds["lon"].values},
            data_vars={
                "obs": (("time", "lat", "lon"), obs.astype(np.float32)),
                "raw": (("time", "lat", "lon"), raw.astype(np.float32)),
                "pred": (("time", "lat", "lon"), pred.astype(np.float32)),
            },
            attrs={**{k: str(v) for k, v in task.items()}, "source_model": _scope_source_label(scope), "units": _target_units(task["target"])},
        )
        ds_pred["obs"].attrs["long_name"] = f"Observed {task['target']}"
        ds_pred["raw"].attrs["long_name"] = f"Raw direct-model reference for {task['target']}"
        ds_pred["pred"].attrs["long_name"] = f"Stacking bias-corrected prediction for {task['target']}"
        ds_pred.to_netcdf(pred_file)

    plot_all_metric_maps(first_gd.ds["lon"].values, first_gd.ds["lat"].values, metric_values, label, dirs["maps"])
    done_file.write_text(json.dumps({"label": label, "status": "done"}, indent=2), encoding="utf-8")
    return {"status": "done", "label": label, "metrics": metrics_file, "summary": summary_file, "pred": pred_file}