import os
from datetime import datetime
from typing import Any, Dict

from pathlib import Path
from urllib.request import urlretrieve

import numpy as np
import xarray as xr
from scipy.ndimage import gaussian_filter

from libs.cca_utils import CCAForecaster, CCAConfig
from libs.data_utils import (
    to_360,
    to_180,
    subset_latlon,
    create_dry_mask_frequency,
)
from libs.metrics_utils import (
    tercile_observations_from_training,
    compute_correlation_map,
    compute_rpss_map,
    compute_grocs_map,
    compute_domain_roc_curves,
    compute_domain_reliability,
    build_metric_mask,
    build_prob_metric_mask,
    compute_mode_correlation_map,
    compute_hss_map,
    )
from libs.visualization_utils import (
    plot_tercile_dominant_forecast,
    plot_skill_map,
    plot_roc_curves,
    plot_reliability_diagrams,
    plot_cca_loading_map,
    plot_combined_cca_scores_by_mode,
)

BASE_URL = "https://ftp.cpc.ncep.noaa.gov/International/seasonal"


def get_ftp_data_dir(data_dir=None):
    if data_dir is None:
        data_dir = "./ftp_data"
    data_dir = Path(data_dir)
    data_dir.mkdir(parents=True, exist_ok=True)
    return data_dir


def remove_old_files(filenames, data_dir=None):
    data_dir = get_ftp_data_dir(data_dir)
    for fname in filenames:
        fpath = data_dir / fname
        if fpath.exists():
            print(f"[FTP] Removing old file: {fpath}")
            fpath.unlink()


def fetch_file(filename, data_dir=None, force_download=True):
    data_dir = get_ftp_data_dir(data_dir)
    local_path = data_dir / filename

    if force_download or (not local_path.exists()):
        url = f"{BASE_URL}/{filename}"
        print(f"[FTP] Downloading: {url}")
        urlretrieve(url, local_path)
    else:
        print(f"[FTP] Using cached: {local_path}")

    return str(local_path)


def fetch_cca_inputs(
    lead,
    x_template,
    y_template,
    t_template,
    data_dir=None,
    force_download=True,
):
    x_name = x_template.format(lead=lead)
    y_name = y_template.format(lead=lead)
    t_name = t_template.format(lead=lead)

    filenames = [x_name, y_name, t_name]

    if force_download:
        remove_old_files(filenames, data_dir)

    xtrain = fetch_file(x_name, data_dir, force_download)
    ytrain = fetch_file(y_name, data_dir, force_download)
    xtest = fetch_file(t_name, data_dir, force_download)

    return xtrain, ytrain, xtest


MONTH_ABBR = ["Jan", "Feb", "Mar", "Apr", "May", "Jun",
              "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]


def build_cca_config(**overrides) -> Dict[str, Any]:
    today = datetime.now()
    config = {
        "verbose": 1,
        "leave_n_out": 3,

        "xtrain_file": " ",
        "ytrain_file": " ",
        "xtest_file": " ",
        "fig_dir": "./fig_dir",
        "ccadat": "./ccadat",

        "x_var": "sst",
        "y_var": "precip",

        # examples: "nmme_sst", "nmme_precip", "ersstv5"
        "x_name": None,

        "x_lat_min": -30,
        "x_lat_max": 30,
        "x_lon_min": 120,
        "x_lon_max": 275,

        "y_lat_min": -10,
        "y_lat_max": 18,
        "y_lon_min": 20,
        "y_lon_max": 52,

        "x_plot_extent": None,
        "y_plot_extent": None,

        "hindcast_start_year": 1991,
        "hindcast_end_year": 2020,

        "apply_dry_mask": False,
        "dry_threshold": 0.5,
        "dry_min_fraction": 2 / 3,

        # Probability post-processing
        "smooth_alpha": 0.75,
        "prob_smooth_sigma_cv": 0.0,
        "prob_smooth_sigma_pred": 0.0,
        "tercile_min_prob_percent": 40.0,

        # Dominant-tercile cleanup
        "majority_filter_size": 3,
        "instability_margin_percent": 5.0,

        # Optional sign flip by mode
        "flip_modes": None,

        # Continuous-map smoothing (visual only)
        "x_map_smooth_sigma": 0.0,
        "y_map_smooth_sigma": 0.0,

        "metric_min_valid_time": 3,
        "prob_min_valid_time": 3,

        "x_eof_range": (1, 3),
        "y_eof_range": (1, 3),
        "cca_mode_range": (1, 3),
        "cv_splits": 5,

        # Probability model options
        "probability_method": "cv_error",
        "error_inflation_factor": 1.5,

        "combo_ignore": (3, 3, 3),

        "ic_year": today.year,
        "ic_month": today.month,
        "lead": 1,

        # Optional raw NMME tercile comparison
        "raw_tercile_file": None,
        "plot_raw_tercile": False,

        # Optional NetCDF forecast output
        "save_tercile_netcdf": True,
        "tercile_netcdf_name": None,

        # Optional NetCDF HSS output
        "save_hss_netcdf": True,
        "hss_netcdf_name": None,
    }
    config.update(overrides)
    return config


def _vprint(config: Dict[str, Any], *args, level: int = 1, **kwargs) -> None:
    if config["verbose"] >= level:
        print(*args, **kwargs)


def add_months(year: int, month: int, offset: int) -> tuple[int, int]:
    total = (year * 12 + (month - 1)) + offset
    new_year = total // 12
    new_month = (total % 12) + 1
    return new_year, new_month


def format_month_label(year: int, month: int, *, include_year: bool = True) -> str:
    mon = MONTH_ABBR[month - 1]
    return f"{mon} {year}" if include_year else mon


def _escape_mathtext_spaces(text: str) -> str:
    return text.replace(" ", r"\ ")


def _format_subtitle_line(text: str) -> str:
    return rf"$\mathregular{{{_escape_mathtext_spaces(text)}}}$"


def _normalize_predictor_name(predictor_name: str | None) -> str:
    return (predictor_name or "").strip().lower()


def _pretty_predictor_label(predictor_name: str | None) -> str | None:
    name = _normalize_predictor_name(predictor_name)

    if not name:
        return None

    if name == "ersstv5":
        return "ERSSTv5"

    if name == "nmme_sst":
        return "Predicted NMME SST"

    if name == "nmme_precip":
        return "Predicted NMME Precip"

    parts = name.replace("-", "_").split("_")
    pretty_parts = []
    for p in parts:
        if p.lower() in {"sst", "ersstv5", "nmme"}:
            pretty_parts.append(p.upper() if p.lower() != "ersstv5" else "ERSSTv5")
        elif p.lower() == "precip":
            pretty_parts.append("Precip")
        else:
            pretty_parts.append(p.capitalize())

    return " ".join(pretty_parts)


def format_ic_label(
    ic_year: int,
    ic_month: int,
    *,
    include_year: bool = True,
    predictor_name: str | None = None,
) -> str:
    return format_month_label(ic_year, ic_month, include_year=include_year)


def format_predictor_usage_label(
    ic_year: int,
    ic_month: int,
    *,
    include_year: bool = True,
    predictor_name: str | None = None,
) -> str | None:
    name = _normalize_predictor_name(predictor_name)

    if not name:
        return None

    if name == "ersstv5":
        py, pm = add_months(ic_year, ic_month, -2)
        return f"Using {format_month_label(py, pm, include_year=include_year)} ERSSTv5"

    pretty = _pretty_predictor_label(name)
    return f"Using {pretty}" if pretty else None


def format_valid_season_label(
    ic_year: int,
    ic_month: int,
    lead: int,
    *,
    include_year: bool = True,
) -> str:
    start_year, start_month = add_months(ic_year, ic_month, lead)
    end_year, end_month = add_months(start_year, start_month, 2)

    start_lab = MONTH_ABBR[start_month - 1]
    end_lab = MONTH_ABBR[end_month - 1]

    if include_year:
        if start_year == end_year:
            return f"{start_lab}–{end_lab} {start_year}"
        return f"{start_lab} {start_year}–{end_lab} {end_year}"

    return f"{start_lab}–{end_lab}"


def build_forecast_title(
    base_title: str,
    ic_year: int,
    ic_month: int,
    lead: int,
    *,
    predictor_name: str | None = None,
) -> str:
    issue_label = format_ic_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )
    predictor_label = format_predictor_usage_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )
    valid_label = format_valid_season_label(
        ic_year,
        ic_month,
        lead,
        include_year=True,
    )

    meta = f"Issue: {issue_label}"
    if predictor_label:
        meta += f" | {predictor_label}"
    meta += f" | Valid: {valid_label}"

    return f"{base_title}\n{_format_subtitle_line(meta)}"


def build_hindcast_title(
    base_title: str,
    ic_year: int,
    ic_month: int,
    lead: int,
    *,
    predictor_name: str | None = None,
) -> str:
    issue_label = format_ic_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )
    predictor_label = format_predictor_usage_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )
    valid_label = format_valid_season_label(
        ic_year,
        ic_month,
        lead,
        include_year=True,
    )

    meta = f"Issue: {issue_label}"
    if predictor_label:
        meta += f" | {predictor_label}"
    meta += f" | Valid: {valid_label}"

    return f"{base_title}\n{_format_subtitle_line(meta)}"


def normalize_probabilities(prob: xr.DataArray, category_dim: str = "category") -> xr.DataArray:
    prob = prob.clip(min=0.0)
    denom = prob.sum(category_dim)
    denom = xr.where(np.abs(denom) < 1e-12, np.nan, denom)
    return prob / denom


def soften_probabilities(
    prob: xr.DataArray,
    alpha: float,
    category_dim: str = "category",
) -> xr.DataArray:
    if alpha <= 0:
        raise ValueError("alpha must be > 0")
    prob = prob.clip(min=1e-12)
    prob = prob ** alpha
    prob = prob / prob.sum(category_dim)
    return prob


def smooth_tercile_probabilities(
    prob_da: xr.DataArray,
    sigma: float = 1.0,
    category_dim: str = "category",
    sample_dim: str = "time",
) -> xr.DataArray:
    if sigma is None or sigma <= 0:
        return prob_da.copy()

    out = prob_da.copy(deep=True)

    def _smooth_field(arr: np.ndarray) -> np.ndarray:
        valid = np.isfinite(arr).astype(float)
        filled = np.where(np.isfinite(arr), arr, 0.0)

        num = gaussian_filter(filled * valid, sigma=sigma, mode="nearest")
        den = gaussian_filter(valid, sigma=sigma, mode="nearest")

        sm = np.full_like(arr, np.nan, dtype=float)
        mask = den > 1e-12
        sm[mask] = num[mask] / den[mask]
        return sm

    if sample_dim in prob_da.dims:
        for t in range(prob_da.sizes[sample_dim]):
            for c in range(prob_da.sizes[category_dim]):
                arr = prob_da.isel({sample_dim: t, category_dim: c}).values
                smoothed = _smooth_field(arr)
                out.loc[{sample_dim: prob_da[sample_dim][t], category_dim: prob_da[category_dim][c]}] = smoothed
    else:
        for c in range(prob_da.sizes[category_dim]):
            arr = prob_da.isel({category_dim: c}).values
            smoothed = _smooth_field(arr)
            out.loc[{category_dim: prob_da[category_dim][c]}] = smoothed

    denom = out.sum(category_dim)
    denom = xr.where(np.abs(denom) < 1e-12, np.nan, denom)
    out = out / denom

    return out


def save_tercile_forecast_netcdf(
    prob: xr.DataArray,
    *,
    output_path: str,
    ic_year: int,
    ic_month: int,
    lead: int,
    predictor_name: str | None = None,
    category_dim: str = "category",
) -> str:
    """
    Save CCA tercile forecast probabilities to NetCDF.

    Output variable:
      probs(lat, lon, category)

    category:
      below, near, above
    """
    prob = normalize_probabilities(prob, category_dim=category_dim)

    # Operational forecast output is normally a single target season.
    # Drop singleton time if present; otherwise keep time as a leading dimension.
    if "time" in prob.dims and prob.sizes["time"] == 1:
        prob = prob.isel(time=0, drop=True)

    category_labels = ["below", "near", "above"]
    prob = prob.assign_coords({category_dim: category_labels})

    if "time" in prob.dims:
        prob = prob.transpose("time", "lat", "lon", category_dim)
    else:
        prob = prob.transpose("lat", "lon", category_dim)

    da = prob.rename("probs")

    valid_label = format_valid_season_label(
        ic_year,
        ic_month,
        lead,
        include_year=True,
    )
    issue_label = format_ic_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )

    ds = da.to_dataset()

    ds.attrs.update({
        "title": "CCA tercile probability forecast",
        "description": "Probabilistic tercile forecast generated from the CCA pipeline.",
        "issue": issue_label,
        "valid_season": valid_label,
        "lead": int(lead),
        "predictor": str(predictor_name),
        "category_order": "below, near, above",
        "probability_units": "fraction",
        "created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    })

    ds["probs"].attrs.update({
        "long_name": "tercile forecast probability",
        "units": "1",
        "valid_range": [0.0, 1.0],
    })

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    if os.path.exists(output_path):
        os.remove(output_path)

    ds.to_netcdf(output_path)

    print(f"Saved tercile forecast NetCDF: {output_path}")
    return output_path


def save_hss_netcdf(
    hss_map: xr.DataArray,
    *,
    output_path: str,
    ic_year: int,
    ic_month: int,
    lead: int,
    predictor_name: str | None = None,
) -> str:
    """
    Save CCA cross-validated Heidke Skill Score map to NetCDF.

    Output variable:
      hss(lat, lon)
    """
    hss_da = hss_map.rename("hss")

    valid_label = format_valid_season_label(
        ic_year,
        ic_month,
        lead,
        include_year=True,
    )
    issue_label = format_ic_label(
        ic_year,
        ic_month,
        include_year=True,
        predictor_name=predictor_name,
    )

    ds = hss_da.to_dataset()

    ds.attrs.update({
        "title": "CCA cross-validated Heidke Skill Score",
        "description": "Gridpoint Heidke Skill Score generated from cross-validated CCA tercile forecasts.",
        "issue": issue_label,
        "valid_season": valid_label,
        "lead": int(lead),
        "predictor": str(predictor_name),
        "score": "Heidke Skill Score",
        "created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    })

    ds["hss"].attrs.update({
        "long_name": "Heidke Skill Score",
        "description": "Cross-validated categorical tercile HSS computed from dominant-probability forecast categories.",
        "units": "1",
        "valid_range": [-1.0, 1.0],
    })

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    if os.path.exists(output_path):
        os.remove(output_path)

    ds.to_netcdf(output_path)

    print(f"Saved HSS NetCDF: {output_path}")
    return output_path


def cross_validated_obs_terciles(
    y_train: xr.DataArray,
    *,
    leave_n_out: int,
    sample_dim: str = "time",
) -> xr.DataArray:
    if sample_dim not in y_train.dims:
        raise ValueError(f"{sample_dim!r} not found in y_train dims={y_train.dims}")

    n = y_train.sizes[sample_dim]
    out = xr.full_like(y_train, np.nan, dtype=float)

    for start in range(0, n, leave_n_out):
        stop = min(start + leave_n_out, n)
        test_idx = np.arange(start, stop)

        train_mask = np.ones(n, dtype=bool)
        train_mask[test_idx] = False
        train_idx = np.where(train_mask)[0]

        ytr = y_train.isel({sample_dim: train_idx})
        yte = y_train.isel({sample_dim: test_idx})

        cls = tercile_observations_from_training(
            y_obs=yte,
            y_train=ytr,
            sample_dim=sample_dim,
        )
        out[test_idx, :, :] = cls.values

    return out.rename("obs_tercile_cv")


def print_probability_diagnostics(prob: xr.DataArray, *, verbose: int = 1) -> None:
    if verbose < 1:
        return

    pvals = np.asarray(prob.values, dtype=float)
    finite_count = int(np.isfinite(pvals).sum())
    print("Finite probability count:", finite_count)

    if finite_count == 0:
        raise ValueError("Probability array contains no finite values after CV/post-processing.")

    print("Prob min/max:", float(np.nanmin(pvals)), float(np.nanmax(pvals)))

    prob_sum = prob.sum("category")
    valid_prob_sum = np.isfinite(prob_sum.values)
    if valid_prob_sum.any():
        print(
            "Prob sum min/max over valid cells:",
            float(np.nanmin(prob_sum.values[valid_prob_sum])),
            float(np.nanmax(prob_sum.values[valid_prob_sum])),
        )
        print(
            "Prob sum mean over valid cells:",
            float(np.nanmean(prob_sum.values[valid_prob_sum])),
        )

    vmax = prob.max("category").values
    if np.isfinite(vmax).any():
        print("Mean max prob:", float(np.nanmean(vmax)))


def _normalize_lon_value(lon: float, lon_mode: str) -> float:
    lon = float(lon)
    if lon_mode == "360":
        return lon % 360.0
    if lon_mode == "180":
        out = ((lon + 180.0) % 360.0) - 180.0
        if np.isclose(out, -180.0) and lon > 0:
            return 180.0
        return out
    raise ValueError(f"Unsupported lon_mode={lon_mode!r}. Expected '180' or '360'.")


def _subset_latlon_wrapped(
    da: xr.DataArray,
    *,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    lon_mode: str,
    name: str,
) -> xr.DataArray:
    if "lat" not in da.dims or "lon" not in da.dims:
        raise ValueError(f"{name} must contain 'lat' and 'lon' dims. Found dims={da.dims}")

    out = da.sortby("lat").sortby("lon")
    out = out.sel(lat=slice(lat_min, lat_max))

    lon_min_n = _normalize_lon_value(lon_min, lon_mode)
    lon_max_n = _normalize_lon_value(lon_max, lon_mode)

    if lon_mode == "360":
        raw_span = float(lon_max) - float(lon_min)
        if np.isclose(raw_span, 360.0) or (np.isclose(lon_min_n, lon_max_n) and not np.isclose(lon_min, lon_max)):
            if out.sizes.get("lat", 0) == 0 or out.sizes.get("lon", 0) == 0:
                raise ValueError(
                    f"{name} subset is empty after latitude selection. "
                    f"Requested lat=({lat_min}, {lat_max}), lon=({lon_min}, {lon_max})"
                )
            return out

    lon_vals = np.asarray(out["lon"].values, dtype=float)
    if lon_vals.size == 0:
        raise ValueError(f"{name} latitude subset is empty before longitude selection.")

    domain_min = float(np.nanmin(lon_vals))
    domain_max = float(np.nanmax(lon_vals))

    if lon_max_n >= lon_min_n:
        out = out.sel(lon=slice(lon_min_n, lon_max_n))
    else:
        left = out.sel(lon=slice(lon_min_n, domain_max))
        right = out.sel(lon=slice(domain_min, lon_max_n))

        if left.sizes.get("lon", 0) == 0 and right.sizes.get("lon", 0) == 0:
            out = out.isel(lon=slice(0, 0))
        elif left.sizes.get("lon", 0) == 0:
            out = right
        elif right.sizes.get("lon", 0) == 0:
            out = left
        else:
            if np.isclose(float(left["lon"].values[-1]), float(right["lon"].values[0])):
                right = right.isel(lon=slice(1, None))
            out = xr.concat([left, right], dim="lon")

    if out.sizes.get("lat", 0) == 0 or out.sizes.get("lon", 0) == 0:
        raise ValueError(
            f"{name} subset is empty. "
            f"Requested lat=({lat_min}, {lat_max}), lon=({lon_min}, {lon_max}) "
            f"using lon_mode={lon_mode}, but got shape={tuple(out.shape)}."
        )

    return out


def _subset_predictor_x(
    da: xr.DataArray,
    *,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    name: str,
) -> xr.DataArray:
    return _subset_latlon_wrapped(
        da,
        lat_min=lat_min,
        lat_max=lat_max,
        lon_min=lon_min,
        lon_max=lon_max,
        lon_mode="360",
        name=name,
    )


def load_raw_nmme_tercile_forecast(
    filepath: str,
    *,
    lead: int,
    ic_year: int,
    ic_month: int,
) -> xr.DataArray:
    """
    Load raw NMME tercile forecast from file with:
      prob_below(target, lat, lon)
      prob_norm(target, lat, lon)
      prob_above(target, lat, lon)

    Uses decode_times=False because target units are not standard-decoded
    calendar dates for our workflow. We map:
      lead=1 -> target index 1
      lead=2 -> target index 2
      lead=3 -> target index 3
    """
    ds = xr.open_dataset(filepath, decode_times=False)

    target_index = int(lead)
    if target_index < 0 or target_index >= ds.sizes["target"]:
        raise ValueError(
            f"Requested lead={lead} but target dimension has size {ds.sizes['target']}."
        )

    below = ds["prob_below"].isel(target=target_index)
    near = ds["prob_norm"].isel(target=target_index)
    above = ds["prob_above"].isel(target=target_index)

    raw = xr.concat(
        [below, near, above],
        dim=xr.IndexVariable("category", ["below", "near", "above"]),
    ).transpose("category", "lat", "lon")

    raw = raw.sortby("lat").sortby("lon")
    raw = to_180(raw)

    # convert from percent to probability
    raw = raw / 100.0
    raw = normalize_probabilities(raw, category_dim="category")

    raw.name = "raw_nmme_probability"
    return raw


def run_cca_pipeline(config: Dict[str, Any]) -> Dict[str, Any]:
    os.makedirs(config["fig_dir"], exist_ok=True)
    os.makedirs(config["ccadat"], exist_ok=True)

    x_plot_extent = config.get("x_plot_extent")
    if x_plot_extent is None:
        x_plot_extent = [config["x_lon_min"], config["x_lon_max"], config["x_lat_min"], config["x_lat_max"]]

    y_plot_extent = config.get("y_plot_extent")
    if y_plot_extent is None:
        y_plot_extent = [config["y_lon_min"], config["y_lon_max"], config["y_lat_min"], config["y_lat_max"]]

    predictor_name = config.get("x_name") or config.get("x_var")

    _vprint(config, "Loading data...")

    x_var = config.get("x_var", "sst")
    y_var = config.get("y_var", "precip")

    ds_xtrain = xr.open_dataset(config["xtrain_file"])
    ds_ytrain = xr.open_dataset(config["ytrain_file"])
    ds_xtest = xr.open_dataset(config["xtest_file"])

    def _get_var(ds, var_name, label):
        if var_name in ds:
            return ds[var_name]
        fallback = list(ds.data_vars)[0]
        print(f"[WARN] '{var_name}' not found in {label}. Using '{fallback}' instead.")
        return ds[fallback]

    X_train = _get_var(ds_xtrain, x_var, "X_train")
    Y_train = _get_var(ds_ytrain, y_var, "Y_train")
    X_test = _get_var(ds_xtest, x_var, "X_test")

    X_train = to_360(X_train.sortby("lat").sortby("lon"))
    X_test = to_360(X_test.sortby("lat").sortby("lon"))
    Y_train = to_180(Y_train.sortby("lat").sortby("lon"))
    X_train = X_train.drop_duplicates("lon").drop_duplicates("lat")
    X_test = X_test.drop_duplicates("lon").drop_duplicates("lat")
    Y_train = Y_train.drop_duplicates("lon").drop_duplicates("lat")

    _vprint(config, "Subsetting regions...")

    X_train = _subset_predictor_x(
        X_train,
        lat_min=config["x_lat_min"],
        lat_max=config["x_lat_max"],
        lon_min=config["x_lon_min"],
        lon_max=config["x_lon_max"],
        name="X_train",
    )
    X_test = _subset_predictor_x(
        X_test,
        lat_min=config["x_lat_min"],
        lat_max=config["x_lat_max"],
        lon_min=config["x_lon_min"],
        lon_max=config["x_lon_max"],
        name="X_test",
    )

    Y_train = subset_latlon(
        Y_train,
        lat_min=config["y_lat_min"],
        lat_max=config["y_lat_max"],
        lon_min=config["y_lon_min"],
        lon_max=config["y_lon_max"],
        lon_mode="180",
        name="Y_train",
    )

    if X_train.size == 0 or Y_train.size == 0:
        raise ValueError("Selected domain returned empty data. Check bounds.")

    _vprint(config, "Subset shapes:")
    _vprint(config, "  X_train:", X_train.shape)
    _vprint(config, "  Y_train:", Y_train.shape)
    _vprint(config, "  X_test :", X_test.shape)
    _vprint(
        config,
        f"  X lon range: {float(X_train.lon.min().item()):.3f} to {float(X_train.lon.max().item()):.3f}"
    )

    if X_train.sizes["time"] != Y_train.sizes["time"]:
        raise ValueError(
            f"X_train and Y_train sample counts differ after subsetting: "
            f"{X_train.sizes['time']} vs {Y_train.sizes['time']}"
        )

    X_train = X_train.assign_coords(time=Y_train.time)

    n_years = Y_train.sizes["time"]
    years = np.arange(config["hindcast_start_year"], config["hindcast_start_year"] + n_years)

    if config["hindcast_end_year"] != years[-1]:
        _vprint(
            config,
            f"Warning: generated hindcast years end at {years[-1]}, "
            f"not {config['hindcast_end_year']}.",
        )

    X_train = X_train.assign_coords(year=("time", years))
    Y_train = Y_train.assign_coords(year=("time", years))

    if config["apply_dry_mask"]:
        dry_mask = create_dry_mask_frequency(
            Y_train,
            sample_dim="time",
            threshold=config["dry_threshold"],
            min_fraction=config["dry_min_fraction"],
        )
        Y_train = Y_train.where(~dry_mask)
        _vprint(config, "Dry masked cells:", int(dry_mask.sum().item()))
    else:
        dry_mask = xr.zeros_like(Y_train.isel(time=0), dtype=bool).rename("dry_mask")
        _vprint(config, "Dry mask disabled.")

    x_mask = xr.DataArray(
        np.all(np.isfinite(X_train.values), axis=0),
        dims=("lat", "lon"),
        coords={"lat": X_train.lat, "lon": X_train.lon},
        name="x_mask",
    )
    y_mask = xr.DataArray(
        np.all(np.isfinite(Y_train.values), axis=0),
        dims=("lat", "lon"),
        coords={"lat": Y_train.lat, "lon": Y_train.lon},
        name="y_mask",
    )

    _vprint(config, "Valid grid cells:")
    _vprint(config, "  X:", int(x_mask.sum().item()))
    _vprint(config, "  Y:", int(y_mask.sum().item()))

    if int(x_mask.sum().item()) == 0:
        raise ValueError("X predictor region has zero valid grid cells after masking.")
    if int(y_mask.sum().item()) == 0:
        raise ValueError("Y predictand region has zero valid grid cells after masking.")

    cca_config = CCAConfig(
        x_eof_range=config["x_eof_range"],
        y_eof_range=config["y_eof_range"],
        cca_mode_range=config["cca_mode_range"],
        cv_splits=config["cv_splits"],
        combo_ignore=config["combo_ignore"],
        probability_method=config["probability_method"],
        error_inflation_factor=config.get("error_inflation_factor", 1.0),
        verbose=config["verbose"],
    )

    model = CCAForecaster(cca_config)

    _vprint(config, "Running cross-validation...")
    Y_cv = model.cross_validated_predict_xarray(
        X_train,
        Y_train,
        leave_n_out=config["leave_n_out"],
        sample_dim="time",
        x_domain_dims=("lat", "lon"),
        y_domain_dims=("lat", "lon"),
        x_mask=x_mask,
        y_mask=y_mask,
        auto_latitude_weights=True,
        x_lat_dim="lat",
        y_lat_dim="lat",
    )

    P_cv = model.cross_validated_predict_tercile_proba_xarray(
        X_train,
        Y_train,
        leave_n_out=config["leave_n_out"],
        sample_dim="time",
        x_domain_dims=("lat", "lon"),
        y_domain_dims=("lat", "lon"),
        x_mask=x_mask,
        y_mask=y_mask,
        auto_latitude_weights=True,
        x_lat_dim="lat",
        y_lat_dim="lat",
    )

    Y_cv.name = "cv_prediction"
    P_cv.name = "cv_probability"

    Y_cv = Y_cv.assign_coords(year=("time", years))
    P_cv = P_cv.assign_coords(year=("time", years))

    _vprint(config, "After CV:")
    _vprint(config, "  Y_cv finite:", int(np.isfinite(Y_cv.values).sum()))
    _vprint(config, "  P_cv finite:", int(np.isfinite(P_cv.values).sum()))

    cv_resid = Y_train.values - Y_cv.values
    model.cv_error_variance_ = np.nanvar(cv_resid, axis=0, ddof=0)

    P_cv = normalize_probabilities(P_cv, category_dim="category")
    P_cv = soften_probabilities(P_cv, alpha=config["smooth_alpha"], category_dim="category")
    P_cv = normalize_probabilities(P_cv, category_dim="category")
    P_cv = smooth_tercile_probabilities(
        P_cv,
        sigma=config.get("prob_smooth_sigma_cv", 0.0),
        category_dim="category",
        sample_dim="time",
    )
    P_cv = normalize_probabilities(P_cv, category_dim="category")

    print_probability_diagnostics(P_cv, verbose=config["verbose"])

    _vprint(config, "Computing metrics...")
    obs_class = cross_validated_obs_terciles(
        Y_train,
        leave_n_out=config["leave_n_out"],
        sample_dim="time",
    )

    metric_mask = build_metric_mask(
        Y_train,
        Y_cv,
        sample_dim="time",
        min_valid_time=config["metric_min_valid_time"],
    ) & y_mask

    prob_mask = build_prob_metric_mask(
        obs_class,
        P_cv,
        sample_dim="time",
        category_dim="category",
        min_valid_time=config["prob_min_valid_time"],
        require_sum_to_one=False,
    ) & y_mask

    Y_train_m = Y_train.where(metric_mask)
    Y_cv_m = Y_cv.where(metric_mask)
    obs_class_m = obs_class.where(prob_mask)
    P_cv_m = P_cv.where(prob_mask)

    corr_map = compute_correlation_map(
        Y_cv_m,
        Y_train_m,
        sample_dim="time",
    ).where(y_mask)
    corr_map.name = "correlation"

    rpss_map = compute_rpss_map(
        P_cv_m,
        obs_class_m,
        category_dim="category",
        sample_dim="time",
        normalize_probabilities=True,
    ).where(y_mask)
    rpss_map.name = "rpss"

    grocs_map = compute_grocs_map(
        P_cv_m,
        obs_class_m,
        sample_dim="time",
        category_dim="category",
    ).where(y_mask)
    grocs_map.name = "grocs"

    hss_map = compute_hss_map(
        P_cv_m,
        obs_class_m,
        sample_dim="time",
        category_dim="category",
        min_valid_time=config["prob_min_valid_time"],
    ).where(y_mask)

    hss_map.name = "hss"

    roc_dict = compute_domain_roc_curves(P_cv_m, obs_class_m, category_dim="category")
    rel_dict = compute_domain_reliability(P_cv_m, obs_class_m, category_dim="category")

    hss_netcdf_path = None

    _vprint(config, "Saving verification plots...")

    plot_skill_map(
        corr_map,
        title=build_hindcast_title(
            f"CCA Cross-Validated Correlation (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        extent=y_plot_extent,
        smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
        output_path=os.path.join(config["fig_dir"], f"cca_corr_l{config['leave_n_out']}o.png"),
    )

    plot_skill_map(
        rpss_map,
        title=build_hindcast_title(
            f"CCA Cross-Validated RPSS (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        extent=y_plot_extent,
        smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
        output_path=os.path.join(config["fig_dir"], f"cca_rpss_l{config['leave_n_out']}o.png"),
    )

    plot_skill_map(
        grocs_map,
        title=build_hindcast_title(
            f"CCA Cross-Validated GROCS (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        extent=y_plot_extent,
        smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
        output_path=os.path.join(config["fig_dir"], f"cca_grocs_l{config['leave_n_out']}o.png"),
    )
     
    plot_skill_map(
        hss_map,
        title=build_hindcast_title(
            f"CCA Cross-Validated HSS (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        extent=y_plot_extent,
        smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
        output_path=os.path.join(config["fig_dir"], f"cca_hss_l{config['leave_n_out']}o.png"),
    )

    if config.get("save_hss_netcdf", True):
        if config.get("hss_netcdf_name"):
            hss_netcdf_name = config["hss_netcdf_name"]
        else:
            dtt = f"{config['ic_year']:04d}{config['ic_month']:02d}"
            xname = _normalize_predictor_name(predictor_name) or "predictor"
            hss_netcdf_name = f"cca_hss_{xname}_lead{config['lead']}_{dtt}.nc"

        hss_netcdf_path = os.path.join(config["ccadat"], hss_netcdf_name)

        save_hss_netcdf(
            hss_map,
            output_path=hss_netcdf_path,
            ic_year=config["ic_year"],
            ic_month=config["ic_month"],
            lead=config["lead"],
            predictor_name=predictor_name,
        )

    plot_roc_curves(
        roc_dict,
        title=build_hindcast_title(
            f"CCA ROC Curves (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        output_path=os.path.join(config["fig_dir"], f"cca_roc_l{config['leave_n_out']}o.png"),
    )

    plot_reliability_diagrams(
        rel_dict,
        title=build_hindcast_title(
            f"CCA Reliability (L{config['leave_n_out']}O)",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        output_path=os.path.join(config["fig_dir"], f"cca_reliability_l{config['leave_n_out']}o.png"),
    )

    _vprint(config, "Fitting final model...")
    model.fit_xarray(
        X_train,
        Y_train,
        sample_dim="time",
        x_domain_dims=("lat", "lon"),
        y_domain_dims=("lat", "lon"),
        x_mask=x_mask,
        y_mask=y_mask,
        auto_latitude_weights=False,
    )

    canon_corrs = np.asarray(model.summary_["canonical_correlations"], dtype=float)

    if not hasattr(model, "cv_error_variance_"):
        model.cv_error_variance_ = np.nanvar(cv_resid, axis=0, ddof=0)

    P_pred = model.predict_tercile_proba_xarray(
        X_test,
        sample_dim="time",
        x_domain_dims=("lat", "lon"),
    ).where(y_mask)

    P_pred = normalize_probabilities(P_pred, category_dim="category")
    P_pred = smooth_tercile_probabilities(
        P_pred,
        sigma=config.get("prob_smooth_sigma_pred", 0.0),
        category_dim="category",
        sample_dim="time",
    )
    P_pred = normalize_probabilities(P_pred, category_dim="category")

    plot_tercile_dominant_forecast(
        P_pred,
        title=build_forecast_title(
            "CCA Dominant Tercile Forecast",
            config["ic_year"],
            config["ic_month"],
            config["lead"],
            predictor_name=predictor_name,
        ),
        extent=y_plot_extent,
        min_prob_percent=config.get("tercile_min_prob_percent", 40.0),
        majority_filter_size=config.get("majority_filter_size", 3),
        instability_margin_percent=config.get("instability_margin_percent", 5.0),
        output_path=os.path.join(config["fig_dir"], "cca_tercile_fcst.png"),
    )

    tercile_netcdf_path = None

    if config.get("save_tercile_netcdf", True):
        if config.get("tercile_netcdf_name"):
            tercile_netcdf_name = config["tercile_netcdf_name"]
        else:
            dtt = f"{config['ic_year']:04d}{config['ic_month']:02d}"
            xname = _normalize_predictor_name(predictor_name) or "predictor"
            tercile_netcdf_name = f"cca_tercile_fcst_{xname}_lead{config['lead']}_{dtt}.nc"

        tercile_netcdf_path = os.path.join(config["ccadat"], tercile_netcdf_name)

        save_tercile_forecast_netcdf(
            P_pred,
            output_path=tercile_netcdf_path,
            ic_year=config["ic_year"],
            ic_month=config["ic_month"],
            lead=config["lead"],
            predictor_name=predictor_name,
            category_dim="category",
        )

    raw_nmme_prob = None
    if config.get("plot_raw_tercile", False) and config.get("raw_tercile_file"):
        raw_nmme_prob = load_raw_nmme_tercile_forecast(
            config["raw_tercile_file"],
            lead=config["lead"],
            ic_year=config["ic_year"],
            ic_month=config["ic_month"],
        )

        plot_tercile_dominant_forecast(
            raw_nmme_prob,
            title=build_forecast_title(
                "Original NMME Dominant Tercile Forecast",
                config["ic_year"],
                config["ic_month"],
                config["lead"],
                predictor_name="nmme_precip",
            ),
            extent=y_plot_extent,
            min_prob_percent=config.get("tercile_min_prob_percent", 40.0),
            majority_filter_size=config.get("majority_filter_size", 3),
            instability_margin_percent=config.get("instability_margin_percent", 5.0),
            output_path=os.path.join(config["fig_dir"], "nmme_raw_tercile_fcst.png"),
        )

    X_load = model.get_x_cca_loadings_xarray()
    Y_load = model.get_y_cca_loadings_xarray().where(y_mask)
    X_scores = model.get_x_cca_scores_xarray()
    Y_scores = model.get_y_cca_scores_xarray()

    flip_modes = config.get("flip_modes", None)
    if flip_modes is not None:
        for m, sign in flip_modes.items():
            if sign not in (-1, 1):
                continue
            if m in X_load["mode"].values:
                X_load.loc[dict(mode=m)] *= sign
                Y_load.loc[dict(mode=m)] *= sign
                X_scores.loc[dict(mode=m)] *= sign
                Y_scores.loc[dict(mode=m)] *= sign

    n_modes = min(3, X_scores.sizes["mode"], Y_scores.sizes["mode"], len(canon_corrs))

    for m in range(1, n_modes + 1):
        r = float(canon_corrs[m - 1])

        plot_cca_loading_map(
            X_load,
            mode=m,
            title=build_hindcast_title(
                f"X CCA Loading - Mode {m} (r={r:.2f})",
                config["ic_year"],
                config["ic_month"],
                config["lead"],
                predictor_name=predictor_name,
            ),
            extent=x_plot_extent,
            smooth_sigma=config.get("x_map_smooth_sigma", 0.0),
            output_path=os.path.join(config["fig_dir"], f"x_loading_mode{m}.png"),
        )

        plot_cca_loading_map(
            Y_load,
            mode=m,
            title=build_hindcast_title(
                f"Y CCA Loading - Mode {m} (r={r:.2f})",
                config["ic_year"],
                config["ic_month"],
                config["lead"],
                predictor_name=predictor_name,
            ),
            extent=y_plot_extent,
            smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
            output_path=os.path.join(config["fig_dir"], f"y_loading_mode{m}.png"),
        )

        X_corr = compute_mode_correlation_map(X_train, X_scores, mode=m)
        Y_corr = compute_mode_correlation_map(Y_train, Y_scores, mode=m).where(y_mask)

        plot_skill_map(
            X_corr,
            title=build_hindcast_title(
                f"X CCA Score Correlation - Mode {m} (r={r:.2f})",
                config["ic_year"],
                config["ic_month"],
                config["lead"],
                predictor_name=predictor_name,
            ),
            extent=x_plot_extent,
            levels=np.linspace(-1, 1, 21),
            smooth_sigma=config.get("x_map_smooth_sigma", 0.0),
            output_path=os.path.join(config["fig_dir"], f"x_corr_mode{m}.png"),
        )

        plot_skill_map(
            Y_corr,
            title=build_hindcast_title(
                f"Y CCA Score Correlation - Mode {m} (r={r:.2f})",
                config["ic_year"],
                config["ic_month"],
                config["lead"],
                predictor_name=predictor_name,
            ),
            extent=y_plot_extent,
            levels=np.linspace(-1, 1, 21),
            smooth_sigma=config.get("y_map_smooth_sigma", 0.0),
            output_path=os.path.join(config["fig_dir"], f"y_corr_mode{m}.png"),
        )

    plot_combined_cca_scores_by_mode(
        X_scores,
        Y_scores,
        years=years,
        canonical_corrs=canon_corrs,
        output_dir=config["fig_dir"],
    )

    return {
        "model": model,
        "canonical_correlations": canon_corrs,
        "X_train": X_train,
        "Y_train": Y_train,
        "X_test": X_test,
        "Y_cv": Y_cv,
        "P_cv": P_cv,
        "P_pred": P_pred,
        "obs_class": obs_class,
        "corr_map": corr_map,
        "rpss_map": rpss_map,
        "grocs_map": grocs_map,
        "hss_map": hss_map,
        "roc_dict": roc_dict,
        "rel_dict": rel_dict,
        "X_scores": X_scores,
        "Y_scores": Y_scores,
        "years": years,
        "x_mask": x_mask,
        "y_mask": y_mask,
        "dry_mask": dry_mask,
        "raw_nmme_prob": raw_nmme_prob,
        "tercile_netcdf_path": tercile_netcdf_path,
        "hss_netcdf_path": hss_netcdf_path,
    }
