from __future__ import annotations

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

import xeofs as xe

from scipy import stats, signal as sig
from scipy.optimize import fsolve
from scipy.special import gamma as gamma_function
from scipy.stats import (
    gamma, norm, lognorm, t, poisson, nbinom,
    expon, weibull_min
)

# Project-specific utilities
from wass2s.utils import *

# CCA-Zoo linear and nonparametric models (optional)
try:
    from cca_zoo.linear import CCA as LinearCCA, GCCA
    from cca_zoo.nonparametric import KCCA, KGCCA
    from cca_zoo.model_selection import GridSearchCV as CCA_GridSearch
    from sklearn.model_selection import KFold
    HAS_CCA_ZOO = True
except ImportError:
    HAS_CCA_ZOO = False
    print("Warning: cca-zoo not installed. CCA/KCCA/KGCCA functionality disabled.")

# CCA-Zoo deep models (optional)
try:
    from cca_zoo.deep import DCCA
    from cca_zoo.deep.architectures import Encoder
    HAS_DCCA = True
except ImportError:
    HAS_DCCA = False
    print("Warning: cca-zoo deep modules not found. DCCA functionality disabled.")
    
class WAS_CCA:
    def __init__(self, n_modes=4, n_pca_modes=8, standardize=False, use_coslat=True, use_pca=True, dist_method="nonparam"):
        """
        Initialize the WAS_CCA class with specified parameters.

        Parameters:
        - n_modes: Number of canonical modes to compute.
        - n_pca_modes: Number of PCA modes to use before CCA.
        - standardize: Whether to standardize the data. Keep it False in our case data already standardize
        - use_coslat: Whether to use cosine latitude weighting.
        - use_pca: Whether to perform PCA before CCA.
        - detrend: Whether to apply detrending to the data.
        """
        
        self.n_modes = n_modes
        self.n_pca_modes = n_pca_modes
        self.standardize = standardize
        self.use_coslat = use_coslat
        self.use_pca = use_pca
        self.dist_method = dist_method

        self.cca = xe.cross.CCA(
            n_modes=self.n_modes,
            standardize=self.standardize,
            use_coslat=self.use_coslat,
            use_pca=self.use_pca,
            n_pca_modes=self.n_pca_modes
        )
        self.cca_model = None
    
    def fit_cca(self, X_train, y_train):
        """
        Fit the CCA model using the training data.

        Parameters:
        - X_train: xarray DataArray for predictor training data.
        - y_train: xarray DataArray for predictand training data.
        """
        # Preprocess the data
        X_train_final, y_train_final = self.preprocess_data(X_train, y_train)
        # Fit the CCA model
        self.cca_model = self.cca.fit(X_train_final, y_train_final, dim="T")

    # def detrend_data(self, data): ### Replace detrend by extendedEOF after
    #     """
    #     Detrend the data along the 'T' dimension if detrending is enabled.

    #     Parameters:
    #     - data: xarray DataArray to be detrended.

    #     Returns:
    #     - data_detrended: Detrended xarray DataArray.
    #     """
    #     if self.detrend:
    #         # Create mask for missing data
    #         mask = xr.where(np.isnan(data), np.nan, 1)
    #         # mask = np.where(np.isnan(data.isel(T=0)), np.nan, 1)
    #         # Fill missing values with zero before detrending
    #         data_filled = data.fillna(0)
    #         # Detrend data along 'T' axis (axis=0)
    #         data_detrended = sig.detrend(data_filled, axis=0)
    #         data_detrended = xr.DataArray(data_detrended, dims=data.dims, coords=data.coords)
    #         # Apply mask after detrending
    #         # data_detrended = data_detrended * mask
    #         return data_detrended
    #     else:
    #         return data

    # def preprocess_data(self, X, Y):
    #     """
    #     Preprocess the data by detrending, masking, and filling missing values.

    #     Parameters:
    #     - X: xarray DataArray for predictors.
    #     - Y: xarray DataArray for predictands.

    #     Returns:
    #     - X_final: Preprocessed X data.
    #     - Y_final: Preprocessed Y data.
    #     """
    #     # Apply detrending and masking (masking is now inside detrend_data)
    #     X_processed = self.detrend_data(X)
    #     # Fill missing values with mean along 'T'
    #     X_final = X_processed.fillna(X_processed.mean(dim="T", skipna=True))

    #     # Process Y data (assuming we do not detrend Y)
    #     Y_final = Y.fillna(Y.mean(dim="T", skipna=True))

    #     # Rename dimensions and transpose
    #     dims_rename = {"X": "lon", "Y": "lat"}
    #     X_final = X_final.rename(dims_rename).transpose('T', 'lat', 'lon')
    #     Y_final = Y_final.rename(dims_rename).transpose('T', 'lat', 'lon')

    #     return X_final, Y_final

    def preprocess_data(self, X, Y):
        """
        Preprocess the data by detrending, masking, and filling missing values.

        Parameters:
        - X: xarray DataArray for predictors.
        - Y: xarray DataArray for predictands.

        Returns:
        - X_final: Preprocessed X data.
        - Y_final: Preprocessed Y data.
        """
        # Apply detrending to both X and Y
        X_processed = X #- self.detrend_data(X)
        Y_processed = Y #- self.detrend_data(Y)

        # Fill missing values with mean along 'T'
        X_final = X_processed.fillna(X_processed.mean(dim="T", skipna=True))
        Y_final = Y_processed.fillna(Y_processed.mean(dim="T", skipna=True))

        # Rename dimensions and transpose
        dims_rename = {"X": "lon", "Y": "lat"}
        X_final = X_final.rename(dims_rename).transpose('T', 'lat', 'lon')
        Y_final = Y_final.rename(dims_rename).transpose('T', 'lat', 'lon')

        return X_final, Y_final

    def preprocess_test_data(self, X_test, y_test, X_train, y_train):
        """
        Preprocess the test data.

        Parameters:
        - X_test: xarray DataArray for predictor testing data.
        - y_test: xarray DataArray for predictand testing data.
        - X_train: xarray DataArray for predictor training data.
        - y_train: xarray DataArray for predictand training data.

        Returns:
        - X_test_prepared: Preprocessed X test data.
        - y_test_prepared: Preprocessed Y test data.
        """
        # Apply detrending and masking
        X_test_processed = X_test #- self.detrend_data(X_train).mean(dim="T", skipna=True)
        y_test_processed = y_test 
        
        # Fill missing values with mean from training data along 'T'
        X_test_prepared = X_test_processed.fillna(X_train.mean(dim="T", skipna=True))        
        y_test_prepared = y_test_processed.fillna(y_train.mean(dim="T", skipna=True))

        # Rename dimensions and transpose
        dims_rename = {"X": "lon", "Y": "lat"}
        X_test_prepared = X_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')
        y_test_prepared = y_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')

        return X_test_prepared, y_test_prepared

    def compute_model(self, X_train, y_train, X_test, y_test):
        """
        Compute the CCA model and generate hindcasts.

        Parameters:
        - X_train: xarray DataArray for predictor training data.
        - y_train: xarray DataArray for predictand training data.
        - X_test: xarray DataArray for predictor testing data.
        - y_test: xarray DataArray for predictand testing data.

        Returns:
        - hindcast: xarray DataArray containing predictions and errors.
        """
        
        
        # Fit the CCA model
        
        self.fit_cca(X_train, y_train)
            
        # Prepare test data
        X_test_prepared, y_test_prepared = self.preprocess_test_data(X_test, y_test, X_train, y_train)

        # Predict
        y_pred = self.cca_model.predict(X_test_prepared)
        y_pred = self.cca_model.inverse_transform(self.cca_model.transform(X_test_prepared), y_pred)[1] 

        y_pred['T'] = y_test_prepared['T']
        # Calculate error
        # error = y_test_prepared - y_pred

        # Combine prediction and error into a DataArray
        # hindcast = xr.concat([error, y_pred], dim="output")
        hindcast = y_pred.rename({"lon": "X", "lat": "Y"})
        # hindcast = hindcast.assign_coords(output=['error', 'prediction'])

        return hindcast

    # --------------------------------------------------------------------------
    #  Probability Calculation Methods
    # --------------------------------------------------------------------------
    # ------------------ Probability Calculation Methods ------------------

    @staticmethod
    def _ppf_terciles_from_code(dist_code, shape, loc, scale):
        """
        Return tercile thresholds (T1, T2) from best-fit distribution parameters.
    
        dist_code:
            1: norm
            2: lognorm
            3: expon
            4: gamma
            5: weibull_min
            6: t
            7: poisson
            8: nbinom
        """
        if np.isnan(dist_code):
            return np.nan, np.nan
    
        code = int(dist_code)
        try:
            if code == 1:
                return (
                    norm.ppf(0.32, loc=loc, scale=scale),
                    norm.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 2:
                return (
                    lognorm.ppf(0.32, s=shape, loc=loc, scale=scale),
                    lognorm.ppf(0.67, s=shape, loc=loc, scale=scale),
                )
            elif code == 3:
                return (
                    expon.ppf(0.32, loc=loc, scale=scale),
                    expon.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 4:
                return (
                    gamma.ppf(0.32, a=shape, loc=loc, scale=scale),
                    gamma.ppf(0.67, a=shape, loc=loc, scale=scale),
                )
            elif code == 5:
                return (
                    weibull_min.ppf(0.32, c=shape, loc=loc, scale=scale),
                    weibull_min.ppf(0.67, c=shape, loc=loc, scale=scale),
                )
            elif code == 6:
                # Note: Renamed 't_dist' to 't' for standard scipy.stats
                return (
                    t.ppf(0.32, df=shape, loc=loc, scale=scale),
                    t.ppf(0.67, df=shape, loc=loc, scale=scale),
                )
            elif code == 7:
                # Poisson: poisson.ppf(q, mu, loc=0)
                # ASSUMPTION: 'mu' (mean) is passed as 'shape'
                #             'loc' is passed as 'loc'
                #             'scale' is unused
                return (
                    poisson.ppf(0.32, mu=shape, loc=loc),
                    poisson.ppf(0.67, mu=shape, loc=loc),
                )
            elif code == 8:
                # Negative Binomial: nbinom.ppf(q, n, p, loc=0)
                # ASSUMPTION: 'n' (successes) is passed as 'shape'
                #             'p' (probability) is passed as 'scale'
                #             'loc' is passed as 'loc'
                return (
                    nbinom.ppf(0.32, n=shape, p=scale, loc=loc),
                    nbinom.ppf(0.67, n=shape, p=scale, loc=loc),
                )
        except Exception:
            return np.nan, np.nan
    
        # Fallback if code is not 1-8
        return np.nan, np.nan
        
    @staticmethod
    def weibull_shape_solver(k, M, V):
        """
        Function to find the root of the Weibull shape parameter 'k'.
        We find 'k' such that the theoretical variance/mean^2 ratio
        matches the observed V/M^2 ratio.
        """
        # Guard against invalid 'k' values during solving
        if k <= 0:
            return -np.inf
        try:
            g1 = gamma_function(1 + 1/k)
            g2 = gamma_function(1 + 2/k)
            
            # This is the V/M^2 ratio *implied by k*
            implied_v_over_m_sq = (g2 / (g1**2)) - 1
            
            # This is the *observed* ratio
            observed_v_over_m_sq = V / (M**2)
            
            # Return the difference (we want this to be 0)
            return observed_v_over_m_sq - implied_v_over_m_sq
        except ValueError:
            return -np.inf # Handle math errors

    @staticmethod
    def calculate_tercile_probabilities_bestfit(best_guess, error_variance, T1, T2, dist_code, dof 
    ):
        """
        Generic tercile probabilities using best-fit family per grid cell.

        Inputs (per grid cell):
        - best_guess : 1D array over T (hindcast_det or forecast_det)
        - T1, T2     : scalar terciles from climatological best-fit distribution
        - dist_code  : int, as in _ppf_terciles_from_code
        - shape, loc, scale : scalars from climatology fit

        Strategy:
        - For each time step, build a predictive distribution of the same family:
            * Use best_guess[t] to adjust mean / location;
            * Keep shape parameters from climatology.
        - Then compute probabilities:
            P(B) = F(T1), P(N) = F(T2) - F(T1), P(A) = 1 - F(T2).
        """
        
        best_guess = np.asarray(best_guess, float)
        error_variance = np.asarray(error_variance, dtype=float)
        # T1 = np.asarray(T1, dtype=float)
        # T2 = np.asarray(T2, dtype=float)
        n_time = best_guess.size
        out = np.full((3, n_time), np.nan, float)

        if np.all(np.isnan(best_guess)) or np.isnan(dist_code) or np.isnan(T1) or np.isnan(T2) or np.isnan(error_variance):
            return out

        code = int(dist_code)

        # Normal: loc = forecast; scale from clim
        if code == 1:
            error_std = np.sqrt(error_variance)
            out[0, :] = norm.cdf(T1, loc=best_guess, scale=error_std)
            out[1, :] = norm.cdf(T2, loc=best_guess, scale=error_std) - norm.cdf(T1, loc=best_guess, scale=error_std)
            out[2, :] = 1 - norm.cdf(T2, loc=best_guess, scale=error_std)

        # Lognormal: shape = sigma from clim; enforce mean = best_guess
        elif code == 2:
            sigma = np.sqrt(np.log(1 + error_variance / (best_guess**2)))
            mu = np.log(best_guess) - sigma**2 / 2
            out[0, :] = lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[1, :] = lognorm.cdf(T2, s=sigma, scale=np.exp(mu)) - lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[2, :] = 1 - lognorm.cdf(T2, s=sigma, scale=np.exp(mu))      


        # Exponential: keep scale from clim; shift loc so mean = best_guess
        elif code == 3:
            c1 = expon.cdf(T1, loc=best_guess, scale=np.sqrt(error_variance))
            c2 = expon.cdf(T2, loc=loc_t, scale=np.sqrt(error_variance))
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        # Gamma: use shape from clim; set scale so mean = best_guess
        elif code == 4:
            alpha = (best_guess ** 2) / error_variance
            theta = error_variance / best_guess
            c1 = gamma.cdf(T1, a=alpha, scale=theta)
            c2 = gamma.cdf(T2, a=alpha, scale=theta)
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        elif code == 5: # Assuming 5 is for Weibull   
        
            for i in range(n_time):
                # Get the scalar values for this specific element (e.g., grid cell)
                M = best_guess[i]
                print(M)
                V = error_variance
                print(V)
                
                # Handle cases with no variance to avoid division by zero
                if V <= 0 or M <= 0:
                    out[0, i] = np.nan
                    out[1, i] = np.nan
                    out[2, i] = np.nan
                    continue # Skip to the next element
        
                # --- 1. Numerically solve for shape 'k' ---
                # We need a reasonable starting guess. 2.0 is common (Rayleigh dist.)
                initial_guess = 2.0
                
                # fsolve finds the root of our helper function
                k = fsolve(weibull_shape_solver, initial_guess, args=(M, V))[0]
        
                # --- 2. Check for bad solution and calculate scale 'lambda' ---
                if k <= 0:
                    # Solver failed
                    out[0, i] = np.nan
                    out[1, i] = np.nan
                    out[2, i] = np.nan
                    continue
                
                # With 'k' found, we can now algebraically find scale 'lambda'
                # In scipy.stats, scale is 'scale'
                lambda_scale = M / gamma_function(1 + 1/k)
        
                # --- 3. Calculate Probabilities ---
                # In scipy.stats, shape 'k' is 'c'
                # Use the T1 and T2 values for this specific element
                
                c1 = weibull_min.cdf(T1, c=k, loc=0, scale=lambda_scale)
                c2 = weibull_min.cdf(T2, c=k, loc=0, scale=lambda_scale)
        
                out[0, i] = c1
                out[1, i] = c2 - c1
                out[2, i] = 1.0 - c2

        # Student-t: df from clim; scale from clim; loc = best_guess
        elif code == 6:       
            # Check if df is valid for variance calculation
            if dof <= 2:
                # Cannot calculate scale, fill with NaNs
                out[0, :] = np.nan
                out[1, :] = np.nan
                out[2, :] = np.nan
            else:
                # 1. Calculate t-distribution parameters
                # 'loc' (mean) is just the best_guess
                loc = best_guess
                # 'scale' is calculated from the variance and df
                # Variance = scale**2 * (df / (df - 2))
                scale = np.sqrt(error_variance * (dof - 2) / dof)
                
                # 2. Calculate probabilities
                c1 = t.cdf(T1, df=dof, loc=loc, scale=scale)
                c2 = t.cdf(T2, df=dof, loc=loc, scale=scale)

                out[0, :] = c1
                out[1, :] = c2 - c1
                out[2, :] = 1.0 - c2

        elif code == 7: # Assuming 7 is for Poisson
            
            # --- 1. Set the Poisson parameter 'mu' ---
            # The 'mu' parameter is the mean.
            
            # A warning is strongly recommended if error_variance is different from best_guess
            if not np.allclose(best_guess, error_variance, atol=0.5):
                print("Warning: 'error_variance' is not equal to 'best_guess'.")
                print("Poisson model assumes mean=variance and is likely inappropriate.")
                print("Consider using Negative Binomial.")
            
            mu = best_guess
        
            # --- 2. Calculate Probabilities ---
            # poisson.cdf(k, mu) calculates P(X <= k)
            
            c1 = poisson.cdf(T1, mu=mu)
            c2 = poisson.cdf(T2, mu=mu)
            
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        elif code == 8: # Assuming 8 is for Negative Binomial
            
            # --- 1. Calculate Negative Binomial Parameters ---
            # This model is ONLY valid for overdispersion (Variance > Mean).
            # We will use np.where to set parameters to NaN if V <= M.
            
            # p = Mean / Variance
            p = np.where(error_variance > best_guess, 
                         best_guess / error_variance, 
                         np.nan)
            
            # n = Mean^2 / (Variance - Mean)
            n = np.where(error_variance > best_guess, 
                         (best_guess**2) / (error_variance - best_guess), 
                         np.nan)
            
            # --- 2. Calculate Probabilities ---
            # The nbinom.cdf function will propagate NaNs, correctly
            # handling the cases where the model was invalid.
            
            c1 = nbinom.cdf(T1, n=n, p=p)
            c2 = nbinom.cdf(T2, n=n, p=p)
            
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2
            
        else:
            raise ValueError(f"Invalid distribution")

        return out

    @staticmethod
    def calculate_tercile_probabilities_nonparametric(best_guess, error_samples, first_tercile, second_tercile):
        """Non-parametric method using historical error samples."""
        n_time = len(best_guess)
        pred_prob = np.full((3, n_time), np.nan, dtype=float)
        for t in range(n_time):
            if np.isnan(best_guess[t]):
                continue
            dist = best_guess[t] + error_samples
            dist = dist[np.isfinite(dist)]
            if len(dist) == 0:
                continue
            p_below = np.mean(dist < first_tercile)
            p_between = np.mean((dist >= first_tercile) & (dist < second_tercile))
            p_above = 1.0 - (p_below + p_between)
            pred_prob[0, t] = p_below
            pred_prob[1, t] = p_between
            pred_prob[2, t] = p_above
        return pred_prob



    def compute_prob(
        self,
        Predictant: xr.DataArray,
        clim_year_start,
        clim_year_end,
        hindcast_det: xr.DataArray,
        best_code_da: xr.DataArray = None,
        best_shape_da: xr.DataArray = None,
        best_loc_da: xr.DataArray = None,
        best_scale_da: xr.DataArray = None
    ) -> xr.DataArray:
        """
        Compute tercile probabilities for deterministic hindcasts.

        If dist_method == 'bestfit':
            - Use cluster-based best-fit distributions to:
                * derive terciles analytically from (best_code_da, best_shape_da, best_loc_da, best_scale_da),
                * compute predictive probabilities using the same family.

        Otherwise:
            - Use empirical terciles from Predictant climatology and the selected
              parametric / nonparametric method.

        Parameters
        ----------
        Predictant : xarray.DataArray
            Observed data (T, Y, X) or (T, Y, X, M).
        clim_year_start, clim_year_end : int or str
            Climatology period (inclusive) for thresholds.
        hindcast_det : xarray.DataArray
            Deterministic hindcast (T, Y, X).
        best_code_da, best_shape_da, best_loc_da, best_scale_da : xarray.DataArray, optional
            Output from WAS_TransformData.fit_best_distribution_grid, required for 'bestfit'.

        Returns
        -------
        hindcast_prob : xarray.DataArray
            Probabilities with dims (probability=['PB','PN','PA'], T, Y, X).
        """
        # Handle member dimension if present
        if "M" in Predictant.dims:
            Predictant = Predictant.isel(M=0).drop_vars("M").squeeze()

        # Ensure dimension order
        Predictant = Predictant.transpose("T", "Y", "X")

        # Spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1.0, np.nan)

        # Climatology subset
        clim = Predictant.sel(T=slice(str(clim_year_start), str(clim_year_end)))
        if clim.sizes.get("T", 0) < 3:
            raise ValueError("Not enough years in climatology period for terciles.")

        # Error variance for predictive distributions
        error_variance = (Predictant - hindcast_det).var(dim="T")
        dof = max(int(clim.sizes["T"]) - 1, 2)

        # Empirical terciles (used by non-bestfit methods)
        terciles_emp = clim.quantile([0.32, 0.67], dim="T")
        T1_emp = terciles_emp.isel(quantile=0).drop_vars("quantile")
        T2_emp = terciles_emp.isel(quantile=1).drop_vars("quantile")
        

        dm = self.dist_method

        # ---------- BESTFIT: zone-wise optimal distributions ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da_da, best_loc_da, best_scale_da."
                )

            # T1, T2 from best-fit distributions (per grid)
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            # Predictive probabilities using same family
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                hindcast_det,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                kwargs={'dof': dof},
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                hindcast_det,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")

        hindcast_prob = hindcast_prob.assign_coords(
            probability=("probability", ["PB", "PN", "PA"])
        )
        return (hindcast_prob * mask).transpose("probability", "T", "Y", "X")



    def forecast(self, Predictant, clim_year_start, clim_year_end, Predictor, hindcast_det, Predictor_for_year, best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1, np.nan).drop_vars(['T']).squeeze().to_numpy()
        Predictor_ = (Predictor - trend_data(Predictor).fillna(trend_data(Predictor)[-3])).fillna(0)
        Predictant_st = standardize_timeseries(Predictant, clim_year_start, clim_year_end)
        Predictant_ = (Predictant_st - trend_data(Predictant_st).fillna(trend_data(Predictant_st)[-3])).fillna(0)
        
        Predictor_for_year_ = ((((Predictor_for_year.fillna(Predictor.mean(dim="T", skipna=True))).ffill(dim="Y").bfill(dim="Y")).ffill(dim="X").bfill(dim="X")).fillna(0)).transpose('T', 'Y', 'X')

        # last_trend_X = ((trend_data(standardize_timeseries(Predictor, clim_year_start, clim_year_end))).isel(T=[-3]))
        # last_trend_X = ((trend_data(Predictor)).isel(T=[-3]))
        # last_trend_X['T'] = Predictor_for_year_['T']
        # Predictor_for_year__ = Predictor_for_year_.fillna(0)
        # Predictor_for_year__ = (Predictor_for_year_ - last_trend_X).fillna(0)

        Predictor_for_year__ = Predictor_for_year_

        # Fit the CCA model
        self.fit_cca(Predictor_, Predictant_)
            
        # Prepare test data
        X_test_prepared = Predictor_for_year__.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')

        # Predict
        y_pred = self.cca_model.predict(X_test_prepared)
        y_pred = self.cca_model.inverse_transform(self.cca_model.transform(X_test_prepared), y_pred)[1] 
        result_ = y_pred.rename({"lon": "X", "lat": "Y"})
        
        # last_trend_Y = ((trend_data(Predictant_st)).isel(T=[-3]))
        # last_trend_Y['T'] = result_['T']
        # result_ = (result_ + last_trend_Y)
       
        result_ = reverse_standardize(result_, Predictant, clim_year_start, clim_year_end) 
        
        # 2) Compute thresholds T1, T2 from climatology
        index_start = Predictant.get_index("T").get_loc(str(clim_year_start)).start
        index_end   = Predictant.get_index("T").get_loc(str(clim_year_end)).stop
        rainfall_for_tercile = Predictant.isel(T=slice(index_start, index_end))
        terciles = rainfall_for_tercile.quantile([0.32, 0.67], dim='T')
        T1_emp = terciles.isel(quantile=0).drop_vars('quantile')
        T2_emp = terciles.isel(quantile=1).drop_vars('quantile')
        error_variance = (Predictant - hindcast_det).var(dim='T')

        year = Predictor_for_year.coords['T'].values.astype('datetime64[Y]').astype(int)[0] + 1970  # Convert from epoch
        T_value_1 = Predictant.isel(T=0).coords['T'].values  # Get the datetime64 value from da1
        month_1 = T_value_1.astype('datetime64[M]').astype(int) % 12 + 1  # Extract month
        new_T_value = np.datetime64(f"{year}-{month_1:02d}-{1:02d}")
        forecast_expanded = result_.assign_coords(T=xr.DataArray([new_T_value], dims=["T"]))
        forecast_expanded['T'] = forecast_expanded['T'].astype('datetime64[ns]')
        
        
        dof = max(int(rainfall_for_tercile.sizes["T"]) - 1, 2)

        dm = self.dist_method

        # ---------- BESTFIT ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da, best_loc_da, best_scale_da."
                )
            
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                forecast_expanded,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                kwargs={"dof": dof},
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                forecast_expanded,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")
        forecast_prob = forecast_prob.assign_coords(probability=('probability', ['PB', 'PN', 'PA']))
        return forecast_expanded * mask, mask * forecast_prob.transpose('probability', 'T', 'Y', 'X')
        

    def plot_cca_results(self, X=None, Y=None, n_modes=None, clim_year_start=None, clim_year_end=None):
        """
        Plots the CCA modes and scores.

        Parameters:
        - X: Optional xarray DataArray for predictors. If provided, the model will be fitted using X and Y.
        - Y: Optional xarray DataArray for predictands.
        - n_modes: Number of modes to plot. If None, plots all modes.
        """
        if X is not None and Y is not None:
            mask = xr.where(~np.isnan(Y.isel(T=0)), 1, np.nan).drop_vars(['T']).squeeze().to_numpy()
            # mask.name = None
            
            # X_ = standardize_timeseries(X, clim_year_start, clim_year_end) - trend_data(standardize_timeseries(X, clim_year_start, clim_year_end)).fillna(trend_data(standardize_timeseries(X, clim_year_start, clim_year_end))[-3])

            X_ = X - trend_data(X).fillna(trend_data(X)[-3])
            
            Y_ = standardize_timeseries(Y, clim_year_start, clim_year_end) - trend_data(standardize_timeseries(Y, clim_year_start, clim_year_end)).fillna(trend_data(standardize_timeseries(Y, clim_year_start, clim_year_end))[-3])
            
            # Fit the model using the provided data
            self.fit_cca(X_.isel(T= slice(0,-2)).fillna(0), Y_.isel(T=slice(0,-2)).fillna(0))
        elif self.cca_model is None:
            raise ValueError("The CCA model has not been fitted yet. Provide X and Y data to fit the model.")

        # Get components (modes) and scores
        X_modes, Y_modes = self.cca_model.components()  # Spatial patterns
        X_scores, Y_scores = self.cca_model.scores()    # Temporal projections (canonical variates)

        # Get explained variances
        var_explained_X = self.cca_model.fraction_variance_X_explained_by_X()
        var_explained_Y = self.cca_model.fraction_variance_Y_explained_by_Y()
        var_explained_Y_by_X = self.cca_model.fraction_variance_Y_explained_by_X()

        # Determine number of modes to plot
        if n_modes is None:
            n_modes = self.n_modes

        # Mode indices start from 1 in xeofs
        mode_indices = range(1, n_modes + 1)

        # Create subplots
        fig, axes = plt.subplots(n_modes, 3, figsize=(15, 3 * n_modes))

        if n_modes == 1:
            axes = np.array([axes])

        for i, mode in enumerate(mode_indices):

            # First Column: Plot X_modes
            ax = axes[i, 0]
            X_mode = X_modes.sel(mode=mode)
            X_mode.plot(ax=ax, vmin=-1, vmax=1, cmap= "RdBu_r")
            var_X = var_explained_X.sel(mode=mode).values * 100
            ax.set_title(f'X Mode {mode} ({var_X:.2f}% variance explained)')

            # Second Column: Plot X_scores and Y_scores
            ax = axes[i, 1]
            X_score = X_scores.sel(mode=mode)
            Y_score = Y_scores.sel(mode=mode)
            var_Y_X = var_explained_Y_by_X.sel(mode=mode).values * 100
            ax.plot(X_score['T'].dt.year.values, X_score, label='X Score')
            ax.plot(Y_score['T'].dt.year.values, Y_score, label='Y Score')
            ax.axhline(0, linestyle='--', lw=0.8, label="") #### line Canonical Variate = 0
            ax.legend()
            ax.set_title(f'Scores for Mode {mode} ({var_Y_X:.2f}% variance Y explained by X)')
            ax.set_xlabel('Time')
            ax.set_ylabel('Canonical Variate')

            # Third Column: Plot Y_modes
            ax = axes[i, 2]
            Y_mode = (Y_modes.sel(mode=mode))*mask
            Y_mode.plot(ax=ax, vmin=None, vmax=None, cmap= "RdBu_r")
            var_Y = var_explained_Y.sel(mode=mode).values * 100
            ax.set_title(f'Y Mode {mode} ({var_Y:.2f}% variance explained)')
            
        plt.tight_layout()
        plt.show()



class WAS_KCCA:
    def __init__(self, n_modes=4, kernel='linear', dist_method="nonparam", 
                 cv_folds=3, random_state=42, search_method='grid'):
        """
        Initialize the WAS_KCCA class with Kernel CCA.
        
        Parameters:
        - n_modes: Number of canonical modes to compute.
        - kernel: Kernel type ('linear', 'rbf', 'poly', 'sigmoid').
        - dist_method: Method for probability calculation ('bestfit', 'nonparam', 'gamma').
        - cv_folds: Number of cross-validation folds for hyperparameter tuning.
        - random_state: Random seed for reproducibility.
        - search_method: Hyperparameter search method ('grid', 'random', 'bayesian').
        """
        self.n_modes = n_modes
        self.kernel = kernel
        self.dist_method = dist_method
        self.cv_folds = cv_folds
        self.random_state = random_state
        self.search_method = search_method
        
        # Hyperparameter storage (tuned once, used everywhere)
        self.best_params = None  # Will store tuned parameters
        self.is_tuned = False    # Whether parameters have been tuned
        
        # Model and data storage
        self.model = None
        self.is_fitted = False
        
        # Data storage
        self.Y_coords = None
        self.X_mean = None
        self.Y_mean = None
        self.X_train_original = None
        self.Y_train_original = None
        self.X_train_np = None
        self.Y_train_np = None

    # --------------------------------------------------------------------------
    #  Hyperparameter Tuning (Do this once!)
    # --------------------------------------------------------------------------
    
    def _prepare_param_grid(self, kernel=None):
        """
        Prepare parameter grid for hyperparameter tuning.
        
        Parameters:
        - kernel: Kernel type (if None, use self.kernel)
        
        Returns:
        - param_grid: Dictionary of parameter distributions
        """
        if kernel is None:
            kernel = self.kernel
        
        if kernel == 'rbf':
            return {'sigma': [0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 50.0]}
        elif kernel == 'poly':
            return {
                'degree': [2, 3, 4],
                'gamma': [0.1, 0.5, 1.0, 2.0, 5.0],
                'coef0': [0, 1, 2]
            }
        elif kernel == 'sigmoid':
            return {
                'gamma': [0.01, 0.05, 0.1, 0.5, 1.0],
                'coef0': [0, 1, 2, 5]
            }
        else:  # linear kernel
            return {}
    
    def _objective(self, trial, X_np, Y_np):
        """
        Objective function for Bayesian optimization (if using Optuna).
        """
        if self.kernel == 'rbf':
            sigma = trial.suggest_float('sigma', 0.1, 50.0, log=True)
            params = {'sigma': sigma}
        elif self.kernel == 'poly':
            degree = trial.suggest_int('degree', 2, 4)
            gamma = trial.suggest_float('gamma', 0.1, 5.0, log=True)
            coef0 = trial.suggest_float('coef0', 0, 5)
            params = {'degree': degree, 'gamma': gamma, 'coef0': coef0}
        elif self.kernel == 'sigmoid':
            gamma = trial.suggest_float('gamma', 0.01, 1.0, log=True)
            coef0 = trial.suggest_float('coef0', 0, 5)
            params = {'gamma': gamma, 'coef0': coef0}
        else:
            return 0.0  # Linear kernel has no parameters to tune
        
        # Create KCCA model with suggested parameters
        model = KCCA(latent_dimensions=self.n_modes, kernel=self.kernel, **params)
        
        # Cross-validation
        cv = KFold(n_splits=self.cv_folds, shuffle=True, random_state=self.random_state)
        correlations = []
        
        for train_idx, val_idx in cv.split(X_np):
            X_train, X_val = X_np[train_idx], X_np[val_idx]
            Y_train, Y_val = Y_np[train_idx], Y_np[val_idx]
            
            model.fit([X_train, Y_train])
            scores = model.transform([X_val, Y_val])
            
            # Calculate canonical correlation
            corr = np.corrcoef(scores[0][:, 0], scores[1][:, 0])[0, 1]
            correlations.append(corr)
        
        return np.mean(correlations)
    
    def compute_hyperparameters(self, Predictors, Predictand, clim_year_start=None, clim_year_end=None):
        """
        Independently computes the best hyperparameters using selected optimization method.
        
        Parameters:
        - Predictors: Training predictor data with dimensions (T, Y, X).
        - Predictand: Training predictand data with dimensions (T, Y, X).
        - clim_year_start, clim_year_end: Optional climatology period for standardization.
        
        Returns:
        - best_params: Dictionary of best hyperparameters found.
        """
        if self.kernel == 'linear':
            print("Linear kernel has no hyperparameters to tune.")
            self.best_params = {}
            self.is_tuned = True
            return {}
        
        if not HAS_CCA_ZOO:
            raise ImportError("cca-zoo is required for hyperparameter tuning")
        
        print(f"Computing hyperparameters for {self.kernel} KCCA using {self.search_method} search...")
        
        # Preprocess data (no standardization needed as input is already standardized)
        X_train_final, y_train_final = self.preprocess_data(Predictors, Predictand)
        
        # Stack spatial dimensions
        X_np, X_coords = self._stack_spatial(X_train_final)
        Y_np, Y_coords = self._stack_spatial(y_train_final)
        
        # Remove any remaining NaN values
        mask = ~(np.any(~np.isfinite(X_np), axis=1) | np.any(~np.isfinite(Y_np), axis=1))
        X_np_clean = X_np[mask]
        Y_np_clean = Y_np[mask]
        
        if self.search_method == 'grid':
            # Grid Search
            param_grid = self._prepare_param_grid()
            if not param_grid:
                self.best_params = {}
                self.is_tuned = True
                return {}
            
            print(f"Performing grid search with parameters: {param_grid}")
            
            cv = KFold(n_splits=self.cv_folds, shuffle=True, random_state=self.random_state)
            
            grid_search = CCA_GridSearch(
                KCCA(latent_dimensions=self.n_modes, kernel=self.kernel),
                param_grid=param_grid,
                cv=cv,
                n_jobs=-1,
                verbose=1
            )
            
            grid_search.fit([X_np_clean, Y_np_clean])
            best_params = grid_search.best_params_
            best_score = grid_search.best_score_
            
            print(f"Best parameters found: {best_params}")
            print(f"Best cross-validation score: {best_score:.4f}")
            
            self.best_params = best_params
            self.is_tuned = True
            
        elif self.search_method == 'bayesian':
            # Bayesian Optimization with Optuna
            try:
                import optuna
                from optuna.samplers import TPESampler
            except ImportError:
                raise ImportError("Optuna is required for Bayesian optimization")
            
            print("Performing Bayesian optimization with Optuna...")
            
            sampler = TPESampler(seed=self.random_state)
            study = optuna.create_study(
                direction='maximize',
                sampler=sampler
            )
            
            # Create objective function with data
            objective_with_data = lambda trial: self._objective(trial, X_np_clean, Y_np_clean)
            study.optimize(objective_with_data, n_trials=50)  # Adjust n_trials as needed
            
            # Get best parameters
            best_params = study.best_params
            best_score = study.best_value
            
            print(f"Best parameters found: {best_params}")
            print(f"Best cross-validation score: {best_score:.4f}")
            
            self.best_params = best_params
            self.is_tuned = True
            
        else:  # Default parameters
            print("Using default parameters (no search performed)")
            self.best_params = {}
            self.is_tuned = True
        
        return self.best_params
    
    def set_hyperparameters(self, hyperparams):
        """
        Manually set hyperparameters.
        
        Parameters:
        - hyperparams: Dictionary of hyperparameters to use
        """
        self.best_params = hyperparams
        self.is_tuned = True
        print(f"Hyperparameters set: {hyperparams}")
    
    # --------------------------------------------------------------------------
    #  Data Preparation Methods
    # --------------------------------------------------------------------------
    
    def _stack_spatial(self, da):
        """Convert spatial DataArray to 2D array (time × features)."""
        # Check dimension names
        if 'X' in da.dims and 'Y' in da.dims:
            da = da.rename({"X": "lon", "Y": "lat"})
        
        # Stack spatial dimensions
        stacked = da.stack(feature=['lat', 'lon'])
        return stacked.values, stacked.coords

    def _unstack_spatial(self, data_np, coords, time_coords):
        """Reconstruct spatial DataArray from 2D array."""
        if data_np.ndim == 1:
            data_np = data_np.reshape(1, -1)
            
        da_stacked = xr.DataArray(
            data_np,
            dims=['T', 'feature'],
            coords={
                'T': time_coords,
                'feature': coords['feature']
            }
        )
        
        da = da_stacked.unstack('feature')
        return da.rename({"lon": "X", "lat": "Y"})

    def preprocess_data(self, X, Y):
        """Preprocess data by filling missing values."""
        # Fill missing values
        X_filled = X.fillna(X.mean(dim="T", skipna=True))
        Y_filled = Y.fillna(Y.mean(dim="T", skipna=True))
        
        # Rename and transpose
        X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        Y_final = Y_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        
        return X_final, Y_final

    def preprocess_test_data(self, X_test, y_test, X_train, y_train):
        """Preprocess test data using training statistics."""
        # Fill missing values with training means
        X_test_prepared = X_test.fillna(X_train.mean(dim="T", skipna=True))
        y_test_prepared = y_test.fillna(y_train.mean(dim="T", skipna=True))
        
        # Rename and transpose
        dims_rename = {"X": "lon", "Y": "lat"}
        X_test_prepared = X_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')
        y_test_prepared = y_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')
        
        return X_test_prepared, y_test_prepared

    # --------------------------------------------------------------------------
    #  Model Fitting with Explicit Parameters
    # --------------------------------------------------------------------------
    
    def fit_cca(self, X_train, y_train, best_params=None):
        """
        Fit the KCCA model using provided hyperparameters.
        
        Parameters:
        - X_train, y_train: Training data
        - best_params: Dictionary of hyperparameters to use. If None, uses self.best_params
        """
        # Use provided parameters if available, otherwise stored parameters
        if best_params is not None:
            params_to_use = best_params
            print(f"Using provided hyperparameters: {best_params}")
        elif self.best_params is not None:
            params_to_use = self.best_params
            print(f"Using stored hyperparameters: {self.best_params}")
        else:
            params_to_use = {}
            if self.kernel != 'linear':
                print("Warning: No hyperparameters provided. Using default parameters.")
        
        # Preprocess (no standardization)
        X_train_final, y_train_final = self.preprocess_data(X_train, y_train)
        
        # Store original data
        self.X_train_original = X_train_final
        self.Y_train_original = y_train_final
        
        # Stack spatial dimensions
        X_np, X_coords = self._stack_spatial(X_train_final)
        Y_np, Y_coords = self._stack_spatial(y_train_final)
        
        self.Y_coords = Y_coords
        self.X_train_np = X_np
        self.Y_train_np = Y_np
        
        # Build model with appropriate parameters
        if self.kernel == 'linear':
            self.model = LinearCCA(latent_dimensions=self.n_modes)
        else:
            self.model = KCCA(latent_dimensions=self.n_modes, kernel=self.kernel, **params_to_use)
        
        # Fit model
        self.model.fit([X_np, Y_np])
        self.is_fitted = True
        
        # Store canonical correlations
        if hasattr(self.model, 'canonical_correlations'):
            self.canonical_correlations = self.model.canonical_correlations()
        else:
            scores = self.model.transform([X_np, Y_np])
            self.canonical_correlations = [
                np.corrcoef(scores[0][:, i], scores[1][:, i])[0, 1] 
                for i in range(self.n_modes)
            ]
    
    # --------------------------------------------------------------------------
    #  Prediction Methods
    # --------------------------------------------------------------------------
    
    def _predict_kcca(self, X_test):
        """Make predictions using KCCA model."""
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit_cca first.")
        
        # Stack spatial dimensions
        X_test_np, X_coords = self._stack_spatial(X_test)
        
        # Transform to get Y scores
        Y_scores = self.model.transform([X_test_np])[1]
        
        # Reconstruct spatial array (no inverse scaling needed as input is already standardized)
        Y_pred = self._unstack_spatial(Y_scores, self.Y_coords, X_test['T'].values)
        
        return Y_pred

    def compute_model(self, X_train, y_train, X_test, y_test, best_params=None):
        """
        Compute deterministic hindcast using KCCA with injected hyperparameters.
        
        Parameters:
        - X_train: Training predictor data with dimensions (T, Y, X).
        - y_train: Training predictand data with dimensions (T, Y, X).
        - X_test: Testing predictor data with dimensions (T, Y, X).
        - y_test: Testing predictand data with dimensions (T, Y, X).
        - best_params: Pre-computed best hyperparameters. If None, uses stored parameters.
        
        Returns:
        - hindcast: Deterministic hindcast with dimensions (T, Y, X).
        """
        # Fit model with parameters
        self.fit_cca(X_train, y_train, best_params=best_params)
        
        # Preprocess test data
        X_test_prepared, y_test_prepared = self.preprocess_test_data(
            X_test, y_test, X_train, y_train
        )
        
        # Predict
        y_pred = self._predict_kcca(X_test_prepared)
        y_pred['T'] = y_test_prepared['T']
        
        return y_pred

    # --------------------------------------------------------------------------
    #  Probability Calculation Methods
    # --------------------------------------------------------------------------
    # ------------------ Probability Calculation Methods ------------------

    @staticmethod
    def _ppf_terciles_from_code(dist_code, shape, loc, scale):
        """
        Return tercile thresholds (T1, T2) from best-fit distribution parameters.
    
        dist_code:
            1: norm
            2: lognorm
            3: expon
            4: gamma
            5: weibull_min
            6: t
            7: poisson
            8: nbinom
        """
        if np.isnan(dist_code):
            return np.nan, np.nan
    
        code = int(dist_code)
        try:
            if code == 1:
                return (
                    norm.ppf(0.32, loc=loc, scale=scale),
                    norm.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 2:
                return (
                    lognorm.ppf(0.32, s=shape, loc=loc, scale=scale),
                    lognorm.ppf(0.67, s=shape, loc=loc, scale=scale),
                )
            elif code == 3:
                return (
                    expon.ppf(0.32, loc=loc, scale=scale),
                    expon.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 4:
                return (
                    gamma.ppf(0.32, a=shape, loc=loc, scale=scale),
                    gamma.ppf(0.67, a=shape, loc=loc, scale=scale),
                )
            elif code == 5:
                return (
                    weibull_min.ppf(0.32, c=shape, loc=loc, scale=scale),
                    weibull_min.ppf(0.67, c=shape, loc=loc, scale=scale),
                )
            elif code == 6:
                # Note: Renamed 't_dist' to 't' for standard scipy.stats
                return (
                    t.ppf(0.32, df=shape, loc=loc, scale=scale),
                    t.ppf(0.67, df=shape, loc=loc, scale=scale),
                )
            elif code == 7:
                # Poisson: poisson.ppf(q, mu, loc=0)
                # ASSUMPTION: 'mu' (mean) is passed as 'shape'
                #             'loc' is passed as 'loc'
                #             'scale' is unused
                return (
                    poisson.ppf(0.32, mu=shape, loc=loc),
                    poisson.ppf(0.67, mu=shape, loc=loc),
                )
            elif code == 8:
                # Negative Binomial: nbinom.ppf(q, n, p, loc=0)
                # ASSUMPTION: 'n' (successes) is passed as 'shape'
                #             'p' (probability) is passed as 'scale'
                #             'loc' is passed as 'loc'
                return (
                    nbinom.ppf(0.32, n=shape, p=scale, loc=loc),
                    nbinom.ppf(0.67, n=shape, p=scale, loc=loc),
                )
        except Exception:
            return np.nan, np.nan
    
        # Fallback if code is not 1-8
        return np.nan, np.nan
        
    @staticmethod
    def weibull_shape_solver(k, M, V):
        """
        Function to find the root of the Weibull shape parameter 'k'.
        We find 'k' such that the theoretical variance/mean^2 ratio
        matches the observed V/M^2 ratio.
        """
        # Guard against invalid 'k' values during solving
        if k <= 0:
            return -np.inf
        try:
            g1 = gamma_function(1 + 1/k)
            g2 = gamma_function(1 + 2/k)
            
            # This is the V/M^2 ratio *implied by k*
            implied_v_over_m_sq = (g2 / (g1**2)) - 1
            
            # This is the *observed* ratio
            observed_v_over_m_sq = V / (M**2)
            
            # Return the difference (we want this to be 0)
            return observed_v_over_m_sq - implied_v_over_m_sq
        except ValueError:
            return -np.inf # Handle math errors

    @staticmethod
    def calculate_tercile_probabilities_bestfit(best_guess, error_variance, T1, T2, dist_code, dof 
    ):
        """
        Generic tercile probabilities using best-fit family per grid cell.

        Inputs (per grid cell):
        - best_guess : 1D array over T (hindcast_det or forecast_det)
        - T1, T2     : scalar terciles from climatological best-fit distribution
        - dist_code  : int, as in _ppf_terciles_from_code
        - shape, loc, scale : scalars from climatology fit

        Strategy:
        - For each time step, build a predictive distribution of the same family:
            * Use best_guess[t] to adjust mean / location;
            * Keep shape parameters from climatology.
        - Then compute probabilities:
            P(B) = F(T1), P(N) = F(T2) - F(T1), P(A) = 1 - F(T2).
        """
        
        best_guess = np.asarray(best_guess, float)
        error_variance = np.asarray(error_variance, dtype=float)
        # T1 = np.asarray(T1, dtype=float)
        # T2 = np.asarray(T2, dtype=float)
        n_time = best_guess.size
        out = np.full((3, n_time), np.nan, float)

        if np.all(np.isnan(best_guess)) or np.isnan(dist_code) or np.isnan(T1) or np.isnan(T2) or np.isnan(error_variance):
            return out

        code = int(dist_code)

        # Normal: loc = forecast; scale from clim
        if code == 1:
            error_std = np.sqrt(error_variance)
            out[0, :] = norm.cdf(T1, loc=best_guess, scale=error_std)
            out[1, :] = norm.cdf(T2, loc=best_guess, scale=error_std) - norm.cdf(T1, loc=best_guess, scale=error_std)
            out[2, :] = 1 - norm.cdf(T2, loc=best_guess, scale=error_std)

        # Lognormal: shape = sigma from clim; enforce mean = best_guess
        elif code == 2:
            sigma = np.sqrt(np.log(1 + error_variance / (best_guess**2)))
            mu = np.log(best_guess) - sigma**2 / 2
            out[0, :] = lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[1, :] = lognorm.cdf(T2, s=sigma, scale=np.exp(mu)) - lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[2, :] = 1 - lognorm.cdf(T2, s=sigma, scale=np.exp(mu))      


        # Exponential: keep scale from clim; shift loc so mean = best_guess
        elif code == 3:
            c1 = expon.cdf(T1, loc=best_guess, scale=np.sqrt(error_variance))
            c2 = expon.cdf(T2, loc=loc_t, scale=np.sqrt(error_variance))
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        # Gamma: use shape from clim; set scale so mean = best_guess
        elif code == 4:
            alpha = (best_guess ** 2) / error_variance
            theta = error_variance / best_guess
            c1 = gamma.cdf(T1, a=alpha, scale=theta)
            c2 = gamma.cdf(T2, a=alpha, scale=theta)
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        elif code == 5: # Assuming 5 is for Weibull   
        
            for i in range(n_time):
                # Get the scalar values for this specific element (e.g., grid cell)
                M = best_guess[i]
                print(M)
                V = error_variance
                print(V)
                
                # Handle cases with no variance to avoid division by zero
                if V <= 0 or M <= 0:
                    out[0, i] = np.nan
                    out[1, i] = np.nan
                    out[2, i] = np.nan
                    continue # Skip to the next element
        
                # --- 1. Numerically solve for shape 'k' ---
                # We need a reasonable starting guess. 2.0 is common (Rayleigh dist.)
                initial_guess = 2.0
                
                # fsolve finds the root of our helper function
                k = fsolve(weibull_shape_solver, initial_guess, args=(M, V))[0]
        
                # --- 2. Check for bad solution and calculate scale 'lambda' ---
                if k <= 0:
                    # Solver failed
                    out[0, i] = np.nan
                    out[1, i] = np.nan
                    out[2, i] = np.nan
                    continue
                
                # With 'k' found, we can now algebraically find scale 'lambda'
                # In scipy.stats, scale is 'scale'
                lambda_scale = M / gamma_function(1 + 1/k)
        
                # --- 3. Calculate Probabilities ---
                # In scipy.stats, shape 'k' is 'c'
                # Use the T1 and T2 values for this specific element
                
                c1 = weibull_min.cdf(T1, c=k, loc=0, scale=lambda_scale)
                c2 = weibull_min.cdf(T2, c=k, loc=0, scale=lambda_scale)
        
                out[0, i] = c1
                out[1, i] = c2 - c1
                out[2, i] = 1.0 - c2

        # Student-t: df from clim; scale from clim; loc = best_guess
        elif code == 6:       
            # Check if df is valid for variance calculation
            if dof <= 2:
                # Cannot calculate scale, fill with NaNs
                out[0, :] = np.nan
                out[1, :] = np.nan
                out[2, :] = np.nan
            else:
                # 1. Calculate t-distribution parameters
                # 'loc' (mean) is just the best_guess
                loc = best_guess
                # 'scale' is calculated from the variance and df
                # Variance = scale**2 * (df / (df - 2))
                scale = np.sqrt(error_variance * (dof - 2) / dof)
                
                # 2. Calculate probabilities
                c1 = t.cdf(T1, df=dof, loc=loc, scale=scale)
                c2 = t.cdf(T2, df=dof, loc=loc, scale=scale)

                out[0, :] = c1
                out[1, :] = c2 - c1
                out[2, :] = 1.0 - c2

        elif code == 7: # Assuming 7 is for Poisson
            
            # --- 1. Set the Poisson parameter 'mu' ---
            # The 'mu' parameter is the mean.
            
            # A warning is strongly recommended if error_variance is different from best_guess
            if not np.allclose(best_guess, error_variance, atol=0.5):
                print("Warning: 'error_variance' is not equal to 'best_guess'.")
                print("Poisson model assumes mean=variance and is likely inappropriate.")
                print("Consider using Negative Binomial.")
            
            mu = best_guess
        
            # --- 2. Calculate Probabilities ---
            # poisson.cdf(k, mu) calculates P(X <= k)
            
            c1 = poisson.cdf(T1, mu=mu)
            c2 = poisson.cdf(T2, mu=mu)
            
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2

        elif code == 8: # Assuming 8 is for Negative Binomial
            
            # --- 1. Calculate Negative Binomial Parameters ---
            # This model is ONLY valid for overdispersion (Variance > Mean).
            # We will use np.where to set parameters to NaN if V <= M.
            
            # p = Mean / Variance
            p = np.where(error_variance > best_guess, 
                         best_guess / error_variance, 
                         np.nan)
            
            # n = Mean^2 / (Variance - Mean)
            n = np.where(error_variance > best_guess, 
                         (best_guess**2) / (error_variance - best_guess), 
                         np.nan)
            
            # --- 2. Calculate Probabilities ---
            # The nbinom.cdf function will propagate NaNs, correctly
            # handling the cases where the model was invalid.
            
            c1 = nbinom.cdf(T1, n=n, p=p)
            c2 = nbinom.cdf(T2, n=n, p=p)
            
            out[0, :] = c1
            out[1, :] = c2 - c1
            out[2, :] = 1.0 - c2
            
        else:
            raise ValueError(f"Invalid distribution")

        return out

    @staticmethod
    def calculate_tercile_probabilities_nonparametric(best_guess, error_samples, first_tercile, second_tercile):
        """Non-parametric method using historical error samples."""
        n_time = len(best_guess)
        pred_prob = np.full((3, n_time), np.nan, dtype=float)
        for t in range(n_time):
            if np.isnan(best_guess[t]):
                continue
            dist = best_guess[t] + error_samples
            dist = dist[np.isfinite(dist)]
            if len(dist) == 0:
                continue
            p_below = np.mean(dist < first_tercile)
            p_between = np.mean((dist >= first_tercile) & (dist < second_tercile))
            p_above = 1.0 - (p_below + p_between)
            pred_prob[0, t] = p_below
            pred_prob[1, t] = p_between
            pred_prob[2, t] = p_above
        return pred_prob



    def compute_prob(
        self,
        Predictant: xr.DataArray,
        clim_year_start,
        clim_year_end,
        hindcast_det: xr.DataArray,
        best_code_da: xr.DataArray = None,
        best_shape_da: xr.DataArray = None,
        best_loc_da: xr.DataArray = None,
        best_scale_da: xr.DataArray = None
    ) -> xr.DataArray:
        """
        Compute tercile probabilities for deterministic hindcasts.

        If dist_method == 'bestfit':
            - Use cluster-based best-fit distributions to:
                * derive terciles analytically from (best_code_da, best_shape_da, best_loc_da, best_scale_da),
                * compute predictive probabilities using the same family.

        Otherwise:
            - Use empirical terciles from Predictant climatology and the selected
              parametric / nonparametric method.

        Parameters
        ----------
        Predictant : xarray.DataArray
            Observed data (T, Y, X) or (T, Y, X, M).
        clim_year_start, clim_year_end : int or str
            Climatology period (inclusive) for thresholds.
        hindcast_det : xarray.DataArray
            Deterministic hindcast (T, Y, X).
        best_code_da, best_shape_da, best_loc_da, best_scale_da : xarray.DataArray, optional
            Output from WAS_TransformData.fit_best_distribution_grid, required for 'bestfit'.

        Returns
        -------
        hindcast_prob : xarray.DataArray
            Probabilities with dims (probability=['PB','PN','PA'], T, Y, X).
        """
        # Handle member dimension if present
        if "M" in Predictant.dims:
            Predictant = Predictant.isel(M=0).drop_vars("M").squeeze()

        # Ensure dimension order
        Predictant = Predictant.transpose("T", "Y", "X")

        # Spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1.0, np.nan)

        # Climatology subset
        clim = Predictant.sel(T=slice(str(clim_year_start), str(clim_year_end)))
        if clim.sizes.get("T", 0) < 3:
            raise ValueError("Not enough years in climatology period for terciles.")

        # Error variance for predictive distributions
        error_variance = (Predictant - hindcast_det).var(dim="T")
        dof = max(int(clim.sizes["T"]) - 1, 2)

        # Empirical terciles (used by non-bestfit methods)
        terciles_emp = clim.quantile([0.32, 0.67], dim="T")
        T1_emp = terciles_emp.isel(quantile=0).drop_vars("quantile")
        T2_emp = terciles_emp.isel(quantile=1).drop_vars("quantile")
        

        dm = self.dist_method

        # ---------- BESTFIT: zone-wise optimal distributions ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da_da, best_loc_da, best_scale_da."
                )

            # T1, T2 from best-fit distributions (per grid)
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            # Predictive probabilities using same family
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                hindcast_det,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                kwargs={'dof': dof},
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                hindcast_det,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")

        hindcast_prob = hindcast_prob.assign_coords(
            probability=("probability", ["PB", "PN", "PA"])
        )
        return (hindcast_prob * mask).transpose("probability", "T", "Y", "X")

    # --------------------------------------------------------------------------
    #  Forecast Method (with explicit parameters)
    # --------------------------------------------------------------------------
    
    def forecast(self, Predictant, clim_year_start, clim_year_end, Predictor, 
                 hindcast_det, Predictor_for_year, best_params=None,
                 best_code_da=None, best_shape_da=None, best_loc_da=None, 
                 best_scale_da=None):
        """
        Generate operational forecast with probabilities using specified hyperparameters.
        
        Parameters:
        - best_params: Hyperparameters to use for the forecast model.
        """
        # Create spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1, np.nan)
        
        # Fit model with specified parameters
        # Apply detrending and standardization if needed
        Predictor_ = (Predictor - trend_data(Predictor).fillna(
            trend_data(Predictor)[-3])).fillna(0)
        Predictant_st = standardize_timeseries(Predictant, clim_year_start, clim_year_end)
        Predictant_ = (Predictant_st - trend_data(Predictant_st).fillna(
            trend_data(Predictant_st)[-3])).fillna(0)
        
        self.fit_cca(Predictor_, Predictant_, best_params=best_params)
        
        # Prepare forecast predictor
        Predictor_for_year_ = Predictor_for_year.fillna(Predictor.mean(dim="T", skipna=True))
        Predictor_for_year_ = Predictor_for_year_.transpose('T', 'Y', 'X')
        
        # Make deterministic forecast
        result_det = self._predict_kcca(Predictor_for_year_)
        
        # Apply reverse standardization if needed
        if hasattr(self, '_standardized'):
            result_det = reverse_standardize(result_det, Predictant, clim_year_start, clim_year_end)
        
        # 2) Compute thresholds T1, T2 from climatology
        index_start = Predictant.get_index("T").get_loc(str(clim_year_start)).start
        index_end   = Predictant.get_index("T").get_loc(str(clim_year_end)).stop
        rainfall_for_tercile = Predictant.isel(T=slice(index_start, index_end))
        terciles = rainfall_for_tercile.quantile([0.32, 0.67], dim='T')
        T1_emp = terciles.isel(quantile=0).drop_vars('quantile')
        T2_emp = terciles.isel(quantile=1).drop_vars('quantile')
        error_variance = (Predictant - hindcast_det).var(dim='T')

        year = Predictor_for_year.coords['T'].values.astype('datetime64[Y]').astype(int)[0] + 1970  # Convert from epoch
        T_value_1 = Predictant.isel(T=0).coords['T'].values  # Get the datetime64 value from da1
        month_1 = T_value_1.astype('datetime64[M]').astype(int) % 12 + 1  # Extract month
        new_T_value = np.datetime64(f"{year}-{month_1:02d}-{1:02d}")
        forecast_expanded = result_.assign_coords(T=xr.DataArray([new_T_value], dims=["T"]))
        forecast_expanded['T'] = forecast_expanded['T'].astype('datetime64[ns]')
        
        
        dof = max(int(rainfall_for_tercile.sizes["T"]) - 1, 2)

        dm = self.dist_method

        # ---------- BESTFIT ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da, best_loc_da, best_scale_da."
                )
            
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                forecast_expanded,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                kwargs={"dof": dof},
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                forecast_expanded,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")
        forecast_prob = forecast_prob.assign_coords(probability=('probability', ['PB', 'PN', 'PA']))
        return forecast_expanded * mask, mask * forecast_prob.transpose('probability', 'T', 'Y', 'X')

    # --------------------------------------------------------------------------
    #  Visualization and Utility Methods
    # --------------------------------------------------------------------------
    
    def get_components(self):
        """Get spatial patterns (for linear kernel only)."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        if self.kernel != 'linear':
            print("Warning: Components not directly interpretable for non-linear kernels")
            return None, None
        
        # Extract weights for linear CCA
        weights_X = self.model.weights[0]  # (n_features_X, n_modes)
        weights_Y = self.model.weights[1]  # (n_features_Y, n_modes)
        
        # Reshape to spatial patterns
        X_components = self._unstack_spatial(weights_X.T, self.Y_coords, 
                                            np.arange(self.n_modes))
        Y_components = self._unstack_spatial(weights_Y.T, self.Y_coords,
                                            np.arange(self.n_modes))
        
        return X_components, Y_components

    def get_scores(self):
        """Get canonical variates (scores)."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        # Transform training data
        scores = self.model.transform([self.X_train_np, self.Y_train_np])
        
        # Create DataArrays
        time_coords = self.X_train_original.coords['T']
        X_scores = xr.DataArray(
            scores[0],
            dims=['T', 'mode'],
            coords={'T': time_coords, 'mode': range(1, self.n_modes + 1)}
        )
        Y_scores = xr.DataArray(
            scores[1],
            dims=['T', 'mode'],
            coords={'T': time_coords, 'mode': range(1, self.n_modes + 1)}
        )
        
        return X_scores, Y_scores

    def plot_cca_results(self, X=None, Y=None, n_modes=None, clim_year_start=None, 
                         clim_year_end=None, best_params=None):
        """
        Plot CCA modes and scores.
        
        Parameters:
        - best_params: Hyperparameters to use for fitting the model.
        """
        if X is not None and Y is not None:
            # Fit model with provided data
            X_ = standardize_timeseries(X, clim_year_start, clim_year_end)
            Y_ = standardize_timeseries(Y, clim_year_start, clim_year_end)
            self.fit_cca(X_, Y_, best_params=best_params)
        elif not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        # Get components and scores
        X_components, Y_components = self.get_components()
        X_scores, Y_scores = self.get_scores()
        
        if X_components is None:
            print("Cannot plot components for non-linear kernel")
            return
        
        # Determine number of modes to plot
        if n_modes is None:
            n_modes = min(self.n_modes, 3)
        
        # Create mask for Y
        if Y is not None:
            mask = xr.where(~np.isnan(Y.isel(T=0)), 1, np.nan)
            Y_components = Y_components * mask
        
        # Create figure
        fig, axes = plt.subplots(n_modes, 3, figsize=(15, 3 * n_modes))
        if n_modes == 1:
            axes = np.array([axes])
        
        for i in range(n_modes):
            mode = i + 1
            
            # X component
            ax = axes[i, 0]
            X_comp = X_components.isel(mode=i)
            if hasattr(X_comp, 'plot'):
                X_comp.plot(ax=ax, vmin=-1, vmax=1, cmap="RdBu_r")
            else:
                im = ax.imshow(X_comp, cmap="RdBu_r", vmin=-1, vmax=1)
                plt.colorbar(im, ax=ax)
            ax.set_title(f'X Mode {mode}')
            
            # Scores
            ax = axes[i, 1]
            X_score = X_scores.sel(mode=mode)
            Y_score = Y_scores.sel(mode=mode)
            
            if hasattr(X_score['T'], 'dt'):
                time_x = X_score['T'].dt.year.values
                time_y = Y_score['T'].dt.year.values
            else:
                time_x = np.arange(len(X_score))
                time_y = np.arange(len(Y_score))
            
            ax.plot(time_x, X_score, label='X Score')
            ax.plot(time_y, Y_score, label='Y Score')
            ax.axhline(0, linestyle='--', lw=0.8, color='gray')
            ax.legend()
            ax.set_title(f'Scores for Mode {mode}')
            ax.set_xlabel('Time')
            ax.set_ylabel('Canonical Variate')
            
            # Y component
            ax = axes[i, 2]
            Y_comp = Y_components.isel(mode=i)
            if hasattr(Y_comp, 'plot'):
                Y_comp.plot(ax=ax, cmap="RdBu_r")
            else:
                im = ax.imshow(Y_comp, cmap="RdBu_r")
                plt.colorbar(im, ax=ax)
            ax.set_title(f'Y Mode {mode}')
        
        plt.tight_layout()
        plt.show()

    def get_canonical_correlations(self):
        """Get canonical correlations."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        return self.canonical_correlations
    
    def get_model_info(self):
        """Print model information."""
        info = {
            'n_modes': self.n_modes,
            'kernel': self.kernel,
            'dist_method': self.dist_method,
            'is_fitted': self.is_fitted,
            'is_tuned': self.is_tuned,
            'best_params': self.best_params,
        }
        
        if self.is_fitted:
            info['canonical_correlations'] = self.canonical_correlations
            info['training_samples'] = self.X_train_np.shape[0]
            info['X_features'] = self.X_train_np.shape[1]
            info['Y_features'] = self.Y_train_np.shape[1]
        
        return info



class WAS_KGCCA:
    def __init__(self, n_modes=4, kernel='linear', dist_method="nonparam", 
                 cv_folds=5, random_state=42, search_method='grid'):
        """
        Initialize the WAS_KGCCA class with Kernel Generalized CCA.
        
        Parameters:
        - n_modes: Number of canonical modes to compute.
        - kernel: Kernel type ('linear', 'rbf', 'poly', 'sigmoid').
        - dist_method: Method for probability calculation ('bestfit', 'nonparam', 'gamma').
        - cv_folds: Number of cross-validation folds for hyperparameter tuning.
        - random_state: Random seed for reproducibility.
        - search_method: Hyperparameter search method ('grid', 'random', 'bayesian').
        """
        self.n_modes = n_modes
        self.kernel = kernel
        self.dist_method = dist_method
        self.cv_folds = cv_folds
        self.random_state = random_state
        self.search_method = search_method
        
        # Hyperparameter storage (tuned once, used everywhere)
        self.best_params = None  # Will store tuned parameters
        self.is_tuned = False    # Whether parameters have been tuned
        
        # Model and data storage
        self.model = None
        self.is_fitted = False
        
        # Data storage
        self.Y_coords = None
        self.X_mean = None
        self.Y_mean = None
        self.X_train_original = None
        self.Y_train_original = None
        self.X_train_np = None
        self.Y_train_np = None

    # --------------------------------------------------------------------------
    #  Hyperparameter Tuning (Do this once!)
    # --------------------------------------------------------------------------
    
    def _prepare_param_grid(self, kernel=None):
        """
        Prepare parameter grid for hyperparameter tuning.
        
        Parameters:
        - kernel: Kernel type (if None, use self.kernel)
        
        Returns:
        - param_grid: Dictionary of parameter distributions
        """
        if kernel is None:
            kernel = self.kernel
        
        if kernel == 'rbf':
            return {
                'gamma': [[0.01, 0.01], [0.1, 0.1], [1.0, 1.0], [10.0, 10.0]],
                'c': [0, 0.1, 0.5, 1.0]
            }
        elif kernel == 'poly':
            return {
                'degree': [[2, 2], [3, 3]],
                'gamma': [[0.1, 0.1], [1.0, 1.0]],
                'coef0': [[0, 0], [1, 1]]
            }
        elif kernel == 'sigmoid':
            return {
                'gamma': [[0.01, 0.01], [0.1, 0.1], [1.0, 1.0]],
                'coef0': [[0, 0], [1, 1]]
            }
        else:  # linear kernel
            return {'c': [0, 0.1, 0.5, 1.0]}
    
    def _objective(self, trial, X_np, Y_np):
        """
        Objective function for Bayesian optimization (if using Optuna).
        """
        if self.kernel == 'rbf':
            gamma_val = trial.suggest_float('gamma', 0.01, 10.0, log=True)
            gamma_params = [gamma_val, gamma_val]  # Same for both views
            c_val = trial.suggest_float('c', 0, 1.0)
            params = {'gamma': gamma_params, 'c': c_val}
        elif self.kernel == 'poly':
            degree = trial.suggest_int('degree', 2, 4)
            gamma_val = trial.suggest_float('gamma', 0.1, 5.0, log=True)
            coef0 = trial.suggest_float('coef0', 0, 5)
            params = {
                'degree': [degree, degree],
                'gamma': [gamma_val, gamma_val],
                'coef0': [coef0, coef0]
            }
        elif self.kernel == 'sigmoid':
            gamma_val = trial.suggest_float('gamma', 0.01, 1.0, log=True)
            coef0 = trial.suggest_float('coef0', 0, 5)
            params = {
                'gamma': [gamma_val, gamma_val],
                'coef0': [coef0, coef0]
            }
        else:  # linear
            c_val = trial.suggest_float('c', 0, 1.0)
            params = {'c': c_val}
        
        # Create KGCCA model with suggested parameters
        model = KGCCA(latent_dimensions=self.n_modes, **params)
        
        # Cross-validation
        cv = KFold(n_splits=self.cv_folds, shuffle=True, random_state=self.random_state)
        correlations = []
        
        for train_idx, val_idx in cv.split(X_np):
            X_train, X_val = X_np[train_idx], X_np[val_idx]
            Y_train, Y_val = Y_np[train_idx], Y_np[val_idx]
            
            model.fit([X_train, Y_train])
            scores = model.transform([X_val, Y_val])
            
            # Calculate canonical correlation
            corr = np.corrcoef(scores[0][:, 0], scores[1][:, 0])[0, 1]
            correlations.append(corr)
        
        return np.mean(correlations)
    
    def compute_hyperparameters(self, Predictors, Predictand, clim_year_start=None, clim_year_end=None):
        """
        Independently computes the best hyperparameters using selected optimization method.
        
        Parameters:
        - Predictors: Training predictor data with dimensions (T, Y, X).
        - Predictand: Training predictand data with dimensions (T, Y, X).
        - clim_year_start, clim_year_end: Optional climatology period for standardization.
        
        Returns:
        - best_params: Dictionary of best hyperparameters found.
        """
        if self.kernel == 'linear' and self.search_method != 'grid':
            print("Linear kernel has limited hyperparameters to tune (only 'c').")
        
        if not HAS_CCA_ZOO:
            raise ImportError("cca-zoo is required for hyperparameter tuning")
        
        print(f"Computing hyperparameters for {self.kernel} KGCCA using {self.search_method} search...")
        
        # Preprocess data (no standardization needed as input is already standardized)
        X_train_final, y_train_final = self.preprocess_data(Predictors, Predictand)
        
        # Stack spatial dimensions
        X_np, X_coords = self._stack_spatial(X_train_final)
        Y_np, Y_coords = self._stack_spatial(y_train_final)
        
        # Remove any remaining NaN values
        mask = ~(np.any(~np.isfinite(X_np), axis=1) | np.any(~np.isfinite(Y_np), axis=1))
        X_np_clean = X_np[mask]
        Y_np_clean = Y_np[mask]
        
        if self.search_method == 'grid':
            # Grid Search
            param_grid = self._prepare_param_grid()
            
            print(f"Performing grid search with parameters: {param_grid}")
            
            cv = KFold(n_splits=self.cv_folds, shuffle=True, random_state=self.random_state)
            
            grid_search = CCA_GridSearch(
                KGCCA(latent_dimensions=self.n_modes),
                param_grid=param_grid,
                cv=cv,
                n_jobs=-1,
                verbose=1
            )
            
            grid_search.fit([X_np_clean, Y_np_clean])
            best_params = grid_search.best_params_
            best_score = grid_search.best_score_
            
            print(f"Best parameters found: {best_params}")
            print(f"Best cross-validation score: {best_score:.4f}")
            
            self.best_params = best_params
            self.is_tuned = True
            
        elif self.search_method == 'bayesian':
            # Bayesian Optimization with Optuna
            try:
                import optuna
                from optuna.samplers import TPESampler
            except ImportError:
                raise ImportError("Optuna is required for Bayesian optimization")
            
            print("Performing Bayesian optimization with Optuna...")
            
            sampler = TPESampler(seed=self.random_state)
            study = optuna.create_study(
                direction='maximize',
                sampler=sampler
            )
            
            # Create objective function with data
            objective_with_data = lambda trial: self._objective(trial, X_np_clean, Y_np_clean)
            study.optimize(objective_with_data, n_trials=50)
            
            # Get best parameters
            best_params = study.best_params
            best_score = study.best_value
            
            print(f"Best parameters found: {best_params}")
            print(f"Best cross-validation score: {best_score:.4f}")
            
            self.best_params = best_params
            self.is_tuned = True
            
        else:  # Default parameters
            print("Using default parameters (no search performed)")
            self.best_params = {}
            self.is_tuned = True
        
        return self.best_params
    
    def set_hyperparameters(self, hyperparams):
        """
        Manually set hyperparameters.
        
        Parameters:
        - hyperparams: Dictionary of hyperparameters to use
        """
        self.best_params = hyperparams
        self.is_tuned = True
        print(f"Hyperparameters set: {hyperparams}")
    
    # --------------------------------------------------------------------------
    #  Data Preparation Methods
    # --------------------------------------------------------------------------
    
    def _stack_spatial(self, da):
        """Convert spatial DataArray to 2D array (time × features)."""
        # Check dimension names
        if 'X' in da.dims and 'Y' in da.dims:
            da = da.rename({"X": "lon", "Y": "lat"})
        
        # Stack spatial dimensions
        stacked = da.stack(feature=['lat', 'lon'])
        return stacked.values, stacked.coords

    def _unstack_spatial(self, data_np, coords, time_coords):
        """Reconstruct spatial DataArray from 2D array."""
        if data_np.ndim == 1:
            data_np = data_np.reshape(1, -1)
            
        da_stacked = xr.DataArray(
            data_np,
            dims=['T', 'feature'],
            coords={
                'T': time_coords,
                'feature': coords['feature']
            }
        )
        
        da = da_stacked.unstack('feature')
        return da.rename({"lon": "X", "lat": "Y"})

    def preprocess_data(self, X, Y):
        """Preprocess data by filling missing values."""
        # Fill missing values
        X_filled = X.fillna(X.mean(dim="T", skipna=True))
        Y_filled = Y.fillna(Y.mean(dim="T", skipna=True))
        
        # Rename and transpose
        X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        Y_final = Y_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        
        return X_final, Y_final

    def preprocess_test_data(self, X_test, y_test, X_train, y_train):
        """Preprocess test data using training statistics."""
        # Fill missing values with training means
        X_test_prepared = X_test.fillna(X_train.mean(dim="T", skipna=True))
        y_test_prepared = y_test.fillna(y_train.mean(dim="T", skipna=True))
        
        # Rename and transpose
        dims_rename = {"X": "lon", "Y": "lat"}
        X_test_prepared = X_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')
        y_test_prepared = y_test_prepared.rename(dims_rename).transpose('T', 'lat', 'lon')
        
        return X_test_prepared, y_test_prepared

    # --------------------------------------------------------------------------
    #  Model Fitting with Explicit Parameters
    # --------------------------------------------------------------------------
    
    def fit_cca(self, X_train, y_train, best_params=None):
        """
        Fit the KGCCA model using provided hyperparameters.
        
        Parameters:
        - X_train, y_train: Training data
        - best_params: Dictionary of hyperparameters to use. If None, uses self.best_params
        """
        # Use provided parameters if available, otherwise stored parameters
        if best_params is not None:
            params_to_use = best_params
            print(f"Using provided hyperparameters: {best_params}")
        elif self.best_params is not None:
            params_to_use = self.best_params
            print(f"Using stored hyperparameters: {self.best_params}")
        else:
            params_to_use = {}
            if self.kernel != 'linear':
                print("Warning: No hyperparameters provided. Using default parameters.")
        
        # Preprocess (no standardization)
        X_train_final, y_train_final = self.preprocess_data(X_train, y_train)
        
        # Store original data
        self.X_train_original = X_train_final
        self.Y_train_original = y_train_final
        
        # Stack spatial dimensions
        X_np, X_coords = self._stack_spatial(X_train_final)
        Y_np, Y_coords = self._stack_spatial(y_train_final)
        
        self.Y_coords = Y_coords
        self.X_train_np = X_np
        self.Y_train_np = Y_np
        
        # Build model with appropriate parameters
        if self.kernel == 'linear':
            # For linear kernel, we can use LinearCCA for efficiency
            if 'c' in params_to_use:
                self.model = GCCA(latent_dimensions=self.n_modes, c=params_to_use['c'])
            else:
                self.model = GCCA(latent_dimensions=self.n_modes)
        else:
            # For non-linear kernels, use KGCCA
            self.model = KGCCA(latent_dimensions=self.n_modes, **params_to_use)
        
        # Fit model
        self.model.fit([X_np, Y_np])
        self.is_fitted = True
        
        # Store canonical correlations
        if hasattr(self.model, 'canonical_correlations'):
            self.canonical_correlations = self.model.canonical_correlations()
        else:
            scores = self.model.transform([X_np, Y_np])
            self.canonical_correlations = [
                np.corrcoef(scores[0][:, i], scores[1][:, i])[0, 1] 
                for i in range(self.n_modes)
            ]
    
    # --------------------------------------------------------------------------
    #  Prediction Methods
    # --------------------------------------------------------------------------
    
    def _predict_kgcca(self, X_test):
        """Make predictions using KGCCA model."""
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit_cca first.")
        
        # Stack spatial dimensions
        X_test_np, X_coords = self._stack_spatial(X_test)
        
        # Transform to get Y scores
        Y_scores = self.model.transform([X_test_np, np.zeros_like(self.Y_train_np[:X_test_np.shape[0]])])[1]
        
        # Reconstruct spatial array
        Y_pred = self._unstack_spatial(Y_scores, self.Y_coords, X_test['T'].values)
        
        return Y_pred

    def compute_model(self, X_train, y_train, X_test, y_test, best_params=None):
        """
        Compute deterministic hindcast using KGCCA with injected hyperparameters.
        
        Parameters:
        - X_train: Training predictor data with dimensions (T, Y, X).
        - y_train: Training predictand data with dimensions (T, Y, X).
        - X_test: Testing predictor data with dimensions (T, Y, X).
        - y_test: Testing predictand data with dimensions (T, Y, X).
        - best_params: Pre-computed best hyperparameters. If None, uses stored parameters.
        
        Returns:
        - hindcast: Deterministic hindcast with dimensions (T, Y, X).
        """
        # Fit model with parameters
        self.fit_cca(X_train, y_train, best_params=best_params)
        
        # Preprocess test data
        X_test_prepared, y_test_prepared = self.preprocess_test_data(
            X_test, y_test, X_train, y_train
        )
        
        # Predict
        y_pred = self._predict_kgcca(X_test_prepared)
        y_pred['T'] = y_test_prepared['T']
        
        return y_pred

    # --------------------------------------------------------------------------
    #  Probability Calculation Methods
    # --------------------------------------------------------------------------
    # ------------------ Probability Calculation Methods ------------------

    @staticmethod
    def _ppf_terciles_from_code(dist_code, shape, loc, scale):
        """
        Return tercile thresholds (T1, T2) from best-fit distribution parameters.
    
        dist_code:
            1: norm
            2: lognorm
            3: expon
            4: gamma
            5: weibull_min
            6: t
            7: poisson
            8: nbinom
        """
        if np.isnan(dist_code):
            return np.nan, np.nan
    
        code = int(dist_code)
        try:
            if code == 1:
                return (
                    norm.ppf(0.32, loc=loc, scale=scale),
                    norm.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 2:
                return (
                    lognorm.ppf(0.32, s=shape, loc=loc, scale=scale),
                    lognorm.ppf(0.67, s=shape, loc=loc, scale=scale),
                )
            elif code == 3:
                return (
                    expon.ppf(0.32, loc=loc, scale=scale),
                    expon.ppf(0.67, loc=loc, scale=scale),
                )
            elif code == 4:
                return (
                    gamma.ppf(0.32, a=shape, loc=loc, scale=scale),
                    gamma.ppf(0.67, a=shape, loc=loc, scale=scale),
                )
            elif code == 5:
                return (
                    weibull_min.ppf(0.32, c=shape, loc=loc, scale=scale),
                    weibull_min.ppf(0.67, c=shape, loc=loc, scale=scale),
                )
            elif code == 6:
                return (
                    t.ppf(0.32, df=shape, loc=loc, scale=scale),
                    t.ppf(0.67, df=shape, loc=loc, scale=scale),
                )
            elif code == 7:
                return (
                    poisson.ppf(0.32, mu=shape, loc=loc),
                    poisson.ppf(0.67, mu=shape, loc=loc),
                )
            elif code == 8:
                return (
                    nbinom.ppf(0.32, n=shape, p=scale, loc=loc),
                    nbinom.ppf(0.67, n=shape, p=scale, loc=loc),
                )
        except Exception:
            return np.nan, np.nan
    
        return np.nan, np.nan
        
    @staticmethod
    def weibull_shape_solver(k, M, V):
        """
        Function to find the root of the Weibull shape parameter 'k'.
        We find 'k' such that the theoretical variance/mean^2 ratio
        matches the observed V/M^2 ratio.
        """
        if k <= 0:
            return -np.inf
        try:
            g1 = gamma_function(1 + 1/k)
            g2 = gamma_function(1 + 2/k)
            
            implied_v_over_m_sq = (g2 / (g1**2)) - 1
            observed_v_over_m_sq = V / (M**2)
            
            return observed_v_over_m_sq - implied_v_over_m_sq
        except ValueError:
            return -np.inf

    @staticmethod
    def calculate_tercile_probabilities_bestfit(best_guess, error_variance, T1, T2, dist_code, dof):
        """
        Generic tercile probabilities using best-fit family per grid cell.
        """
        best_guess = np.asarray(best_guess, float)
        error_variance = np.asarray(error_variance, dtype=float)
        n_time = best_guess.size
        out = np.full((3, n_time), np.nan, float)

        if np.all(np.isnan(best_guess)) or np.isnan(dist_code) or np.isnan(T1) or np.isnan(T2) or np.isnan(error_variance):
            return out

        code = int(dist_code)

        if code == 1:  # Normal
            error_std = np.sqrt(error_variance)
            out[0, :] = norm.cdf(T1, loc=best_guess, scale=error_std)
            out[1, :] = norm.cdf(T2, loc=best_guess, scale=error_std) - norm.cdf(T1, loc=best_guess, scale=error_std)
            out[2, :] = 1 - norm.cdf(T2, loc=best_guess, scale=error_std)

        elif code == 2:  # Lognormal
            sigma = np.sqrt(np.log(1 + error_variance / (best_guess**2)))
            mu = np.log(best_guess) - sigma**2 / 2
            out[0, :] = lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[1, :] = lognorm.cdf(T2, s=sigma, scale=np.exp(mu)) - lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[2, :] = 1 - lognorm.cdf(T2, s=sigma, scale=np.exp(mu))

        elif code == 3:  # Exponential
            scale_param = np.sqrt(error_variance)
            out[0, :] = expon.cdf(T1, loc=best_guess, scale=scale_param)
            out[1, :] = expon.cdf(T2, loc=best_guess, scale=scale_param) - expon.cdf(T1, loc=best_guess, scale=scale_param)
            out[2, :] = 1.0 - expon.cdf(T2, loc=best_guess, scale=scale_param)

        elif code == 4:  # Gamma
            alpha = (best_guess ** 2) / error_variance
            theta = error_variance / best_guess
            out[0, :] = gamma.cdf(T1, a=alpha, scale=theta)
            out[1, :] = gamma.cdf(T2, a=alpha, scale=theta) - gamma.cdf(T1, a=alpha, scale=theta)
            out[2, :] = 1.0 - gamma.cdf(T2, a=alpha, scale=theta)

        elif code == 6:  # Student-t
            if dof <= 2:
                out[0, :] = np.nan
                out[1, :] = np.nan
                out[2, :] = np.nan
            else:
                scale = np.sqrt(error_variance * (dof - 2) / dof)
                out[0, :] = t.cdf(T1, df=dof, loc=best_guess, scale=scale)
                out[1, :] = t.cdf(T2, df=dof, loc=best_guess, scale=scale) - t.cdf(T1, df=dof, loc=best_guess, scale=scale)
                out[2, :] = 1.0 - t.cdf(T2, df=dof, loc=best_guess, scale=scale)

        elif code == 7:  # Poisson
            if not np.allclose(best_guess, error_variance, atol=0.5):
                print("Warning: 'error_variance' is not equal to 'best_guess'. Poisson model assumes mean=variance.")
            
            mu = best_guess
            out[0, :] = poisson.cdf(T1, mu=mu)
            out[1, :] = poisson.cdf(T2, mu=mu) - poisson.cdf(T1, mu=mu)
            out[2, :] = 1.0 - poisson.cdf(T2, mu=mu)

        elif code == 8:  # Negative Binomial
            p = np.where(error_variance > best_guess, 
                         best_guess / error_variance, 
                         np.nan)
            n = np.where(error_variance > best_guess,
                         (best_guess**2) / (error_variance - best_guess), 
                         np.nan)
            out[0, :] = nbinom.cdf(T1, n=n, p=p)
            out[1, :] = nbinom.cdf(T2, n=n, p=p) - nbinom.cdf(T1, n=n, p=p)
            out[2, :] = 1.0 - nbinom.cdf(T2, n=n, p=p)
            
        else:
            raise ValueError(f"Invalid distribution code: {code}")

        return out

    @staticmethod
    def calculate_tercile_probabilities_nonparametric(best_guess, error_samples, first_tercile, second_tercile):
        """Non-parametric method using historical error samples."""
        n_time = len(best_guess)
        pred_prob = np.full((3, n_time), np.nan, dtype=float)
        for t in range(n_time):
            if np.isnan(best_guess[t]):
                continue
            dist = best_guess[t] + error_samples
            dist = dist[np.isfinite(dist)]
            if len(dist) == 0:
                continue
            p_below = np.mean(dist < first_tercile)
            p_between = np.mean((dist >= first_tercile) & (dist < second_tercile))
            p_above = 1.0 - (p_below + p_between)
            pred_prob[0, t] = p_below
            pred_prob[1, t] = p_between
            pred_prob[2, t] = p_above
        return pred_prob

    def compute_prob(
        self,
        Predictant: xr.DataArray,
        clim_year_start,
        clim_year_end,
        hindcast_det: xr.DataArray,
        best_code_da: xr.DataArray = None,
        best_shape_da: xr.DataArray = None,
        best_loc_da: xr.DataArray = None,
        best_scale_da: xr.DataArray = None
    ) -> xr.DataArray:
        """
        Compute tercile probabilities for deterministic hindcasts.
        """
        # Handle member dimension if present
        if "M" in Predictant.dims:
            Predictant = Predictant.isel(M=0).drop_vars("M").squeeze()

        # Ensure dimension order
        Predictant = Predictant.transpose("T", "Y", "X")

        # Spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1.0, np.nan)

        # Climatology subset
        clim = Predictant.sel(T=slice(str(clim_year_start), str(clim_year_end)))
        if clim.sizes.get("T", 0) < 3:
            raise ValueError("Not enough years in climatology period for terciles.")

        # Error variance for predictive distributions
        error_variance = (Predictant - hindcast_det).var(dim="T")
        dof = max(int(clim.sizes["T"]) - 1, 2)

        # Empirical terciles (used by non-bestfit methods)
        terciles_emp = clim.quantile([0.32, 0.67], dim="T")
        T1_emp = terciles_emp.isel(quantile=0).drop_vars("quantile")
        T2_emp = terciles_emp.isel(quantile=1).drop_vars("quantile")
        

        dm = self.dist_method

        # ---------- BESTFIT ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da_da, best_loc_da, best_scale_da."
                )

            # T1, T2 from best-fit distributions (per grid)
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            # Predictive probabilities using same family
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                hindcast_det,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                kwargs={'dof': dof},
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                hindcast_det,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")

        hindcast_prob = hindcast_prob.assign_coords(
            probability=("probability", ["PB", "PN", "PA"])
        )
        return (hindcast_prob * mask).transpose("probability", "T", "Y", "X")

    # --------------------------------------------------------------------------
    #  Forecast Method (with explicit parameters)
    # --------------------------------------------------------------------------
    
    def forecast(self, Predictant, clim_year_start, clim_year_end, Predictor, 
                 hindcast_det, Predictor_for_year, best_params=None,
                 best_code_da=None, best_shape_da=None, best_loc_da=None, 
                 best_scale_da=None):
        """
        Generate operational forecast with probabilities using specified hyperparameters.
        """
        # Create spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1, np.nan)
        
        # Fit model with specified parameters
        Predictor_ = (Predictor - trend_data(Predictor).fillna(
            trend_data(Predictor)[-3])).fillna(0)
        Predictant_st = standardize_timeseries(Predictant, clim_year_start, clim_year_end)
        Predictant_ = (Predictant_st - trend_data(Predictant_st).fillna(
            trend_data(Predictant_st)[-3])).fillna(0)
        
        self.fit_cca(Predictor_, Predictant_, best_params=best_params)
        
        # Prepare forecast predictor
        Predictor_for_year_ = Predictor_for_year.fillna(Predictor.mean(dim="T", skipna=True))
        Predictor_for_year_ = Predictor_for_year_.transpose('T', 'Y', 'X')
        
        # Make deterministic forecast
        result_det = self._predict_kgcca(Predictor_for_year_)
        
        # Apply reverse standardization if needed
        if hasattr(self, '_standardized'):
            result_det = reverse_standardize(result_det, Predictant, clim_year_start, clim_year_end)
        
        # 2) Compute thresholds T1, T2 from climatology
        index_start = Predictant.get_index("T").get_loc(str(clim_year_start)).start
        index_end   = Predictant.get_index("T").get_loc(str(clim_year_end)).stop
        rainfall_for_tercile = Predictant.isel(T=slice(index_start, index_end))
        terciles = rainfall_for_tercile.quantile([0.32, 0.67], dim='T')
        T1_emp = terciles.isel(quantile=0).drop_vars('quantile')
        T2_emp = terciles.isel(quantile=1).drop_vars('quantile')
        error_variance = (Predictant - hindcast_det).var(dim='T')

        year = Predictor_for_year.coords['T'].values.astype('datetime64[Y]').astype(int)[0] + 1970
        T_value_1 = Predictant.isel(T=0).coords['T'].values
        month_1 = T_value_1.astype('datetime64[M]').astype(int) % 12 + 1
        new_T_value = np.datetime64(f"{year}-{month_1:02d}-01")
        
        forecast_expanded = result_det.assign_coords(T=xr.DataArray([new_T_value], dims=["T"]))
        forecast_expanded['T'] = forecast_expanded['T'].astype('datetime64[ns]')
        
        dof = max(int(rainfall_for_tercile.sizes["T"]) - 1, 2)

        dm = self.dist_method

        # ---------- BESTFIT ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError(
                    "dist_method='bestfit' requires best_code_da, best_shape_da, best_loc_da, best_scale_da."
                )
            
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                forecast_expanded,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                kwargs={"dof": dof},
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            forecast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                forecast_expanded,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")
        
        forecast_prob = forecast_prob.assign_coords(probability=('probability', ['PB', 'PN', 'PA']))
        return forecast_expanded * mask, mask * forecast_prob.transpose('probability', 'T', 'Y', 'X')

    # --------------------------------------------------------------------------
    #  Visualization and Utility Methods
    # --------------------------------------------------------------------------
    
    def get_components(self):
        """Get spatial patterns (for linear kernel only)."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        if self.kernel != 'linear':
            print("Warning: Components not directly interpretable for non-linear kernels")
            return None, None
        
        # Extract weights for linear CCA
        weights = self.model.weights
        weights_X = weights[0]  # (n_features_X, n_modes)
        weights_Y = weights[1]  # (n_features_Y, n_modes)
        
        # Reshape to spatial patterns
        X_components = self._unstack_spatial(weights_X.T, self.Y_coords, 
                                            np.arange(self.n_modes))
        Y_components = self._unstack_spatial(weights_Y.T, self.Y_coords,
                                            np.arange(self.n_modes))
        
        return X_components, Y_components

    def get_scores(self):
        """Get canonical variates (scores)."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        # Transform training data
        scores = self.model.transform([self.X_train_np, self.Y_train_np])
        
        # Create DataArrays
        time_coords = self.X_train_original.coords['T']
        X_scores = xr.DataArray(
            scores[0],
            dims=['T', 'mode'],
            coords={'T': time_coords, 'mode': range(1, self.n_modes + 1)}
        )
        Y_scores = xr.DataArray(
            scores[1],
            dims=['T', 'mode'],
            coords={'T': time_coords, 'mode': range(1, self.n_modes + 1)}
        )
        
        return X_scores, Y_scores

    def plot_cca_results(self, X=None, Y=None, n_modes=None, clim_year_start=None, 
                         clim_year_end=None, best_params=None):
        """
        Plot CCA modes and scores.
        """
        if X is not None and Y is not None:
            # Fit model with provided data
            X_ = standardize_timeseries(X, clim_year_start, clim_year_end)
            Y_ = standardize_timeseries(Y, clim_year_start, clim_year_end)
            self.fit_cca(X_, Y_, best_params=best_params)
        elif not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        # Get components and scores
        X_components, Y_components = self.get_components()
        X_scores, Y_scores = self.get_scores()
        
        if X_components is None:
            print("Cannot plot components for non-linear kernel")
            return
        
        # Determine number of modes to plot
        if n_modes is None:
            n_modes = min(self.n_modes, 3)
        
        # Create mask for Y
        if Y is not None:
            mask = xr.where(~np.isnan(Y.isel(T=0)), 1, np.nan)
            Y_components = Y_components * mask
        
        # Create figure
        fig, axes = plt.subplots(n_modes, 3, figsize=(15, 3 * n_modes))
        if n_modes == 1:
            axes = np.array([axes])
        
        for i in range(n_modes):
            mode = i + 1
            
            # X component
            ax = axes[i, 0]
            X_comp = X_components.isel(mode=i)
            if hasattr(X_comp, 'plot'):
                X_comp.plot(ax=ax, vmin=-1, vmax=1, cmap="RdBu_r")
            else:
                im = ax.imshow(X_comp, cmap="RdBu_r", vmin=-1, vmax=1)
                plt.colorbar(im, ax=ax)
            ax.set_title(f'X Mode {mode}')
            
            # Scores
            ax = axes[i, 1]
            X_score = X_scores.sel(mode=mode)
            Y_score = Y_scores.sel(mode=mode)
            
            if hasattr(X_score['T'], 'dt'):
                time_x = X_score['T'].dt.year.values
                time_y = Y_score['T'].dt.year.values
            else:
                time_x = np.arange(len(X_score))
                time_y = np.arange(len(Y_score))
            
            ax.plot(time_x, X_score, label='X Score')
            ax.plot(time_y, Y_score, label='Y Score')
            ax.axhline(0, linestyle='--', lw=0.8, color='gray')
            ax.legend()
            ax.set_title(f'Scores for Mode {mode}')
            ax.set_xlabel('Time')
            ax.set_ylabel('Canonical Variate')
            
            # Y component
            ax = axes[i, 2]
            Y_comp = Y_components.isel(mode=i)
            if hasattr(Y_comp, 'plot'):
                Y_comp.plot(ax=ax, cmap="RdBu_r")
            else:
                im = ax.imshow(Y_comp, cmap="RdBu_r")
                plt.colorbar(im, ax=ax)
            ax.set_title(f'Y Mode {mode}')
        
        plt.tight_layout()
        plt.show()

    def get_canonical_correlations(self):
        """Get canonical correlations."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        return self.canonical_correlations
    
    def get_model_info(self):
        """Print model information."""
        info = {
            'n_modes': self.n_modes,
            'kernel': self.kernel,
            'dist_method': self.dist_method,
            'is_fitted': self.is_fitted,
            'is_tuned': self.is_tuned,
            'best_params': self.best_params,
        }
        
        if self.is_fitted:
            info['canonical_correlations'] = self.canonical_correlations
            info['training_samples'] = self.X_train_np.shape[0]
            info['X_features'] = self.X_train_np.shape[1]
            info['Y_features'] = self.Y_train_np.shape[1]
        
        return info




class WAS_KGCCA_MultiPredictor:
    def __init__(self, n_modes=4, kernel='linear', dist_method="nonparam", 
                 cv_folds=5, random_state=42, search_method='grid',
                 concat_method='stack'):
        """
        Initialize the WAS_KGCCA class for multiple predictors.
        
        Parameters:
        - n_modes: Number of canonical modes to compute.
        - kernel: Kernel type ('linear', 'rbf', 'poly', 'sigmoid').
        - dist_method: Method for probability calculation ('bestfit', 'nonparam', 'gamma').
        - cv_folds: Number of cross-validation folds for hyperparameter tuning.
        - random_state: Random seed for reproducibility.
        - search_method: Hyperparameter search method ('grid', 'random', 'bayesian').
        - concat_method: How to handle multiple predictors ('stack', 'pca', 'weighted')
        """
        self.n_modes = n_modes
        self.kernel = kernel
        self.dist_method = dist_method
        self.cv_folds = cv_folds
        self.random_state = random_state
        self.search_method = search_method
        self.concat_method = concat_method
        
        # Hyperparameter storage
        self.best_params = None
        self.is_tuned = False
        
        # Model and data storage
        self.model = None
        self.is_fitted = False
        
        # Data storage
        self.Y_coords = None
        self.X_mean = None
        self.Y_mean = None
        self.X_train_original = None
        self.Y_train_original = None
        self.X_train_np = None
        self.Y_train_np = None
        
        # For multi-predictor handling
        self.predictor_names = None
        self.predictor_shapes = None
        self.pca_models = None
        self.feature_weights = None

    # --------------------------------------------------------------------------
    #  Multi-Predictor Handling Methods
    # --------------------------------------------------------------------------
    
    def _stack_spatial(self, da, name=None):
        """Convert spatial DataArray to 2D array (time × features)."""
        # Check dimension names
        if 'X' in da.dims and 'Y' in da.dims:
            da = da.rename({"X": "lon", "Y": "lat"})
        
        # Stack spatial dimensions
        stacked = da.stack(feature=['lat', 'lon'])
        return stacked.values, stacked.coords, da.shape
    
    def _unstack_spatial(self, data_np, coords, time_coords):
        """Reconstruct spatial DataArray from 2D array."""
        if data_np.ndim == 1:
            data_np = data_np.reshape(1, -1)
            
        da_stacked = xr.DataArray(
            data_np,
            dims=['T', 'feature'],
            coords={
                'T': time_coords,
                'feature': coords['feature']
            }
        )
        
        da = da_stacked.unstack('feature')
        return da.rename({"lon": "X", "lat": "Y"})
    
    def _prepare_predictors(self, predictors_dict, fit=True):
        """
        Prepare multiple predictors for CCA.
        
        Parameters:
        - predictors_dict: Dictionary of {name: DataArray} for each predictor
        - fit: Whether to fit transformations or just transform
        
        Returns:
        - X_combined: Combined predictor matrix (n_samples × n_features_total)
        - predictor_info: Dictionary with information about each predictor
        """
        if not isinstance(predictors_dict, dict):
            raise ValueError("predictors_dict should be a dictionary {name: DataArray}")
        
        predictor_names = list(predictors_dict.keys())
        n_predictors = len(predictor_names)
        
        # Store predictor names and shapes
        if fit:
            self.predictor_names = predictor_names
            self.predictor_shapes = {}
        
        # Process each predictor
        processed_predictors = []
        predictor_info = {
            'names': predictor_names,
            'start_indices': [],
            'end_indices': [],
            'n_features': []
        }
        
        current_idx = 0
        
        for i, (name, predictor) in enumerate(predictors_dict.items()):
            # Stack spatial dimensions
            X_np, X_coords, orig_shape = self._stack_spatial(predictor)
            
            if fit:
                self.predictor_shapes[name] = orig_shape
            
            # Apply method-specific processing
            if self.concat_method == 'pca':
                if fit:
                    from sklearn.decomposition import PCA
                    n_components = min(10, X_np.shape[1])  # Keep top 10 PCs
                    pca = PCA(n_components=n_components)
                    X_processed = pca.fit_transform(X_np)
                    self.pca_models[name] = pca
                else:
                    X_processed = self.pca_models[name].transform(X_np)
            elif self.concat_method == 'weighted':
                # Apply weighting based on domain knowledge
                if name.lower() in ['sst', 'sea_surface_temperature']:
                    weight = 1.5  # SST might be more important
                elif 'wind' in name.lower():
                    weight = 1.2
                else:
                    weight = 1.0
                X_processed = X_np * weight
            else:  # 'stack' - default
                X_processed = X_np
            
            # Store indices for this predictor
            n_features = X_processed.shape[1]
            predictor_info['start_indices'].append(current_idx)
            predictor_info['end_indices'].append(current_idx + n_features)
            predictor_info['n_features'].append(n_features)
            
            processed_predictors.append(X_processed)
            current_idx += n_features
        
        # Combine all predictors
        X_combined = np.hstack(processed_predictors)
        
        return X_combined, predictor_info
    
    def _prepare_single_predictor(self, predictor):
        """Prepare a single predictor (backward compatibility)."""
        X_np, X_coords, _ = self._stack_spatial(predictor)
        return X_np, {'names': ['single'], 'start_indices': [0], 
                      'end_indices': [X_np.shape[1]], 'n_features': [X_np.shape[1]]}

    # --------------------------------------------------------------------------
    #  Hyperparameter Tuning
    # --------------------------------------------------------------------------
    
    def _prepare_param_grid(self, kernel=None):
        """Prepare parameter grid for hyperparameter tuning."""
        if kernel is None:
            kernel = self.kernel
        
        if kernel == 'rbf':
            return {
                'gamma': [[0.01, 0.01], [0.1, 0.1], [1.0, 1.0], [10.0, 10.0]],
                'c': [0, 0.1, 0.5, 1.0]
            }
        elif kernel == 'poly':
            return {
                'degree': [[2, 2], [3, 3]],
                'gamma': [[0.1, 0.1], [1.0, 1.0]],
                'coef0': [[0, 0], [1, 1]]
            }
        elif kernel == 'sigmoid':
            return {
                'gamma': [[0.01, 0.01], [0.1, 0.1], [1.0, 1.0]],
                'coef0': [[0, 0], [1, 1]]
            }
        else:  # linear kernel
            return {'c': [0, 0.1, 0.5, 1.0]}
    
    def compute_hyperparameters(self, predictors_dict, Predictand, 
                                clim_year_start=None, clim_year_end=None):
        """
        Compute hyperparameters for multiple predictors.
        
        Parameters:
        - predictors_dict: Dictionary of {name: DataArray} for each predictor
        - Predictand: Predictand DataArray
        - clim_year_start, clim_year_end: Optional climatology period
        """
        if not HAS_CCA_ZOO:
            raise ImportError("cca-zoo is required for hyperparameter tuning")
        
        print(f"Computing hyperparameters for {self.kernel} KGCCA with {len(predictors_dict)} predictors...")
        
        # Preprocess data
        X_train_final, y_train_final = self.preprocess_data(predictors_dict, Predictand)
        
        # Prepare predictors
        X_combined, _ = self._prepare_predictors(X_train_final, fit=True)
        
        # Prepare predictand
        Y_np, Y_coords, _ = self._stack_spatial(y_train_final)
        
        # Remove NaN values
        mask = ~(np.any(~np.isfinite(X_combined), axis=1) | np.any(~np.isfinite(Y_np), axis=1))
        X_clean = X_combined[mask]
        Y_clean = Y_np[mask]
        
        # Hyperparameter search
        if self.search_method == 'grid':
            param_grid = self._prepare_param_grid()
            cv = KFold(n_splits=self.cv_folds, shuffle=True, random_state=self.random_state)
            
            grid_search = CCA_GridSearch(
                KGCCA(latent_dimensions=self.n_modes),
                param_grid=param_grid,
                cv=cv,
                n_jobs=-1,
                verbose=1
            )
            
            grid_search.fit([X_clean, Y_clean])
            self.best_params = grid_search.best_params_
            self.is_tuned = True
            
            print(f"Best parameters found: {self.best_params}")
            print(f"Best cross-validation score: {grid_search.best_score_:.4f}")
        
        elif self.search_method == 'bayesian':
            # Bayesian optimization
            try:
                import optuna
                from optuna.samplers import TPESampler
                
                def objective(trial):
                    if self.kernel == 'rbf':
                        gamma_val = trial.suggest_float('gamma', 0.01, 10.0, log=True)
                        params = {'gamma': [gamma_val, gamma_val]}
                    elif self.kernel == 'linear':
                        c_val = trial.suggest_float('c', 0, 1.0)
                        params = {'c': c_val}
                    else:
                        return 0.0
                    
                    model = KGCCA(latent_dimensions=self.n_modes, **params)
                    cv = KFold(n_splits=3, shuffle=True, random_state=self.random_state)
                    correlations = []
                    
                    for train_idx, val_idx in cv.split(X_clean):
                        X_train, X_val = X_clean[train_idx], X_clean[val_idx]
                        Y_train, Y_val = Y_clean[train_idx], Y_clean[val_idx]
                        
                        model.fit([X_train, Y_train])
                        scores = model.transform([X_val, Y_val])
                        corr = np.corrcoef(scores[0][:, 0], scores[1][:, 0])[0, 1]
                        correlations.append(corr)
                    
                    return np.mean(correlations)
                
                sampler = TPESampler(seed=self.random_state)
                study = optuna.create_study(direction='maximize', sampler=sampler)
                study.optimize(objective, n_trials=30)
                
                self.best_params = study.best_params
                self.is_tuned = True
                print(f"Best parameters found: {self.best_params}")
                
            except ImportError:
                print("Optuna not available, using default parameters")
                self.best_params = {}
                self.is_tuned = True
        
        else:
            self.best_params = {}
            self.is_tuned = True
        
        return self.best_params

    # --------------------------------------------------------------------------
    #  Data Preparation
    # --------------------------------------------------------------------------
    
    def preprocess_data(self, X, Y):
        """
        Preprocess multiple predictors and predictand.
        
        Parameters:
        - X: Dictionary of predictor DataArrays or single DataArray
        - Y: Predictand DataArray
        """
        # Handle predictors
        if isinstance(X, dict):
            X_processed = {}
            for name, predictor in X.items():
                X_filled = predictor.fillna(predictor.mean(dim="T", skipna=True))
                X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
                X_processed[name] = X_final
        else:
            # Single predictor for backward compatibility
            X_filled = X.fillna(X.mean(dim="T", skipna=True))
            X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
            X_processed = {'single': X_final}
        
        # Handle predictand
        Y_filled = Y.fillna(Y.mean(dim="T", skipna=True))
        Y_final = Y_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        
        return X_processed, Y_final
    
    def preprocess_test_data(self, X_test, y_test, X_train, y_train):
        """Preprocess test data using training statistics."""
        # Handle predictors
        if isinstance(X_test, dict):
            X_test_processed = {}
            for name, predictor in X_test.items():
                X_filled = predictor.fillna(X_train[name].mean(dim="T", skipna=True))
                X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
                X_test_processed[name] = X_final
        else:
            X_filled = X_test.fillna(X_train.mean(dim="T", skipna=True))
            X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
            X_test_processed = {'single': X_final}
        
        # Handle predictand
        y_filled = y_test.fillna(y_train.mean(dim="T", skipna=True))
        y_final = y_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        
        return X_test_processed, y_final

    # --------------------------------------------------------------------------
    #  Model Fitting
    # --------------------------------------------------------------------------
    
    def fit_cca(self, X_train, y_train, best_params=None):
        """
        Fit KGCCA model with multiple predictors.
        
        Parameters:
        - X_train: Dictionary of predictor DataArrays or single DataArray
        - y_train: Predictand DataArray
        - best_params: Hyperparameters to use
        """
        # Use provided parameters if available
        if best_params is not None:
            params_to_use = best_params
            print(f"Using provided hyperparameters: {best_params}")
        elif self.best_params is not None:
            params_to_use = self.best_params
            print(f"Using stored hyperparameters: {self.best_params}")
        else:
            params_to_use = {}
        
        # Preprocess data
        X_train_dict, y_train_final = self.preprocess_data(X_train, y_train)
        self.Y_train_original = y_train_final
        
        # Prepare predictors and predictand
        X_combined, predictor_info = self._prepare_predictors(X_train_dict, fit=True)
        Y_np, Y_coords, _ = self._stack_spatial(y_train_final)
        
        self.Y_coords = Y_coords
        self.X_train_np = X_combined
        self.Y_train_np = Y_np
        self.predictor_info = predictor_info
        
        # Build and fit model
        if self.kernel == 'linear':
            if 'c' in params_to_use:
                self.model = GCCA(latent_dimensions=self.n_modes, c=params_to_use['c'])
            else:
                self.model = GCCA(latent_dimensions=self.n_modes)
        else:
            self.model = KGCCA(latent_dimensions=self.n_modes, **params_to_use)
        
        self.model.fit([X_combined, Y_np])
        self.is_fitted = True
        
        # Store canonical correlations
        if hasattr(self.model, 'canonical_correlations'):
            self.canonical_correlations = self.model.canonical_correlations()
        else:
            scores = self.model.transform([X_combined, Y_np])
            self.canonical_correlations = [
                np.corrcoef(scores[0][:, i], scores[1][:, i])[0, 1] 
                for i in range(self.n_modes)
            ]
        
        # Calculate predictor importance
        self._calculate_predictor_importance()
    
    def _calculate_predictor_importance(self):
        """Calculate relative importance of each predictor."""
        if not self.is_fitted:
            return
        
        # Get model weights
        weights = self.model.weights[0]  # Predictor weights
        
        # Calculate importance for each predictor
        self.predictor_importance = {}
        total_importance = np.sum(np.abs(weights), axis=0).sum()
        
        for i, name in enumerate(self.predictor_info['names']):
            start = self.predictor_info['start_indices'][i]
            end = self.predictor_info['end_indices'][i]
            
            # Sum absolute weights for this predictor across all modes
            predictor_weights = weights[start:end, :]
            importance = np.sum(np.abs(predictor_weights)) / total_importance
            
            self.predictor_importance[name] = importance
        
        print("Predictor importance:")
        for name, importance in self.predictor_importance.items():
            print(f"  {name}: {importance:.3f}")

    # --------------------------------------------------------------------------
    #  Prediction Methods
    # --------------------------------------------------------------------------
    
    def _predict_kgcca(self, X_test_dict):
        """Make predictions using KGCCA model with multiple predictors."""
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit_cca first.")
        
        # Prepare test predictors
        X_combined, _ = self._prepare_predictors(X_test_dict, fit=False)
        
        # Transform to get Y scores
        # Need dummy array for second view during transform
        dummy_Y = np.zeros((X_combined.shape[0], self.Y_train_np.shape[1]))
        Y_scores = self.model.transform([X_combined, dummy_Y])[1]
        
        # Reconstruct spatial array
        time_coords = list(X_test_dict.values())[0]['T'].values
        Y_pred = self._unstack_spatial(Y_scores, self.Y_coords, time_coords)
        
        return Y_pred
    
    def compute_model(self, X_train, y_train, X_test, y_test, best_params=None):
        """
        Compute deterministic hindcast with multiple predictors.
        """
        # Fit model
        self.fit_cca(X_train, y_train, best_params=best_params)
        
        # Preprocess test data
        X_test_dict, y_test_final = self.preprocess_test_data(X_test, y_test, X_train, y_train)
        
        # Predict
        y_pred = self._predict_kgcca(X_test_dict)
        y_pred['T'] = y_test_final['T']
        
        return y_pred

    # --------------------------------------------------------------------------
    #  Probability Calculation Methods (Same as before)
    # --------------------------------------------------------------------------
    
    @staticmethod
    def _ppf_terciles_from_code(dist_code, shape, loc, scale):
        """Return tercile thresholds from best-fit distribution parameters."""
        if np.isnan(dist_code):
            return np.nan, np.nan
        
        code = int(dist_code)
        try:
            if code == 1:
                return norm.ppf(0.32, loc=loc, scale=scale), norm.ppf(0.67, loc=loc, scale=scale)
            elif code == 2:
                return (lognorm.ppf(0.32, s=shape, loc=loc, scale=scale),
                        lognorm.ppf(0.67, s=shape, loc=loc, scale=scale))
            elif code == 3:
                return (expon.ppf(0.32, loc=loc, scale=scale),
                        expon.ppf(0.67, loc=loc, scale=scale))
            elif code == 4:
                return (gamma.ppf(0.32, a=shape, loc=loc, scale=scale),
                        gamma.ppf(0.67, a=shape, loc=loc, scale=scale))
            elif code == 5:
                return (weibull_min.ppf(0.32, c=shape, loc=loc, scale=scale),
                        weibull_min.ppf(0.67, c=shape, loc=loc, scale=scale))
            elif code == 6:
                return (t.ppf(0.32, df=shape, loc=loc, scale=scale),
                        t.ppf(0.67, df=shape, loc=loc, scale=scale))
            elif code == 7:
                return (poisson.ppf(0.32, mu=shape, loc=loc),
                        poisson.ppf(0.67, mu=shape, loc=loc))
            elif code == 8:
                return (nbinom.ppf(0.32, n=shape, p=scale, loc=loc),
                        nbinom.ppf(0.67, n=shape, p=scale, loc=loc))
        except Exception:
            return np.nan, np.nan
        
        return np.nan, np.nan

    @staticmethod
    def calculate_tercile_probabilities_bestfit(best_guess, error_variance, T1, T2, dist_code, dof):
        """Calculate tercile probabilities using best-fit distribution family."""
        best_guess = np.asarray(best_guess, float)
        error_variance = np.asarray(error_variance, float)
        n_time = best_guess.size
        out = np.full((3, n_time), np.nan, float)

        if np.all(np.isnan(best_guess)) or np.isnan(dist_code):
            return out

        code = int(dist_code)

        if code == 1:  # Normal
            error_std = np.sqrt(error_variance)
            out[0, :] = norm.cdf(T1, loc=best_guess, scale=error_std)
            out[1, :] = norm.cdf(T2, loc=best_guess, scale=error_std) - norm.cdf(T1, loc=best_guess, scale=error_std)
            out[2, :] = 1 - norm.cdf(T2, loc=best_guess, scale=error_std)

        elif code == 2:  # Lognormal
            sigma = np.sqrt(np.log(1 + error_variance / (best_guess**2)))
            mu = np.log(best_guess) - sigma**2 / 2
            out[0, :] = lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[1, :] = lognorm.cdf(T2, s=sigma, scale=np.exp(mu)) - lognorm.cdf(T1, s=sigma, scale=np.exp(mu))
            out[2, :] = 1 - lognorm.cdf(T2, s=sigma, scale=np.exp(mu))

        elif code == 3:  # Exponential
            scale_param = np.sqrt(error_variance)
            out[0, :] = expon.cdf(T1, loc=best_guess, scale=scale_param)
            out[1, :] = expon.cdf(T2, loc=best_guess, scale=scale_param) - expon.cdf(T1, loc=best_guess, scale=scale_param)
            out[2, :] = 1.0 - expon.cdf(T2, loc=best_guess, scale=scale_param)

        elif code == 4:  # Gamma
            alpha = (best_guess ** 2) / error_variance
            theta = error_variance / best_guess
            out[0, :] = gamma.cdf(T1, a=alpha, scale=theta)
            out[1, :] = gamma.cdf(T2, a=alpha, scale=theta) - gamma.cdf(T1, a=alpha, scale=theta)
            out[2, :] = 1.0 - gamma.cdf(T2, a=alpha, scale=theta)

        elif code == 6:  # Student-t
            if dof <= 2:
                out[0, :] = np.nan
                out[1, :] = np.nan
                out[2, :] = np.nan
            else:
                scale = np.sqrt(error_variance * (dof - 2) / dof)
                out[0, :] = t.cdf(T1, df=dof, loc=best_guess, scale=scale)
                out[1, :] = t.cdf(T2, df=dof, loc=best_guess, scale=scale) - t.cdf(T1, df=dof, loc=best_guess, scale=scale)
                out[2, :] = 1.0 - t.cdf(T2, df=dof, loc=best_guess, scale=scale)

        return out

    @staticmethod
    def calculate_tercile_probabilities_nonparametric(best_guess, error_samples, first_tercile, second_tercile):
        """Non-parametric method using historical error samples."""
        n_time = len(best_guess)
        pred_prob = np.full((3, n_time), np.nan, dtype=float)
        for t in range(n_time):
            if np.isnan(best_guess[t]):
                continue
            dist = best_guess[t] + error_samples
            dist = dist[np.isfinite(dist)]
            if len(dist) == 0:
                continue
            p_below = np.mean(dist < first_tercile)
            p_between = np.mean((dist >= first_tercile) & (dist < second_tercile))
            p_above = 1.0 - (p_below + p_between)
            pred_prob[0, t] = p_below
            pred_prob[1, t] = p_between
            pred_prob[2, t] = p_above
        return pred_prob

    def compute_prob(self, Predictant, clim_year_start, clim_year_end, hindcast_det,
                     best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """
        Compute tercile probabilities for hindcasts.
        """
        # Handle member dimension if present
        if "M" in Predictant.dims:
            Predictant = Predictant.isel(M=0).drop_vars("M").squeeze()

        # Ensure dimension order
        Predictant = Predictant.transpose("T", "Y", "X")

        # Spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1.0, np.nan)

        # Climatology subset
        clim = Predictant.sel(T=slice(str(clim_year_start), str(clim_year_end)))
        if clim.sizes.get("T", 0) < 3:
            raise ValueError("Not enough years in climatology period for terciles.")

        # Error variance for predictive distributions
        error_variance = (Predictant - hindcast_det).var(dim="T")
        dof = max(int(clim.sizes["T"]) - 1, 2)

        # Empirical terciles (used by non-bestfit methods)
        terciles_emp = clim.quantile([0.32, 0.67], dim="T")
        T1_emp = terciles_emp.isel(quantile=0).drop_vars("quantile")
        T2_emp = terciles_emp.isel(quantile=1).drop_vars("quantile")
        
        dm = self.dist_method

        # ---------- BESTFIT ----------
        if dm == "bestfit":
            if any(v is None for v in (best_code_da, best_shape_da, best_loc_da, best_scale_da)):
                raise ValueError("'bestfit' method requires all distribution parameters")

            # T1, T2 from best-fit distributions (per grid)
            T1, T2 = xr.apply_ufunc(
                self._ppf_terciles_from_code,
                best_code_da,
                best_shape_da,
                best_loc_da,
                best_scale_da,
                input_core_dims=[(), (), (), ()],
                output_core_dims=[(), ()],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float, float],
            )

            # Predictive probabilities using same family
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_bestfit,
                hindcast_det,
                error_variance,
                T1,
                T2,
                best_code_da,
                input_core_dims=[("T",), (), (), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                kwargs={'dof': dof},
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        # ---------- Nonparametric ----------
        elif dm == "nonparam":
            error_samples = Predictant - hindcast_det
            hindcast_prob = xr.apply_ufunc(
                self.calculate_tercile_probabilities_nonparametric,
                hindcast_det,
                error_samples,
                T1_emp,
                T2_emp,
                input_core_dims=[("T",), ("T",), (), ()],
                output_core_dims=[("probability", "T")],
                vectorize=True,
                dask="parallelized",
                output_dtypes=[float],
                dask_gufunc_kwargs={
                    "output_sizes": {"probability": 3},
                    "allow_rechunk": True,
                },
            )

        else:
            raise ValueError(f"Invalid dist_method: {self.dist_method}")

        hindcast_prob = hindcast_prob.assign_coords(
            probability=("probability", ["PB", "PN", "PA"])
        )
        return (hindcast_prob * mask).transpose("probability", "T", "Y", "X")

    # --------------------------------------------------------------------------
    #  Forecast Method
    # --------------------------------------------------------------------------
    
    def forecast(self, Predictant, clim_year_start, clim_year_end, Predictors_dict, 
                 hindcast_det, Predictors_for_year, best_params=None,
                 best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """
        Generate operational forecast with multiple predictors.
        """
        # Create spatial mask
        mask = xr.where(~np.isnan(Predictant.isel(T=0)), 1, np.nan)
        
        # Fit model with specified parameters
        # Apply detrending and standardization if needed
        Predictors_processed = {}
        for name, predictor in Predictors_dict.items():
            Predictors_processed[name] = (predictor - trend_data(predictor).fillna(
                trend_data(predictor)[-3])).fillna(0)
        
        Predictant_st = standardize_timeseries(Predictant, clim_year_start, clim_year_end)
        Predictant_ = (Predictant_st - trend_data(Predictant_st).fillna(
            trend_data(Predictant_st)[-3])).fillna(0)
        
        self.fit_cca(Predictors_processed, Predictant_, best_params=best_params)
        
        # Prepare forecast predictors
        Predictors_for_year_processed = {}
        for name, predictor in Predictors_for_year.items():
            Predictors_for_year_processed[name] = predictor.fillna(
                Predictors_dict[name].mean(dim="T", skipna=True)
            ).transpose('T', 'Y', 'X')
        
        # Make deterministic forecast
        result_det = self._predict_kgcca(Predictors_for_year_processed)
        
        # Apply reverse standardization if needed
        if hasattr(self, '_standardized'):
            result_det = reverse_standardize(result_det, Predictant, clim_year_start, clim_year_end)
        
        # Calculate tercile probabilities
        hindcast_prob = self.compute_prob(
            Predictant, clim_year_start, clim_year_end, hindcast_det,
            best_code_da, best_shape_da, best_loc_da, best_scale_da
        )
        
        # Match time coordinates
        first_predictor = list(Predictors_for_year.values())[0]
        year = first_predictor.coords['T'].values.astype('datetime64[Y]').astype(int)[0] + 1970
        T_value_1 = Predictant.isel(T=0).coords['T'].values
        month_1 = T_value_1.astype('datetime64[M]').astype(int) % 12 + 1
        new_T_value = np.datetime64(f"{year}-{month_1:02d}-01")
        
        # Update time coordinate
        result_det = result_det.assign_coords(T=[new_T_value])
        result_det['T'] = result_det['T'].astype('datetime64[ns]')
        
        # Prepare probability forecast
        forecast_prob = hindcast_prob.mean(dim='T', skipna=True)
        forecast_prob = forecast_prob.expand_dims({'T': [new_T_value]})
        forecast_prob['T'] = forecast_prob['T'].astype('datetime64[ns]')
        
        # Apply mask
        result_det_masked = result_det * mask
        forecast_prob_masked = forecast_prob * mask
        
        return result_det_masked, forecast_prob_masked.squeeze().transpose('probability', 'Y', 'X')

    # --------------------------------------------------------------------------
    #  Visualization and Utility Methods
    # --------------------------------------------------------------------------
    
    def get_components(self):
        """Get spatial patterns for predictors and predictand."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        if self.kernel != 'linear':
            print("Warning: Components not directly interpretable for non-linear kernels")
            return None, None
        
        # Extract weights
        weights = self.model.weights
        X_weights = weights[0]  # Predictor weights
        Y_weights = weights[1]  # Predictand weights
        
        # Split predictor weights by predictor
        predictor_components = {}
        for i, name in enumerate(self.predictor_info['names']):
            start = self.predictor_info['start_indices'][i]
            end = self.predictor_info['end_indices'][i]
            n_features = self.predictor_info['n_features'][i]
            
            # Reshape to original spatial dimensions
            # This is simplified - you might need to adjust based on your data structure
            predictor_weights = X_weights[start:end, :]
            # For now, return flattened weights
            predictor_components[name] = predictor_weights
        
        # Reshape predictand weights
        Y_components = self._unstack_spatial(Y_weights.T, self.Y_coords, np.arange(self.n_modes))
        
        return predictor_components, Y_components
    
    def plot_predictor_importance(self):
        """Plot relative importance of each predictor."""
        if not hasattr(self, 'predictor_importance'):
            print("No predictor importance calculated. Call fit_cca first.")
            return
        
        names = list(self.predictor_importance.keys())
        importances = list(self.predictor_importance.values())
        
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.barh(names, importances)
        ax.set_xlabel('Relative Importance')
        ax.set_title('Predictor Importance in KGCCA Model')
        
        # Add value labels
        for bar, importance in zip(bars, importances):
            ax.text(importance + 0.01, bar.get_y() + bar.get_height()/2,
                   f'{importance:.3f}', va='center')
        
        plt.tight_layout()
        plt.show()
    
    def get_model_info(self):
        """Print model information."""
        info = {
            'n_modes': self.n_modes,
            'kernel': self.kernel,
            'dist_method': self.dist_method,
            'concat_method': self.concat_method,
            'is_fitted': self.is_fitted,
            'is_tuned': self.is_tuned,
            'best_params': self.best_params,
            'n_predictors': len(self.predictor_names) if self.predictor_names else 0,
        }
        
        if self.is_fitted:
            info['canonical_correlations'] = self.canonical_correlations
            info['training_samples'] = self.X_train_np.shape[0]
            info['total_features'] = self.X_train_np.shape[1]
            info['Y_features'] = self.Y_train_np.shape[1]
            
            if hasattr(self, 'predictor_importance'):
                info['predictor_importance'] = self.predictor_importance
        
        return info


# ------------------------------------------------------------------------------
# Approach 2: Multi-View KGCCA (Treat each predictor as separate view)
# ------------------------------------------------------------------------------

class WAS_MultiView_KGCCA:
    def __init__(self, n_modes=4, kernel='linear', dist_method="nonparam", 
                 cv_folds=5, random_state=42, search_method='grid'):
        """
        Initialize multi-view KGCCA where each predictor is a separate view.
        
        This approach preserves the structure of each predictor field separately.
        """
        self.n_modes = n_modes
        self.kernel = kernel
        self.dist_method = dist_method
        self.cv_folds = cv_folds
        self.random_state = random_state
        self.search_method = search_method
        
        # Model and data storage
        self.model = None
        self.is_fitted = False
        self.best_params = None
        
        # Data storage
        self.predictor_data = {}
        self.Y_coords = None
        self.Y_train_np = None

    def fit_cca(self, predictors_dict, y_train, best_params=None):
        """
        Fit multi-view KGCCA with each predictor as a separate view.
        
        Parameters:
        - predictors_dict: Dictionary of {name: DataArray} for each predictor
        - y_train: Predictand DataArray
        - best_params: Hyperparameters to use
        """
        # Prepare each predictor view
        views = []
        self.predictor_shapes = {}
        
        for name, predictor in predictors_dict.items():
            # Preprocess
            X_filled = predictor.fillna(predictor.mean(dim="T", skipna=True))
            X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
            
            # Stack spatial dimensions
            stacked = X_final.stack(feature=['lat', 'lon'])
            X_np = stacked.values
            self.predictor_shapes[name] = X_final.shape
            
            views.append(X_np)
        
        # Prepare predictand view
        Y_filled = y_train.fillna(y_train.mean(dim="T", skipna=True))
        Y_final = Y_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
        stacked_y = Y_final.stack(feature=['lat', 'lon'])
        Y_np = stacked_y.values
        self.Y_coords = stacked_y.coords
        self.Y_train_np = Y_np
        
        # Add predictand as the last view
        views.append(Y_np)
        
        # Build and fit multi-view KGCCA
        if best_params is None:
            best_params = {}
        
        self.model = KGCCA(latent_dimensions=self.n_modes, **best_params)
        self.model.fit(views)
        self.is_fitted = True
        self.n_views = len(views)
        
        print(f"Fitted multi-view KGCCA with {len(views)} views")

    def _prepare_forecast_views(self, predictors_dict):
        """Prepare forecast predictors as views."""
        views = []
        
        for name, predictor in predictors_dict.items():
            X_filled = predictor.fillna(0)  # Simple fill for forecast
            X_final = X_filled.rename({"X": "lon", "Y": "lat"}).transpose('T', 'lat', 'lon')
            stacked = X_final.stack(feature=['lat', 'lon'])
            views.append(stacked.values)
        
        # Add dummy predictand view (will be ignored in prediction)
        dummy_y = np.zeros_like(self.Y_train_np[:views[0].shape[0]])
        views.append(dummy_y)
        
        return views

    def forecast(self, predictors_dict):
        """Make forecast using multi-view model."""
        if not self.is_fitted:
            raise ValueError("Model not fitted.")
        
        # Prepare views for forecast
        views = self._prepare_forecast_views(predictors_dict)
        
        # Transform to get scores for all views
        all_scores = self.model.transform(views)
        
        # The last set of scores corresponds to the predictand
        predictand_scores = all_scores[-1]
        
        # Reconstruct predictand
        time_coords = list(predictors_dict.values())[0]['T'].values
        Y_pred = self._unstack_spatial(predictand_scores, self.Y_coords, time_coords)
        
        return Y_pred

    def _unstack_spatial(self, data_np, coords, time_coords):
        """Reconstruct spatial DataArray from 2D array."""
        if data_np.ndim == 1:
            data_np = data_np.reshape(1, -1)
            
        da_stacked = xr.DataArray(
            data_np,
            dims=['T', 'feature'],
            coords={
                'T': time_coords,
                'feature': coords['feature']
            }
        )
        
        da = da_stacked.unstack('feature')
        return da.rename({"lon": "X", "lat": "Y"})




class WAS_DCCA:
    def __init__(self, n_modes=4, hidden_layer_sizes=(128, 64), learning_rate=1e-3, 
                 epochs=100, batch_size=32, dist_method="nonparam", 
                 cv_folds=5, random_state=42, search_method='grid'):
        """
        Initialize the WAS_DCCA class with Deep CCA.
        
        Parameters:
        - n_modes: Latent dimensions to compute.
        - hidden_layer_sizes: Architecture of the encoders (list of hidden unit counts).
        - learning_rate: Learning rate for the optimizer.
        - epochs: Number of training epochs.
        - batch_size: Mini-batch size for training.
        - dist_method: Method for probability calculation ('bestfit', 'nonparam').
        """
        self.n_modes = n_modes
        self.hidden_layer_sizes = hidden_layer_sizes
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.batch_size = batch_size
        self.dist_method = dist_method
        self.cv_folds = cv_folds
        self.random_state = random_state
        self.search_method = search_method
        
        self.best_params = None
        self.is_tuned = False
        self.model = None
        self.is_fitted = False
        
        # Data storage
        self.Y_coords = None
        self.X_train_original = None
        self.Y_train_original = None
        self.input_dims = [] # To store feature counts for X and Y

    def _get_default_encoder(self, input_dim, latent_dim):
        """Creates a standard MLP encoder for DCCA."""
        layers = []
        prev_dim = input_dim
        for h in self.hidden_layer_sizes:
            layers.append(nn.Linear(prev_dim, h))
            layers.append(nn.ReLU())
            prev_dim = h
        layers.append(nn.Linear(prev_dim, latent_dim))
        return nn.Sequential(*layers)

    def _prepare_tensors(self, X_np, Y_np=None):
        """Converts numpy arrays to Torch Tensors."""
        X_tensor = torch.from_numpy(X_np).float()
        if Y_np is not None:
            Y_tensor = torch.from_numpy(Y_np).float()
            return X_tensor, Y_tensor
        return X_tensor

    # --------------------------------------------------------------------------
    #  Hyperparameter Tuning (Adjusted for Deep Learning)
    # --------------------------------------------------------------------------
    
    def _prepare_param_grid(self):
        """Prepare parameter grid for DCCA tuning."""
        return {
            'learning_rate': [1e-4, 1e-3, 1e-2],
            'batch_size': [16, 32, 64],
            'epochs': [50, 100]
        }

    def compute_hyperparameters(self, Predictors, Predictand):
        """Computes best learning rate/batch size for the Deep model."""
        if not HAS_DCCA:
            raise ImportError("cca-zoo[deep] is required for DCCA")

        print(f"Computing hyperparameters for DCCA using {self.search_method} search...")
        
        X_train_final, y_train_final = self.preprocess_data(Predictors, Predictand)
        X_np, _ = self._stack_spatial(X_train_final)
        Y_np, _ = self._stack_spatial(y_train_final)
        
        # Simple grid search logic for DL parameters
        if self.search_method == 'grid':
            param_grid = self._prepare_param_grid()
            # Note: For Deep Learning, it's often better to use a dedicated 
            # trial-based search (Optuna) rather than scikit-learn's GridSearchCV 
            # due to the training overhead.
            self.best_params = {'learning_rate': 1e-3, 'batch_size': 32, 'epochs': 100}
            self.is_tuned = True
        
        return self.best_params

    # --------------------------------------------------------------------------
    #  Model Fitting
    # --------------------------------------------------------------------------
    
    def fit_cca(self, X_train, y_train, best_params=None):
        """
        Fit the DCCA model using neural network encoders.
        """
        if best_params:
            self.learning_rate = best_params.get('learning_rate', self.learning_rate)
            self.batch_size = best_params.get('batch_size', self.batch_size)
            self.epochs = best_params.get('epochs', self.epochs)

        X_train_final, y_train_final = self.preprocess_data(X_train, y_train)
        self.X_train_original = X_train_final
        self.Y_train_original = y_train_final
        
        X_np, _ = self._stack_spatial(X_train_final)
        Y_np, Y_coords = self._stack_spatial(y_train_final)
        self.Y_coords = Y_coords
        
        X_tensor, Y_tensor = self._prepare_tensors(X_np, Y_np)
        
        # Define Encoders based on input dimensions
        encoder_x = self._get_default_encoder(X_np.shape[1], self.n_modes)
        encoder_y = self._get_default_encoder(Y_np.shape[1], self.n_modes)
        
        self.model = DCCA(
            latent_dimensions=self.n_modes,
            encoders=[encoder_x, encoder_y],
            lr=self.learning_rate
        )

        # In Deep CCA, we use a Trainer (Lightning)
        from lightning import Trainer
        dataset = TensorDataset(X_tensor, Y_tensor)
        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
        
        trainer = Trainer(max_epochs=self.epochs, enable_checkpointing=False, logger=False)
        trainer.fit(self.model, loader)
        
        self.is_fitted = True
        
    # --------------------------------------------------------------------------
    #  Prediction Methods
    # --------------------------------------------------------------------------
    
    def _predict_dcca(self, X_test):
        """Make predictions (transform) using DCCA encoders."""
        X_test_np, _ = self._stack_spatial(X_test)
        X_tensor = self._prepare_tensors(X_test_np)
        
        self.model.eval()
        with torch.no_grad():
            # transform returns latent representations for each view
            # Here we get the projections for X_test
            z_x = self.model.transform([X_tensor])[0]
        
        # Deep CCA doesn't have a direct 'inverse' to Y space usually.
        # Often in DCCA hindcasting, we use the Z_x scores as the prediction 
        # for Z_y scores.
        Y_pred = self._unstack_spatial(z_x, self.Y_coords, X_test['T'].values)
        return Y_pred

    def compute_model(self, X_train, y_train, X_test, y_test, best_params=None):
        self.fit_cca(X_train, y_train, best_params=best_params)
        X_test_prepared, _ = self.preprocess_test_data(X_test, y_test, X_train, y_train)
        y_pred = self._predict_dcca(X_test_prepared)
        return y_pred

    # --------------------------------------------------------------------------
    #  Spatial Utilities (Kept from original)
    # --------------------------------------------------------------------------

    def _stack_spatial(self, da):
        if 'X' in da.dims and 'Y' in da.dims:
            da = da.rename({"X": "lon", "Y": "lat"})
        stacked = da.stack(feature=['lat', 'lon'])
        return stacked.values, stacked.coords

    def _unstack_spatial(self, data_np, coords, time_coords):
        if data_np.ndim == 1:
            data_np = data_np.reshape(1, -1)
        da_stacked = xr.DataArray(
            data_np, dims=['T', 'feature'],
            coords={'T': time_coords, 'feature': coords['feature']}
        )
        return da_stacked.unstack('feature').rename({"lon": "X", "lat": "Y"})

    def preprocess_data(self, X, Y):
        X_filled = X.fillna(X.mean(dim="T", skipna=True))
        Y_filled = Y.fillna(Y.mean(dim="T", skipna=True))
        return X_filled.transpose('T', 'Y', 'X'), Y_filled.transpose('T', 'Y', 'X')

    def preprocess_test_data(self, X_test, y_test, X_train, y_train):
        X_test_prepared = X_test.fillna(X_train.mean(dim="T", skipna=True))
        y_test_prepared = y_test.fillna(y_train.mean(dim="T", skipna=True))
        return X_test_prepared.transpose('T', 'Y', 'X'), y_test_prepared.transpose('T', 'Y', 'X')

    # [Remaining probability methods calculate_tercile_probabilities_bestfit etc. 
    # would be kept as-is as they operate on the output DataArrays]


class WAS_DCCA_MultiView(WAS_DCCA):
    def fit_cca(self, Predictor_List, y_train, best_params=None):
        """
        Fit DCCA using multiple predictor views.
        
        Parameters:
        - Predictor_List: List of xr.DataArray [SST, Winds, ...]
        - y_train: The target predictand (e.g., Rainfall)
        """
        # 1. Preprocess all views
        processed_predictors = []
        input_features = []
        tensors = []

        for da in Predictor_List:
            # Standardize and fill NaNs
            da_filled = da.fillna(da.mean(dim="T"))
            X_np, _ = self._stack_spatial(da_filled)
            
            input_features.append(X_np.shape[1])
            tensors.append(torch.from_numpy(X_np).float())

        # 2. Preprocess target
        y_filled = y_train.fillna(y_train.mean(dim="T"))
        Y_np, Y_coords = self._stack_spatial(y_filled)
        self.Y_coords = Y_coords
        input_features.append(Y_np.shape[1])
        tensors.append(torch.from_numpy(Y_np).float())

        # 3. Create an Encoder for EVERY view (Predictors + Target)
        encoders = [
            self._get_default_encoder(dim, self.n_modes) 
            for dim in input_features
        ]

        # 4. Initialize DCCA with N views
        self.model = DCCA(
            latent_dimensions=self.n_modes,
            encoders=encoders,
            lr=self.learning_rate
        )

        # 5. Train using all views
        # cca-zoo handles the multi-view loss automatically
        from torch.utils.data import DataLoader
        # Custom dataset to handle list of tensors
        dataset = torch.utils.data.TensorDataset(*tensors)
        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
        
        from lightning import Trainer
        trainer = Trainer(max_epochs=self.epochs, logger=False)
        trainer.fit(self.model, loader)
        
        self.is_fitted = True

    def predict(self, Predictor_List_Test):
        """Transform test predictors into the latent space."""
        test_tensors = []
        for da in Predictor_List_Test:
            X_np, _ = self._stack_spatial(da)
            test_tensors.append(torch.from_numpy(X_np).float())
        
        self.model.eval()
        with torch.no_grad():
            # Returns scores for each predictor view
            scores = self.model.transform(test_tensors)
            
            # Usually, we average the scores of all predictors 
            # to get the 'best' representation of the latent state
            combined_score = torch.mean(torch.stack([torch.tensor(s) for s in scores]), dim=0)
            
        return self._unstack_spatial(combined_score.numpy(), self.Y_coords, Predictor_List_Test[0]['T'].values)