from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence, Tuple, Union

import numpy as np
import xarray as xr
from scipy.linalg import svd
from scipy.stats import spearmanr, t
from sklearn.model_selection import KFold


@dataclass
class CCAConfig:
    # ----------------------------------
    # Mode search ranges
    # ----------------------------------
    x_eof_range: Tuple[int, int] = (1, 5)
    y_eof_range: Tuple[int, int] = (1, 5)
    cca_mode_range: Tuple[int, int] = (1, 5)

    # ----------------------------------
    # Cross-validation
    # ----------------------------------
    cv_splits: int = 5
    leave_n_out: int = 1
    shuffle: bool = False
    random_state: Optional[int] = None

    # ----------------------------------
    # Probability generation
    # ----------------------------------
    probability_method: str = "constant_error"
    error_inflation_factor: float = 1.0

    # ----------------------------------
    # Weights
    # ----------------------------------
    x_weights: Optional[np.ndarray] = None
    y_weights: Optional[np.ndarray] = None

    # ----------------------------------
    # Numerical stability
    # ----------------------------------
    epsilon: float = 1e-8

    # ----------------------------------
    # Optional override (skip mode search)
    # ----------------------------------
    combo_ignore: Optional[Tuple[int, int, int]] = None

    # ----------------------------------
    # Optional extensions
    # ----------------------------------
    selection_metric: str = "spearman"
    top_n_models: int = 1
    verbose: int = 1

    def __post_init__(self) -> None:
        valid_probability_methods = {
            "constant_error",
            "adjusted_error",
            "cv_error",
            "inflated_error",
        }
        if self.probability_method not in valid_probability_methods:
            raise ValueError(
                f"probability_method must be one of {valid_probability_methods}, "
                f"got {self.probability_method!r}"
            )

        for name, value in {
            "x_eof_range": self.x_eof_range,
            "y_eof_range": self.y_eof_range,
            "cca_mode_range": self.cca_mode_range,
        }.items():
            if len(value) != 2 or value[0] < 1 or value[1] < value[0]:
                raise ValueError(
                    f"{name} must be a tuple like (min_modes, max_modes) with 1 <= min <= max. "
                    f"Got {value!r}"
                )

        if self.cv_splits < 2:
            raise ValueError("cv_splits must be at least 2")
        if self.leave_n_out < 1:
            raise ValueError("leave_n_out must be at least 1")
        if self.epsilon <= 0:
            raise ValueError("epsilon must be > 0")
        if self.error_inflation_factor <= 0:
            raise ValueError("error_inflation_factor must be > 0")


class CCAForecaster:
    """
    Canonical Correlation Analysis forecaster with:
      - NumPy and xarray interfaces
      - optional latitude weighting
      - cross-validated deterministic and tercile-probability hindcasts
      - loading/score diagnostics
    """

    def __init__(self, config: Optional[CCAConfig] = None):
        self.config = config or CCAConfig()
        self._is_fitted = False
        self._xarray_metadata: Dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Logging
    # ------------------------------------------------------------------
    def _log(self, msg: str, level: int = 1) -> None:
        if self.config.verbose >= level:
            print(msg)

    # ------------------------------------------------------------------
    # Public API: numpy
    # ------------------------------------------------------------------
    def fit(self, X: np.ndarray, Y: np.ndarray) -> "CCAForecaster":
        X, Y = self._validate_inputs(X, Y)
        self._train_targets = Y.copy()
        self._xarray_metadata = {}

        self._log("[CCA] Starting fit...", 1)

        if self.config.combo_ignore is not None:
            best = self.config.combo_ignore
            self._log(f"[CCA] Using combo_ignore={best}", 1)
        else:
            best = self._select_best_modes(X, Y)

        self.x_modes_requested_, self.y_modes_requested_, self.cca_modes_requested_ = best
        self._cross_validated_error_variance(X, Y, best)
        self._fit_core(X, Y, *best, calc_loadings=True)

        self._is_fitted = True
        self._log("[CCA] Final fit complete.", 1)
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        self._check_is_fitted()
        return self._predict_from_fitted_core(X)

    def predict_quantile_exceedance(
        self,
        X: np.ndarray,
        quantiles: Union[float, Sequence[float]],
    ) -> np.ndarray:
        self._check_is_fitted()
        return self._predict_quantile_exceedance_from_fitted_core(X, quantiles)

    def predict_tercile_proba(
        self,
        X: np.ndarray,
        lower_threshold: Optional[np.ndarray] = None,
        upper_threshold: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        self._check_is_fitted()
        return self._predict_tercile_proba_from_fitted_core(
            X,
            lower_threshold=lower_threshold,
            upper_threshold=upper_threshold,
        )

    def prediction_error_variance(self, X: np.ndarray) -> np.ndarray:
        self._check_is_fitted()
        return self._prediction_error_variance_from_fitted_core(X)

    # ------------------------------------------------------------------
    # Public API: xarray
    # ------------------------------------------------------------------
    def fit_xarray(
        self,
        X: xr.DataArray,
        Y: xr.DataArray,
        *,
        sample_dim: str = "time",
        x_domain_dims: Optional[Sequence[str]] = None,
        y_domain_dims: Optional[Sequence[str]] = None,
        x_mask: Optional[xr.DataArray] = None,
        y_mask: Optional[xr.DataArray] = None,
        auto_latitude_weights: bool = False,
        x_lat_dim: str = "lat",
        y_lat_dim: str = "lat",
    ) -> "CCAForecaster":
        X2d, Y2d, metadata, x_weights, y_weights = self._prepare_xarray_pair(
            X,
            Y,
            sample_dim=sample_dim,
            x_domain_dims=x_domain_dims,
            y_domain_dims=y_domain_dims,
            x_mask=x_mask,
            y_mask=y_mask,
            auto_latitude_weights=auto_latitude_weights,
            x_lat_dim=x_lat_dim,
            y_lat_dim=y_lat_dim,
        )

        old_xw = self.config.x_weights
        old_yw = self.config.y_weights

        try:
            if auto_latitude_weights:
                self.config.x_weights = x_weights
                self.config.y_weights = y_weights

            self._log(
                f"[CCA] Samples={X2d.shape[0]}, X domains={X2d.shape[1]}, Y domains={Y2d.shape[1]}",
                1,
            )
            self.fit(X2d, Y2d)
            self._xarray_metadata = metadata
            return self
        finally:
            self.config.x_weights = old_xw
            self.config.y_weights = old_yw

    def predict_xarray(
        self,
        X: xr.DataArray,
        *,
        sample_dim: Optional[str] = None,
        x_domain_dims: Optional[Sequence[str]] = None,
    ) -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()

        X2d, sample_coords = self._flatten_new_xarray_predictor(
            X,
            sample_dim=sample_dim,
            x_domain_dims=x_domain_dims,
        )
        Y2d = self._predict_from_fitted_core(X2d)
        return self._rebuild_y_xarray(Y2d, sample_coords, name="prediction")

    def predict_tercile_proba_xarray(
        self,
        X: xr.DataArray,
        *,
        sample_dim: Optional[str] = None,
        x_domain_dims: Optional[Sequence[str]] = None,
        category_dim: str = "category",
        category_labels: Sequence[str] = ("below", "near", "above"),
    ) -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()

        X2d, sample_coords = self._flatten_new_xarray_predictor(
            X,
            sample_dim=sample_dim,
            x_domain_dims=x_domain_dims,
        )
        probs = self._predict_tercile_proba_from_fitted_core(X2d)
        return self._rebuild_probability_xarray(
            probs,
            sample_coords,
            category_dim=category_dim,
            category_labels=category_labels,
            name="probability",
        )

    # ------------------------------------------------------------------
    # Cross-validated hindcast outputs
    # ------------------------------------------------------------------
    def cross_validated_predict(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        *,
        leave_n_out: int = 1,
    ) -> np.ndarray:
        X, Y = self._validate_inputs(X, Y)
        best = (
            self.config.combo_ignore
            if self.config.combo_ignore is not None
            else self._select_best_modes(X, Y)
        )
        return self._cross_validated_predict_with_modes(
            X,
            Y,
            x_modes=best[0],
            y_modes=best[1],
            cca_modes=best[2],
            leave_n_out=leave_n_out,
        )

    def cross_validated_predict_tercile_proba(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        *,
        leave_n_out: int = 1,
    ) -> np.ndarray:
        X, Y = self._validate_inputs(X, Y)
        best = (
            self.config.combo_ignore
            if self.config.combo_ignore is not None
            else self._select_best_modes(X, Y)
        )

        probs = np.full((3, X.shape[0], Y.shape[1]), np.nan, dtype=float)

        for test_idx in self._iter_leave_n_out_indices(X.shape[0], leave_n_out=leave_n_out):
            train_idx = self._complement_indices(X.shape[0], test_idx)

            xtr = X[train_idx]
            ytr = Y[train_idx]

            self._fit_core(
                xtr,
                ytr,
                best[0],
                best[1],
                best[2],
                calc_loadings=False,
            )

            ytr_pred = self._cross_validated_predict_with_modes(
                xtr,
                ytr,
                x_modes=best[0],
                y_modes=best[1],
                cca_modes=best[2],
                leave_n_out=1,
            )

            dof = max(xtr.shape[0] - best[2] - 1, 1)
            self.hindcast_error_variance_ = np.sum((ytr_pred - ytr) ** 2, axis=0) / dof
            self._degrees_of_freedom = dof
            self._train_targets = ytr.copy()

            probs[:, test_idx, :] = self._predict_tercile_proba_from_fitted_core(X[test_idx])

        return probs

    def cross_validated_predict_xarray(
        self,
        X: xr.DataArray,
        Y: xr.DataArray,
        *,
        leave_n_out: int = 1,
        sample_dim: str = "time",
        x_domain_dims: Optional[Sequence[str]] = None,
        y_domain_dims: Optional[Sequence[str]] = None,
        x_mask: Optional[xr.DataArray] = None,
        y_mask: Optional[xr.DataArray] = None,
        auto_latitude_weights: bool = False,
        x_lat_dim: str = "lat",
        y_lat_dim: str = "lat",
    ) -> xr.DataArray:
        X2d, Y2d, metadata, x_weights, y_weights = self._prepare_xarray_pair(
            X,
            Y,
            sample_dim=sample_dim,
            x_domain_dims=x_domain_dims,
            y_domain_dims=y_domain_dims,
            x_mask=x_mask,
            y_mask=y_mask,
            auto_latitude_weights=auto_latitude_weights,
            x_lat_dim=x_lat_dim,
            y_lat_dim=y_lat_dim,
        )

        old_xw = self.config.x_weights
        old_yw = self.config.y_weights

        try:
            if auto_latitude_weights:
                self.config.x_weights = x_weights
                self.config.y_weights = y_weights

            preds = self.cross_validated_predict(X2d, Y2d, leave_n_out=leave_n_out)
            self._xarray_metadata = metadata

            sample_coords = {sample_dim: X[sample_dim].values}
            return self._rebuild_y_xarray(preds, sample_coords, name="cv_prediction")
        finally:
            self.config.x_weights = old_xw
            self.config.y_weights = old_yw

    def cross_validated_predict_tercile_proba_xarray(
        self,
        X: xr.DataArray,
        Y: xr.DataArray,
        *,
        leave_n_out: int = 1,
        sample_dim: str = "time",
        x_domain_dims: Optional[Sequence[str]] = None,
        y_domain_dims: Optional[Sequence[str]] = None,
        x_mask: Optional[xr.DataArray] = None,
        y_mask: Optional[xr.DataArray] = None,
        auto_latitude_weights: bool = False,
        x_lat_dim: str = "lat",
        y_lat_dim: str = "lat",
        category_dim: str = "category",
        category_labels: Sequence[str] = ("below", "near", "above"),
    ) -> xr.DataArray:
        X2d, Y2d, metadata, x_weights, y_weights = self._prepare_xarray_pair(
            X,
            Y,
            sample_dim=sample_dim,
            x_domain_dims=x_domain_dims,
            y_domain_dims=y_domain_dims,
            x_mask=x_mask,
            y_mask=y_mask,
            auto_latitude_weights=auto_latitude_weights,
            x_lat_dim=x_lat_dim,
            y_lat_dim=y_lat_dim,
        )

        old_xw = self.config.x_weights
        old_yw = self.config.y_weights

        try:
            if auto_latitude_weights:
                self.config.x_weights = x_weights
                self.config.y_weights = y_weights

            probs = self.cross_validated_predict_tercile_proba(X2d, Y2d, leave_n_out=leave_n_out)
            self._xarray_metadata = metadata

            sample_coords = {sample_dim: X[sample_dim].values}
            return self._rebuild_probability_xarray(
                probs,
                sample_coords,
                category_dim=category_dim,
                category_labels=category_labels,
                name="cv_probability",
            )
        finally:
            self.config.x_weights = old_xw
            self.config.y_weights = old_yw

    # ------------------------------------------------------------------
    # CCA diagnostics
    # ------------------------------------------------------------------
    def get_x_cca_loadings(self) -> np.ndarray:
        self._check_is_fitted()
        if not hasattr(self, "x_cca_loadings_"):
            raise RuntimeError("X CCA loadings are not available.")
        return self.x_cca_loadings_.copy()

    def get_y_cca_loadings(self) -> np.ndarray:
        self._check_is_fitted()
        if not hasattr(self, "y_cca_loadings_"):
            raise RuntimeError("Y CCA loadings are not available.")
        return self.y_cca_loadings_.copy()

    def get_x_cca_scores(self) -> np.ndarray:
        self._check_is_fitted()
        if not hasattr(self, "x_cca_scores_"):
            raise RuntimeError("X CCA scores are not available.")
        return self.x_cca_scores_.copy()

    def get_y_cca_scores(self) -> np.ndarray:
        self._check_is_fitted()
        if not hasattr(self, "y_cca_scores_"):
            raise RuntimeError("Y CCA scores are not available.")
        return self.y_cca_scores_.copy()

    def get_x_cca_loadings_xarray(self, mode_dim: str = "mode") -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()
        if not hasattr(self, "x_cca_loadings_"):
            raise RuntimeError("X CCA loadings are not available.")

        x_domain_dims = self._xarray_metadata["x_domain_dims"]
        valid_x = self._xarray_metadata["valid_x"]
        x_domain_index = self._xarray_metadata["x_domain_index"]

        full = np.full((x_domain_index.size, self.x_cca_loadings_.shape[1]), np.nan, dtype=float)
        full[valid_x, :] = self.x_cca_loadings_

        da = xr.DataArray(
            full,
            dims=("x_domain", mode_dim),
            coords={
                "x_domain": x_domain_index,
                mode_dim: np.arange(1, self.x_cca_loadings_.shape[1] + 1),
            },
            name="x_cca_loadings",
        )
        return da.unstack("x_domain").transpose(*x_domain_dims, mode_dim)

    def get_y_cca_loadings_xarray(self, mode_dim: str = "mode") -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()
        if not hasattr(self, "y_cca_loadings_"):
            raise RuntimeError("Y CCA loadings are not available.")

        y_domain_dims = self._xarray_metadata["y_domain_dims"]
        y_domain_index = self._xarray_metadata["y_domain_index"]
        valid_y = self._xarray_metadata["valid_y"]

        full = np.full((y_domain_index.size, self.y_cca_loadings_.shape[1]), np.nan, dtype=float)
        full[valid_y, :] = self.y_cca_loadings_

        da = xr.DataArray(
            full,
            dims=("y_domain", mode_dim),
            coords={
                "y_domain": y_domain_index,
                mode_dim: np.arange(1, self.y_cca_loadings_.shape[1] + 1),
            },
            name="y_cca_loadings",
        )
        return da.unstack("y_domain").transpose(*y_domain_dims, mode_dim)

    def get_x_cca_scores_xarray(self, mode_dim: str = "mode") -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()
        if not hasattr(self, "x_cca_scores_"):
            raise RuntimeError("X CCA scores are not available.")

        sample_dim = self._xarray_metadata["sample_dim"]
        sample_coords = self._xarray_metadata["sample_coords"]

        return xr.DataArray(
            self.x_cca_scores_,
            dims=(sample_dim, mode_dim),
            coords={
                sample_dim: sample_coords,
                mode_dim: np.arange(1, self.x_cca_scores_.shape[1] + 1),
            },
            name="x_cca_scores",
        )

    def get_y_cca_scores_xarray(self, mode_dim: str = "mode") -> xr.DataArray:
        self._check_is_fitted()
        self._require_xarray_fit()
        if not hasattr(self, "y_cca_scores_"):
            raise RuntimeError("Y CCA scores are not available.")

        sample_dim = self._xarray_metadata["sample_dim"]
        sample_coords = self._xarray_metadata["sample_coords"]

        return xr.DataArray(
            self.y_cca_scores_,
            dims=(sample_dim, mode_dim),
            coords={
                sample_dim: sample_coords,
                mode_dim: np.arange(1, self.y_cca_scores_.shape[1] + 1),
            },
            name="y_cca_scores",
        )

    @property
    def summary_(self) -> Dict[str, Any]:
        self._check_is_fitted()
        return {
            "x_modes_requested": self.x_modes_requested_,
            "y_modes_requested": self.y_modes_requested_,
            "cca_modes_requested": self.cca_modes_requested_,
            "x_modes_used": self.x_modes_,
            "y_modes_used": self.y_modes_,
            "cca_modes_used": self.cca_modes_,
            "canonical_correlations": self.canonical_correlations_.copy(),
            "x_explained_variance_fraction": self.x_explained_variance_fraction_.copy(),
            "y_explained_variance_fraction": self.y_explained_variance_fraction_.copy(),
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _require_xarray_fit(self) -> None:
        if not self._xarray_metadata:
            raise RuntimeError("This model was not fitted with fit_xarray(). Use the NumPy interface instead.")

    def _iter_leave_n_out_indices(self, n_samples: int, *, leave_n_out: int = 1):
        if leave_n_out < 1:
            raise ValueError("leave_n_out must be >= 1")
        for start in range(0, n_samples, leave_n_out):
            stop = min(start + leave_n_out, n_samples)
            yield np.arange(start, stop)

    def _complement_indices(self, n_samples: int, held_out_idx: np.ndarray) -> np.ndarray:
        mask = np.ones(n_samples, dtype=bool)
        mask[held_out_idx] = False
        return np.where(mask)[0]

    def _validate_inputs(self, X: np.ndarray, Y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        X = np.asarray(X, dtype=float)
        Y = np.asarray(Y, dtype=float)

        if X.ndim != 2 or Y.ndim != 2:
            raise ValueError("X and Y must both be 2D arrays of shape (n_samples, n_domains)")
        if X.shape[0] != Y.shape[0]:
            raise ValueError("X and Y must have the same number of samples")
        if X.shape[0] < 3:
            raise ValueError("At least 3 samples are required")
        if X.shape[1] < 1 or Y.shape[1] < 1:
            raise ValueError("X and Y must each contain at least one domain")
        if np.any(~np.isfinite(X)) or np.any(~np.isfinite(Y)):
            raise ValueError("X and Y must not contain NaN/Inf values. Mask or remove invalid domains first.")

        self._validate_weights(X.shape[1], Y.shape[1])
        return X, Y

    def _validate_predictors(self, X: np.ndarray) -> np.ndarray:
        X = np.asarray(X, dtype=float)
        if X.ndim != 2:
            raise ValueError("X must be a 2D array of shape (n_samples, n_domains)")
        if X.shape[1] != self.x_mean_.shape[0]:
            raise ValueError(f"X has {X.shape[1]} domains but model expects {self.x_mean_.shape[0]}")
        if np.any(~np.isfinite(X)):
            raise ValueError("Prediction X contains NaN/Inf values")
        return X

    def _validate_weights(self, n_x_domains: int, n_y_domains: int) -> None:
        if self.config.x_weights is not None:
            xw = np.asarray(self.config.x_weights, dtype=float)
            if xw.shape != (n_x_domains,):
                raise ValueError("x_weights must have shape (n_x_domains,)")
            if np.any(~np.isfinite(xw)):
                raise ValueError("x_weights must be finite")

        if self.config.y_weights is not None:
            yw = np.asarray(self.config.y_weights, dtype=float)
            if yw.shape != (n_y_domains,):
                raise ValueError("y_weights must have shape (n_y_domains,)")
            if np.any(~np.isfinite(yw)):
                raise ValueError("y_weights must be finite")

    def _validate_mask_for_dims(
        self,
        mask: xr.DataArray,
        domain_dims: Sequence[str],
        mask_name: str,
    ) -> None:
        if not isinstance(mask, xr.DataArray):
            raise TypeError(f"{mask_name} must be an xarray.DataArray")
        missing = [d for d in domain_dims if d not in mask.dims]
        if missing:
            raise ValueError(
                f"{mask_name} is missing required domain dims {missing}. Found dims={mask.dims}"
            )

    def _prepare_xarray_pair(
        self,
        X: xr.DataArray,
        Y: xr.DataArray,
        *,
        sample_dim: str,
        x_domain_dims: Optional[Sequence[str]],
        y_domain_dims: Optional[Sequence[str]],
        x_mask: Optional[xr.DataArray],
        y_mask: Optional[xr.DataArray],
        auto_latitude_weights: bool,
        x_lat_dim: str,
        y_lat_dim: str,
    ) -> Tuple[np.ndarray, np.ndarray, Dict[str, Any], Optional[np.ndarray], Optional[np.ndarray]]:
        if not isinstance(X, xr.DataArray) or not isinstance(Y, xr.DataArray):
            raise TypeError("fit_xarray expects xarray.DataArray inputs")
        if sample_dim not in X.dims or sample_dim not in Y.dims:
            raise ValueError(f"sample_dim={sample_dim!r} must exist in both X and Y")

        x_domain_dims = (
            tuple(x_domain_dims)
            if x_domain_dims is not None
            else tuple(d for d in X.dims if d != sample_dim)
        )
        y_domain_dims = (
            tuple(y_domain_dims)
            if y_domain_dims is not None
            else tuple(d for d in Y.dims if d != sample_dim)
        )

        if len(x_domain_dims) == 0:
            raise ValueError("X must have at least one domain dimension besides sample_dim")
        if len(y_domain_dims) == 0:
            raise ValueError("Y must have at least one domain dimension besides sample_dim")

        missing_x = [d for d in x_domain_dims if d not in X.dims]
        missing_y = [d for d in y_domain_dims if d not in Y.dims]
        if missing_x:
            raise ValueError(f"X is missing domain dims: {missing_x}")
        if missing_y:
            raise ValueError(f"Y is missing domain dims: {missing_y}")

        if x_mask is not None:
            self._validate_mask_for_dims(x_mask, x_domain_dims, "x_mask")
        if y_mask is not None:
            self._validate_mask_for_dims(y_mask, y_domain_dims, "y_mask")

        X_t = X.transpose(sample_dim, *x_domain_dims)
        Y_t = Y.transpose(sample_dim, *y_domain_dims)

        X_stacked = X_t.stack(x_domain=x_domain_dims)
        Y_stacked = Y_t.stack(y_domain=y_domain_dims)

        X_values = np.asarray(X_stacked.values, dtype=float)
        Y_values = np.asarray(Y_stacked.values, dtype=float)

        if x_mask is not None:
            valid_x = np.asarray(
                x_mask.transpose(*x_domain_dims).stack(x_domain=x_domain_dims).values,
                dtype=bool,
            )
        else:
            valid_x = np.all(np.isfinite(X_values), axis=0)

        if y_mask is not None:
            valid_y = np.asarray(
                y_mask.transpose(*y_domain_dims).stack(y_domain=y_domain_dims).values,
                dtype=bool,
            )
        else:
            valid_y = np.all(np.isfinite(Y_values), axis=0)

        valid_x = valid_x & np.all(np.isfinite(X_values), axis=0)
        valid_y = valid_y & np.all(np.isfinite(Y_values), axis=0)

        if not np.any(valid_x):
            raise ValueError("No valid X domains remain after masking and finite-value screening.")
        if not np.any(valid_y):
            raise ValueError("No valid Y domains remain after masking and finite-value screening.")

        X2d = X_values[:, valid_x]
        Y2d = Y_values[:, valid_y]

        if np.any(~np.isfinite(X2d)) or np.any(~np.isfinite(Y2d)):
            raise ValueError("After masking/stacking, X or Y still contains NaN/Inf values.")

        x_weights = None
        y_weights = None
        if auto_latitude_weights:
            x_weights = self._make_latitude_weights_for_stacked(X_t, x_domain_dims, valid_x, x_lat_dim)
            y_weights = self._make_latitude_weights_for_stacked(Y_t, y_domain_dims, valid_y, y_lat_dim)

        metadata = {
            "sample_dim": sample_dim,
            "x_domain_dims": x_domain_dims,
            "y_domain_dims": y_domain_dims,
            "valid_x": valid_x,
            "valid_y": valid_y,
            "x_domain_index": X_stacked["x_domain"].copy(),
            "y_domain_index": Y_stacked["y_domain"].copy(),
            "y_template": Y_t.isel({sample_dim: 0}, drop=True).copy(),
            "y_stack_name": "y_domain",
            "sample_coords": X[sample_dim].values,
        }

        return X2d, Y2d, metadata, x_weights, y_weights

    def _make_latitude_weights_for_stacked(
        self,
        da: xr.DataArray,
        domain_dims: Sequence[str],
        valid_mask: np.ndarray,
        lat_dim: str,
    ) -> np.ndarray:
        if lat_dim not in da.dims:
            raise ValueError(
                f"auto_latitude_weights=True but lat_dim={lat_dim!r} not found in dims {da.dims}"
            )

        sample_dim = self._infer_sample_dim(da)
        template = da.isel({sample_dim: 0}, drop=True)

        lat_radians = np.deg2rad(np.asarray(template[lat_dim].values, dtype=float))
        lat_weights_1d = np.sqrt(np.clip(np.cos(lat_radians), 0.0, None))

        weight_da = xr.DataArray(
            lat_weights_1d,
            dims=(lat_dim,),
            coords={lat_dim: template[lat_dim].values},
        )
        weight_da = weight_da.broadcast_like(template)
        stacked = weight_da.transpose(*domain_dims).stack(domain=domain_dims)
        weights = np.asarray(stacked.values, dtype=float)[valid_mask]

        if np.any(~np.isfinite(weights)):
            raise ValueError("Latitude weights contain NaN/Inf after stacking.")
        return weights

    def _infer_sample_dim(self, da: xr.DataArray) -> str:
        if "time" in da.dims:
            return "time"
        return da.dims[0]

    def _flatten_new_xarray_predictor(
        self,
        X: xr.DataArray,
        *,
        sample_dim: Optional[str],
        x_domain_dims: Optional[Sequence[str]],
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        if not isinstance(X, xr.DataArray):
            raise TypeError("Expected an xarray.DataArray")

        sample_dim = sample_dim or self._xarray_metadata["sample_dim"]
        x_domain_dims = (
            tuple(x_domain_dims)
            if x_domain_dims is not None
            else self._xarray_metadata["x_domain_dims"]
        )

        if sample_dim not in X.dims:
            raise ValueError(f"sample_dim={sample_dim!r} not found in X dims={X.dims}")

        missing = [d for d in x_domain_dims if d not in X.dims]
        if missing:
            raise ValueError(f"Prediction X is missing expected domain dims: {missing}")

        X_t = X.transpose(sample_dim, *x_domain_dims)
        X_stacked = X_t.stack(x_domain=x_domain_dims)
        valid_x = self._xarray_metadata["valid_x"]

        if X_stacked.sizes["x_domain"] != valid_x.size:
            raise ValueError(
                "Prediction X domain layout does not match the fitted layout. "
                "Check domain dims, coordinates, and preprocessing."
            )

        X2d = np.asarray(X_stacked.values[:, valid_x], dtype=float)
        if np.any(~np.isfinite(X2d)):
            raise ValueError("Prediction X contains NaN/Inf values in retained domains")

        return X2d, {sample_dim: X_t[sample_dim].values}

    def _rebuild_y_xarray(
        self,
        Y2d: np.ndarray,
        sample_coords: Dict[str, Any],
        *,
        name: str,
    ) -> xr.DataArray:
        sample_dim = self._xarray_metadata["sample_dim"]
        y_domain_dims = self._xarray_metadata["y_domain_dims"]
        y_stack_name = self._xarray_metadata["y_stack_name"]
        y_domain_index = self._xarray_metadata["y_domain_index"]
        valid_y = self._xarray_metadata["valid_y"]
        y_template = self._xarray_metadata["y_template"]

        full = np.full((Y2d.shape[0], y_domain_index.size), np.nan, dtype=float)
        full[:, valid_y] = Y2d

        da = xr.DataArray(
            full,
            dims=(sample_dim, y_stack_name),
            coords={
                sample_dim: sample_coords[sample_dim],
                y_stack_name: y_domain_index,
            },
            name=name,
        )
        da = da.unstack(y_stack_name).transpose(sample_dim, *y_domain_dims)
        da.attrs.update(y_template.attrs)
        return da

    def _rebuild_probability_xarray(
        self,
        probs: np.ndarray,
        sample_coords: Dict[str, Any],
        *,
        category_dim: str,
        category_labels: Sequence[str],
        name: str,
    ) -> xr.DataArray:
        if probs.shape[0] != len(category_labels):
            raise ValueError(
                f"Number of category labels ({len(category_labels)}) does not match "
                f"probability array first dimension ({probs.shape[0]})."
            )

        sample_dim = self._xarray_metadata["sample_dim"]
        y_domain_dims = self._xarray_metadata["y_domain_dims"]
        y_stack_name = self._xarray_metadata["y_stack_name"]
        y_domain_index = self._xarray_metadata["y_domain_index"]
        valid_y = self._xarray_metadata["valid_y"]

        full = np.full((probs.shape[0], probs.shape[1], y_domain_index.size), np.nan, dtype=float)
        full[:, :, valid_y] = probs

        da = xr.DataArray(
            full,
            dims=(category_dim, sample_dim, y_stack_name),
            coords={
                category_dim: list(category_labels),
                sample_dim: sample_coords[sample_dim],
                y_stack_name: y_domain_index,
            },
            name=name,
        )
        return da.unstack(y_stack_name).transpose(category_dim, sample_dim, *y_domain_dims)

    def _apply_weights(self, arr: np.ndarray, weights: Optional[np.ndarray]) -> np.ndarray:
        return arr if weights is None else arr * weights[None, :]

    def _remove_weights(self, arr: np.ndarray, weights: Optional[np.ndarray]) -> np.ndarray:
        if weights is None:
            return arr
        safe_weights = np.where(np.abs(weights) < self.config.epsilon, np.nan, weights)
        return arr / safe_weights[None, :]

    def _standardize_X(self, X: np.ndarray) -> np.ndarray:
        Xs = (X - self.x_mean_[None, :]) / self.x_std_[None, :]
        return self._apply_weights(Xs, self.config.x_weights)

    def _standardize_Y(self, Y: np.ndarray) -> np.ndarray:
        Ys = (Y - self.y_mean_[None, :]) / self.y_std_[None, :]
        return self._apply_weights(Ys, self.config.y_weights)

    def _destandardize_Y(self, Ys_weighted: np.ndarray) -> np.ndarray:
        Ys = self._remove_weights(Ys_weighted, self.config.y_weights)
        return Ys * self.y_std_[None, :] + self.y_mean_[None, :]

    def _flip_rows(self, Vt: np.ndarray) -> np.ndarray:
        Vt = Vt.copy()
        for i in range(Vt.shape[0]):
            row = Vt[i, :]
            finite = row[np.isfinite(row)]
            if finite.size == 0:
                continue
            if np.abs(np.nanmax(finite)) < np.abs(np.nanmin(finite)):
                Vt[i, :] *= -1.0
        return Vt

    def _safe_standard_deviation(self, arr: np.ndarray, name: str) -> np.ndarray:
        std = np.nanstd(arr, axis=0, ddof=0)
        bad = std < self.config.epsilon
        if np.any(bad):
            n_bad = int(np.sum(bad))
            raise ValueError(f"{name} contains {n_bad} domain(s) with near-zero standard deviation.")
        return std

    def _select_best_modes(self, X: np.ndarray, Y: np.ndarray) -> Tuple[int, int, int]:
        best_score = -np.inf
        best_combo = None
        counter = 0

        self._log("[CCA] Selecting best modes...", 1)

        for x_modes in range(self.config.x_eof_range[0], self.config.x_eof_range[1] + 1):
            for y_modes in range(self.config.y_eof_range[0], self.config.y_eof_range[1] + 1):
                max_cca = min(x_modes, y_modes, self.config.cca_mode_range[1])
                min_cca = self.config.cca_mode_range[0]
                if min_cca > max_cca:
                    continue

                for cca_modes in range(min_cca, max_cca + 1):
                    counter += 1
                    self._log(
                        f"[CCA] Testing combo {counter}: x={x_modes}, y={y_modes}, cca={cca_modes}",
                        2,
                    )
                    score = self._cross_validated_skill(X, Y, x_modes, y_modes, cca_modes)
                    self._log(f"      score={score:.4f}", 2)
                    if score > best_score:
                        best_score = score
                        best_combo = (x_modes, y_modes, cca_modes)

        if best_combo is None:
            raise RuntimeError("Could not find a valid combination of modes")

        self._log(f"[CCA] Best combo: {best_combo}, score={best_score:.4f}", 1)
        return best_combo

    def _cross_validated_skill(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        x_modes: int,
        y_modes: int,
        cca_modes: int,
    ) -> float:
        n_samples = X.shape[0]
        n_splits = min(self.config.cv_splits, n_samples)
        if n_splits < 2:
            raise ValueError("Need at least 2 splits for cross-validation")

        kf = KFold(
            n_splits=n_splits,
            shuffle=self.config.shuffle,
            random_state=self.config.random_state if self.config.shuffle else None,
        )

        preds = []
        obs = []
        for train_idx, test_idx in kf.split(np.arange(n_samples)):
            self._fit_core(X[train_idx], Y[train_idx], x_modes, y_modes, cca_modes, calc_loadings=False)
            preds.append(self._predict_from_fitted_core(X[test_idx]))
            obs.append(Y[test_idx])

        preds = np.vstack(preds)
        obs = np.vstack(obs)

        score, _ = spearmanr(preds.ravel(), obs.ravel())
        if not np.isfinite(score):
            return -np.inf
        return float(score)

    def _cross_validated_error_variance(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        modes: Tuple[int, int, int],
    ) -> None:
        preds = self._cross_validated_predict_with_modes(
            X,
            Y,
            x_modes=modes[0],
            y_modes=modes[1],
            cca_modes=modes[2],
            leave_n_out=1,
        )
        resid = preds - Y
        dof = max(X.shape[0] - modes[2] - 1, 1)
        self.hindcast_error_variance_ = np.sum(resid ** 2, axis=0) / dof
        self.cv_error_variance_ = np.nanvar(resid, axis=0, ddof=0)
        self._degrees_of_freedom = dof

    def _cross_validated_predict_with_modes(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        *,
        x_modes: int,
        y_modes: int,
        cca_modes: int,
        leave_n_out: int = 1,
    ) -> np.ndarray:
        preds = np.full(Y.shape, np.nan, dtype=float)

        for test_idx in self._iter_leave_n_out_indices(X.shape[0], leave_n_out=leave_n_out):
            train_idx = self._complement_indices(X.shape[0], test_idx)

            self._fit_core(
                X[train_idx],
                Y[train_idx],
                x_modes,
                y_modes,
                cca_modes,
                calc_loadings=False,
            )
            preds[test_idx, :] = self._predict_from_fitted_core(X[test_idx])

        return preds

    def _fit_core(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        x_modes_requested: int,
        y_modes_requested: int,
        cca_modes_requested: int,
        calc_loadings: bool,
    ) -> None:
        self.x_mean_ = np.nanmean(X, axis=0)
        self.y_mean_ = np.nanmean(Y, axis=0)
        self.x_std_ = self._safe_standard_deviation(X, "X")
        self.y_std_ = self._safe_standard_deviation(Y, "Y")

        Xs = self._standardize_X(X)
        Ys = self._standardize_Y(Y)

        _, sx, Vtx = svd(Xs, full_matrices=False)
        Vtx = self._flip_rows(Vtx)
        sx = sx[sx > self.config.epsilon]
        if sx.size == 0:
            raise RuntimeError("No valid singular values found in X")

        self.x_modes_ = min(x_modes_requested, sx.size)
        self.x_eof_singular_values_ = sx[: self.x_modes_]
        self.x_eof_loadings_ = Vtx[: self.x_modes_, :]
        self.x_explained_variance_fraction_ = sx[: self.x_modes_] / sx.sum()
        x_scores = Xs @ self.x_eof_loadings_.T

        _, sy, Vty = svd(Ys, full_matrices=False)
        Vty = self._flip_rows(Vty)
        sy = sy[sy > self.config.epsilon]
        if sy.size == 0:
            raise RuntimeError("No valid singular values found in Y")

        self.y_modes_ = min(y_modes_requested, sy.size)
        self.y_eof_singular_values_ = sy[: self.y_modes_]
        self.y_eof_loadings_ = Vty[: self.y_modes_, :]
        self.y_explained_variance_fraction_ = sy[: self.y_modes_] / sy.sum()
        y_scores = Ys @ self.y_eof_loadings_.T

        x_pc = x_scores / self.x_eof_singular_values_
        y_pc = y_scores / self.y_eof_singular_values_

        C = x_pc.T @ y_pc
        Uc, sc, Vtc = svd(C, full_matrices=False)
        sc = sc[sc > self.config.epsilon]
        if sc.size == 0:
            raise RuntimeError("No valid canonical singular values found")

        self.cca_modes_ = min(cca_modes_requested, sc.size, self.x_modes_, self.y_modes_)
        self.cca_left_ = Uc[:, : self.cca_modes_]
        self.cca_right_ = Vtc[: self.cca_modes_, :]
        self.canonical_correlations_ = sc[: self.cca_modes_]

        if calc_loadings:
            self.x_cca_scores_ = x_pc @ self.cca_left_
            self.y_cca_scores_ = y_pc @ self.cca_right_.T
            self.x_cca_loadings_ = self.x_eof_loadings_.T @ (
                np.diag(self.x_eof_singular_values_) @ self.cca_left_
            )
            self.y_cca_loadings_ = self.y_eof_loadings_.T @ (
                np.diag(self.y_eof_singular_values_) @ self.cca_right_.T
            )

    def _predict_from_fitted_core(self, X: np.ndarray) -> np.ndarray:
        X = self._validate_predictors(X)

        Xs = self._standardize_X(X)
        x_scores = Xs @ self.x_eof_loadings_.T
        x_pc = x_scores / self.x_eof_singular_values_
        x_cca = x_pc @ self.cca_left_

        y_pc = x_cca @ self.cca_right_
        y_scores = y_pc * self.y_eof_singular_values_
        Ys = y_scores @ self.y_eof_loadings_

        return self._destandardize_Y(Ys)

    def _prediction_error_variance_from_fitted_core(self, X: np.ndarray) -> np.ndarray:
        X = self._validate_predictors(X)

        Xs = self._standardize_X(X)
        x_scores = Xs @ self.x_eof_loadings_.T
        x_pc = x_scores / self.x_eof_singular_values_
        x_cca = x_pc @ self.cca_left_

        method = self.config.probability_method

        if method == "constant_error":
            base = self.hindcast_error_variance_

        elif method == "cv_error":
            base = getattr(self, "cv_error_variance_", self.hindcast_error_variance_)

        elif method == "inflated_error":
            base = self.config.error_inflation_factor * self.hindcast_error_variance_

        elif method == "adjusted_error":
            leverage = np.sum(x_cca ** 2, axis=1, keepdims=True)
            return (1.0 + leverage) * self.hindcast_error_variance_[None, :]

        else:
            raise ValueError(f"Unknown probability_method: {method!r}")

        return np.broadcast_to(base[None, :], (X.shape[0], base.shape[0])).copy()

    def _predict_quantile_exceedance_from_fitted_core(
        self,
        X: np.ndarray,
        quantiles: Union[float, Sequence[float]],
    ) -> np.ndarray:
        X = self._validate_predictors(X)
        mu = self._predict_from_fitted_core(X)
        pred_var = self._prediction_error_variance_from_fitted_core(X)
        dof = self._degrees_of_freedom

        if np.isscalar(quantiles):
            quantiles = [float(quantiles)]

        out = []
        for q in quantiles:
            if not 0 < q < 1:
                raise ValueError(f"Each quantile must be between 0 and 1, got {q}")
            threshold = np.nanquantile(self._train_targets, q, axis=0)
            prob = t.sf(threshold, dof, loc=mu, scale=np.sqrt(pred_var))
            out.append(prob)

        return np.stack(out, axis=0)

    def _predict_tercile_proba_from_fitted_core(
        self,
        X: np.ndarray,
        lower_threshold: Optional[np.ndarray] = None,
        upper_threshold: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        X = self._validate_predictors(X)
        mu = self._predict_from_fitted_core(X)
        pred_var = self._prediction_error_variance_from_fitted_core(X)
        dof = self._degrees_of_freedom

        if lower_threshold is None:
            lower_threshold = np.nanquantile(self._train_targets, 1.0 / 3.0, axis=0)
        else:
            lower_threshold = np.asarray(lower_threshold, dtype=float)

        if upper_threshold is None:
            upper_threshold = np.nanquantile(self._train_targets, 2.0 / 3.0, axis=0)
        else:
            upper_threshold = np.asarray(upper_threshold, dtype=float)

        if lower_threshold.shape != (mu.shape[1],):
            raise ValueError(
                f"lower_threshold must have shape ({mu.shape[1]},), got {lower_threshold.shape}"
            )
        if upper_threshold.shape != (mu.shape[1],):
            raise ValueError(
                f"upper_threshold must have shape ({mu.shape[1]},), got {upper_threshold.shape}"
            )

        below = t.cdf(lower_threshold, dof, loc=mu, scale=np.sqrt(pred_var))
        above = t.sf(upper_threshold, dof, loc=mu, scale=np.sqrt(pred_var))
        near = 1.0 - below - above
        near = np.clip(near, 0.0, 1.0)

        probs = np.stack([below, near, above], axis=0)
        probs = np.clip(probs, 0.0, 1.0)

        denom = np.sum(probs, axis=0, keepdims=True)
        bad = denom <= self.config.epsilon
        denom = np.where(bad, 1.0, denom)
        probs = probs / denom

        return probs

    def _check_is_fitted(self) -> None:
        if not self._is_fitted:
            raise RuntimeError(
                "This CCAForecaster is not fitted yet. Call fit(X, Y) or fit_xarray(X, Y) first."
            )


def fit_cca_forecaster(
    X: np.ndarray,
    Y: np.ndarray,
    *,
    x_eof_range: Tuple[int, int] = (1, 5),
    y_eof_range: Tuple[int, int] = (1, 5),
    cca_mode_range: Tuple[int, int] = (1, 5),
    cv_splits: int = 5,
    leave_n_out: int = 1,
    probability_method: str = "constant_error",
    error_inflation_factor: float = 1.0,
    x_weights: Optional[np.ndarray] = None,
    y_weights: Optional[np.ndarray] = None,
    combo_ignore: Optional[Tuple[int, int, int]] = None,
    shuffle: bool = False,
    random_state: Optional[int] = None,
    epsilon: float = 1e-8,
    selection_metric: str = "spearman",
    top_n_models: int = 1,
    verbose: int = 1,
) -> CCAForecaster:
    config = CCAConfig(
        x_eof_range=x_eof_range,
        y_eof_range=y_eof_range,
        cca_mode_range=cca_mode_range,
        cv_splits=cv_splits,
        leave_n_out=leave_n_out,
        shuffle=shuffle,
        random_state=random_state,
        probability_method=probability_method,
        error_inflation_factor=error_inflation_factor,
        x_weights=x_weights,
        y_weights=y_weights,
        epsilon=epsilon,
        combo_ignore=combo_ignore,
        selection_metric=selection_metric,
        top_n_models=top_n_models,
        verbose=verbose,
    )
    return CCAForecaster(config=config).fit(X, Y)
