# Core Python Libraries
import os
from pathlib import Path
import calendar
from datetime import datetime
from dateutil.relativedelta import relativedelta
from collections import defaultdict

# Numerical and Data Manipulation Libraries
import numpy as np
import pandas as pd
import xarray as xr
import dask.array as da

# Visualization Libraries
import matplotlib.pyplot as plt
import seaborn as sns
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap, BoundaryNorm

# Climate Data API
import cdsapi

# Machine Learning and Statistical Analysis Libraries
from minisom import MiniSom
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.cluster import KMeans, AgglomerativeClustering
from scipy.stats import pearsonr, gamma, lognorm, stats
import scipy.signal as sig

# EOF Analysis and Verification Libraries
import xeofs as xe
import xskillscore as xs
from wass2s.utils import *

# Core Python Libraries
import os
from pathlib import Path
import calendar
from datetime import datetime
from dateutil.relativedelta import relativedelta
from collections import defaultdict

# Numerical and Data Manipulation Libraries
import numpy as np
import pandas as pd
import xarray as xr
import dask.array as da

# Visualization Libraries
import matplotlib.pyplot as plt
import seaborn as sns
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap, BoundaryNorm

# Climate Data API
import cdsapi

# Machine Learning and Statistical Analysis Libraries
from minisom import MiniSom
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.cluster import KMeans, AgglomerativeClustering
from scipy.stats import pearsonr, gamma, lognorm, stats
import scipy.signal as sig

# EOF Analysis and Verification Libraries
import xeofs as xe
import xskillscore as xs
from wass2s.utils import *
from wass2s.was_bias_correction import *

class WAS_Analog:
    """Analog-based forecasting toolkit for seasonal climate applications.

    This class orchestrates the end-to-end workflow required to build analog ensembles
    for Sea-Surface Temperature (SST) predictors and to translate them into deterministic
    and probabilistic rainfall forecasts over West Africa or any user-defined domain.
    Supports five analog-selection strategies: Self-Organizing Maps (SOM),
    correlation ranking, Principal-Component (EOF) similarity, K-Means clustering,
    and Agglomerative clustering.

    Parameters
    ----------
    dir_to_save : str
        Directory path to save downloaded and processed data files.
    year_start : int
        Starting year for historical data.
    year_forecast : int
        Target forecast year.
    predictor_vars : list of dict, optional
        List of dictionaries specifying predictor variables, each containing
        'reanalysis_name', 'model_name', 'variable', and 'area'.
        Default is a list with NOAA SST and ERA5 SP, VGRD_850 variables.
    method_analog : str, optional
        Analog method to use ("som", "cor_based", "bias_based, "pca_based", "kmeans_based", "agglomerative_based"). Default is "som".
    month_of_initialization : int, optional
        Month of initialization for forecasts (1-12). If None, uses previous month.
    lead_time : list, optional
        List of lead times in months. If None, defaults to [1, 2, 3, 4, 5].
    ensemble_mean : str, optional
        Method for ensemble mean ("mean" or "median"). Default is "mean".
    rolling : int, optional
        Window size for rolling mean. Default is 3.
    standardize : bool, optional
        If True, standardize data; otherwise, compute anomalies. Default is True.
    multivariateEOF : bool, optional
        If True, use multivariate EOF analysis. Default is False.
    eof_explained_var : float, optional
        Fraction of variance to explain with EOF modes. Default is 0.95.
    clim_year_start : int, optional
        Start year for climatology period.
    clim_year_end : int, optional
        End year for climatology period.
    index_compute : list, optional
        List of climate indices to compute (e.g., ['NINO34', 'DMI']).
    some_grid_size : tuple, optional
        Grid size for SOM (rows, columns). Default is (None, None) for automatic sizing.
    some_learning_rate : float, optional
        Learning rate for SOM training. Default is 0.5.
    radius : float, optional
        Neighborhood radius for SOM analog search. Default is 1.0.
    some_neighborhood_function : str, optional
        Neighborhood function for SOM ("gaussian", "mexican_hat"). Default is "gaussian".
    some_sigma : float, optional
        Initial neighborhood radius for SOM. Default is 1.0.
    some_num_iteration : int, optional
        Number of iterations for SOM training. Default is 2000.
    dist_method : str, optional
        Method for probability calculation ("gamma", "t", "normal", "lognormal", "nonparam").
        Default is "gamma".
    n_clusters : int, optional
        Number of clusters for K-Means and Agglomerative clustering. Default is 4.
    affinity : str, optional
        Distance metric for Agglomerative clustering (e.g., 'euclidean', 'manhattan'). Default is 'euclidean'.
    linkage : str, optional
        Linkage criterion for Agglomerative clustering (e.g., 'ward', 'complete', 'average'). Default is 'ward'.

    Notes
    -----
    - Performance Considerations: For large datasets (e.g., many years or high-dimensional data), the AgglomerativeClustering method (`agglomerative_based`) may be significantly slower than K-Means (`kmeans_based`) or SOM (`som`) due to its computational complexity. Users should consider testing with smaller datasets or reducing the number of clusters when using AgglomerativeClustering.
    """

    def __init__(self, dir_to_save, year_start, year_forecast,
                 predictor_vars=[{'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'SST', 'area': [60, -180, -60, 180]},
                                {'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'SP', 'area': [60, -180, -60, 180]},
                                {'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'VGRD_850', 'area': [60, -180, -60, 180]}],
                 method_analog="som", month_of_initialization=None,
                 lead_time=None, ensemble_mean="mean", rolling=3, standardize=True, detrend=False, multivariateEOF=False,
                 eof_explained_var=0.95, clim_year_start=None, clim_year_end=None,
                 index_compute=None, some_grid_size=(None, None), some_learning_rate=0.5, radius=1.0,
                 some_neighborhood_function='gaussian', some_sigma=1.0, some_num_iteration=2000, dist_method="gamma",
                 n_clusters=4, affinity='euclidean', linkage='ward'):
        
        self.dir_to_save = dir_to_save
        self.year_start = year_start
        self.year_forecast = year_forecast
        self.predictor_vars = predictor_vars
        self.method_analog = method_analog
        self.month_of_initialization = month_of_initialization
        self.lead_time = lead_time
        self.ensemble_mean = ensemble_mean
        self.eof_explained_var = eof_explained_var
        self.multivariateEOF = multivariateEOF
        self.rolling = rolling
        self.standardize = standardize
        self.detrend = False
        self.clim_year_start = clim_year_start
        self.clim_year_end = clim_year_end
        self.index_compute = index_compute
        self.some_grid_size = some_grid_size
        self.some_learning_rate = some_learning_rate
        self.radius = radius
        self.some_neighborhood_function = some_neighborhood_function
        self.some_sigma = some_sigma
        self.some_num_iteration = some_num_iteration
        self.dist_method = dist_method
        self.n_clusters = n_clusters
        self.affinity = affinity
        self.linkage = linkage

    def calc_index(self, indices, sst):
        """Calculate climate indices from SST data."""
        sst_indices_name = {
            "NINO34": ("Nino3.4", -170, -120, -5, 5),
            "NINO12": ("Niño1+2", -90, -80, -10, 0),
            "NINO3": ("Nino3", -150, -90, -5, 5),
            "NINO4": ("Nino4", -150, 160, -5, 5),
            "NINO_Global": ("ALL NINO Zone", -80, 160, -10, 5),
            "TNA": ("Tropical Northern Atlantic Index", -55, -15, 5, 25),
            "TSA": ("Tropical Southern Atlantic Index", -30, 10, -20, 0),
            "NAT": ("North Atlantic Tropical", -40, -20, 5, 20),
            "SAT": ("South Atlantic Tropical", -15, 5, -20, 5),
            "TASI": ("NAT-SAT", None, None, None, None),
            "WTIO": ("Western Tropical Indian Ocean (WTIO)", 50, 70, -10, 10),
            "SETIO": ("Southeastern Tropical Indian Ocean (SETIO)", 90, 110, -10, 0),
            "DMI": ("WTIO - SETIO", None, None, None, None),
            "MB": ("Mediterranean Basin", 0, 50, 30, 42),
            "M1": ("M1", -50, 5, -50, -25),
            "M2": ("M2", -75, -10, 25, 50),
            "M3": ("M3", -175, -125, 25, 50),
            "M4": ("M4", -175, -125, -50, -25),
        }
        
        predictor = {}
        for idx in sst_indices_name:
            if idx in ["TASI", "DMI"]:
                continue
            _, lon_min, lon_max, lat_min, lat_max = sst_indices_name[idx]
            sst_region = sst.sel(X=slice(lon_min, lon_max), Y=slice(lat_min, lat_max)).mean(dim=["X", "Y"], skipna=True)
            predictor[idx] = sst_region

        predictor["TASI"] = predictor["NAT"] - predictor["SAT"]
        predictor["DMI"] = predictor["WTIO"] - predictor["SETIO"]
        
        selected_indices = {i: predictor[i] for i in indices if i in predictor}
        data_vars = {key: ds.rename(key) for key, ds in selected_indices.items()}
        return xr.Dataset(data_vars)

    def _postprocess_ersst(self, ds, var_name):
        """Post-process ERSST dataset."""
        ds = ds.drop_vars('zlev', errors='ignore').squeeze()
        keep_vars = [var_name, 'T', 'X', 'Y']
        drop_vars = [v for v in ds.variables if v not in keep_vars]
        return ds.drop_vars(drop_vars, errors="ignore")

    # def download_reanalysis_(self, force_download=False):
    #         """
    #         Download reanalysis data (ERA5 and NOAA ERSST).
    #         Saves raw MONTHLY data (1-12) without aggregation, preserving seasonality.
    #         """
    #         year_end = self.year_forecast
            
    #         # We download ALL months to preserve the full seasonal cycle for the analog search
    #         months_int = list(range(1, 13))
    #         # Naming convention: usually we just denote this as the full year range
    #         # If you want to put "JanDec" in the filename, you can, or just leave it generic.
    #         season_str = "JanDec" 
            
    #         variables = [item['variable'] for item in self.predictor_vars]
    #         centers = [item['reanalysis_name'] for item in self.predictor_vars]
    #         areas = [item['area'] for item in self.predictor_vars]
            
    #         # Variable Mappings
    #         variables_1 = {
    #             "PRCP": "total_precipitation", "TEMP": "2m_temperature",
    #             "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
    #             "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
    #             "UGRD10": "10m_u_component_of_wind", "VGRD10": "10m_v_component_of_wind",
    #             "SST": "sea_surface_temperature", "SLP": "mean_sea_level_pressure",
    #             "DSWR": "surface_solar_radiation_downwards", "DLWR": "surface_thermal_radiation_downwards",
    #             "NOLR": "top_net_thermal_radiation",
    #         }
    #         variables_2 = {
    #             "HUSS_1000": "specific_humidity", "HUSS_925": "specific_humidity", "HUSS_850": "specific_humidity",
    #             "UGRD_1000": "u_component_of_wind", "UGRD_925": "u_component_of_wind", "UGRD_850": "u_component_of_wind",
    #             "VGRD_1000": "v_component_of_wind", "VGRD_925": "v_component_of_wind", "VGRD_850": "v_component_of_wind",
    #         }
    
    #         dir_to_save = Path(self.dir_to_save)
    #         dir_to_save.mkdir(parents=True, exist_ok=True)
    #         store_file_path = {}
    
    #         for center, var, area in zip(centers, variables, areas):
    #             # Filename includes "JanDec" or similar to indicate full monthly data
    #             combined_output_path = dir_to_save / f"{center}_{var}_{self.year_start}_{year_end}_{season_str}_{area[0]}{area[1]}{area[2]}{area[3]}.nc" 
    
    #             if not force_download and combined_output_path.exists():
    #                 print(f"{combined_output_path} exists. Skipping download.")
    #                 store_file_path[var] = xr.open_dataset(combined_output_path)
    #                 continue
    
    #             # =================================================================
    #             # CASE 1: NOAA ERSST (Direct NCEI Download)
    #             # =================================================================
    #             if f"{center}.{var}" == "NOAA.SST":
    #                 try:
    #                     print(f"Starting NCEI Direct Download for NOAA ERSST: {self.year_start}-{year_end}")
    #                     cache_dir = dir_to_save / "ersst_cache"
    #                     cache_dir.mkdir(exist_ok=True)
                        
    #                     downloaded_paths = []
    #                     # v5 is standard, check if v6 is required/available
    #                     base_url = "https://www.ncei.noaa.gov/data/sea-surface-temperature-extended-reconstructed/v6/access/" 
    
    #                     for s_year in range(self.year_start, year_end + 1):
    #                         for month in months_int:
    #                             fname = f"ersst.v6.{s_year}{month:02d}.nc"
    #                             local_path = cache_dir / fname
                                
    #                             if not local_path.exists():
    #                                 url = f"{base_url}/{fname}"
    #                                 # print(f"Downloading {fname}...") # reduced verbosity
    #                                 r = requests.get(url, timeout=30)
    #                                 r.raise_for_status()
    #                                 with open(local_path, 'wb') as f:
    #                                     f.write(r.content)
    #                             downloaded_paths.append(local_path)
    
    #                     # Process ERSST
    #                     ds = xr.open_mfdataset(downloaded_paths, combine='by_coords')
    #                     ds = ds[["sst"]].drop_vars("lev", errors="ignore").squeeze()
    
    #                     # Standardize Coordinates (0-360 -> -180/180)
    #                     if 'lon' in ds.coords:
    #                         ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
    #                         ds = ds.sortby('lon')
                        
    #                     # Rename to standard (X, Y, T, SST)
    #                     rename_dict = {}
    #                     if 'lat' in ds.coords: rename_dict['lat'] = 'Y'
    #                     if 'lon' in ds.coords: rename_dict['lon'] = 'X'
    #                     if 'time' in ds.coords: rename_dict['time'] = 'T'
    #                     if 'sst' in ds.variables: rename_dict['sst'] = 'SST'
    #                     ds = ds.rename(rename_dict)
                        
    #                     # Slice Area [N, W, S, E]
    #                     ds = ds.sortby(['X', 'Y'])
    #                     ds = ds.sel(X=slice(area[1], area[3]), Y=slice(area[2], area[0]))
    
    #                     # NO AGGREGATION: Just ensure time is sorted
    #                     ds = ds.sortby('T')
                        
    #                     # Final Cleanup
    #                     ds = ds.rename({"SST": "sst"}) 
    #                     ds.to_netcdf(combined_output_path)
    #                     store_file_path[var] = ds
    #                     print(f"Saved NOAA ERSST data to {combined_output_path}")
    #                     continue
    #                 except Exception as e:
    #                     print(f"Failed to download NOAA.SST via NCEI: {e}")
    #                     import traceback
    #                     traceback.print_exc()
    #                     continue
    
    #             # =================================================================
    #             # CASE 2: ERA5 / CDSAPI
    #             # =================================================================
    #             combined_datasets = []
    #             client = cdsapi.Client()
                
    #             is_pressure = var in variables_2
    #             dataset_name = "reanalysis-era5-pressure-levels-monthly-means" if is_pressure else "reanalysis-era5-single-levels-monthly-means"
    #             api_var = variables_2[var] if is_pressure else variables_1.get(var, var)
    #             press_level = var.split("_")[1] if is_pressure else None
    
    #             for year in range(self.year_start, year_end + 1):
    #                 yearly_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                
    #                 if not force_download and yearly_file_path.exists():
    #                     try:
    #                         ds = xr.open_dataset(yearly_file_path).load()
    #                         if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
    #                         combined_datasets.append(ds)
    #                         continue
    #                     except Exception as e:
    #                         print(f"Failed to load {yearly_file_path}: {e}")
                
    #                 try:
    #                     request = {
    #                         "product_type": "monthly_averaged_reanalysis",
    #                         "variable": api_var,
    #                         "year": str(year),
    #                         "month": [f"{m:02d}" for m in months_int],
    #                         "time": "00:00",
    #                         "area": area,
    #                         "format": "netcdf",
    #                     }
    #                     if is_pressure:
    #                         request["pressure_level"] = press_level
    
    #                     print(f"Downloading {var} for {year} from {center}...")
    #                     client.retrieve(dataset_name, request).download(str(yearly_file_path))
                    
    #                     with xr.open_dataset(yearly_file_path) as ds:
    #                         if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
    #                         ds = ds.load()
    #                         combined_datasets.append(ds)
    #                 except Exception as e:
    #                     print(f"Failed to download/process {var} for {year}: {e}")
    #                     continue
            
    #             if combined_datasets:
    #                 print(f"Concatenating {var} datasets...")
    #                 combined_ds = xr.concat(combined_datasets, dim="T")
    #                 combined_ds = combined_ds.drop_vars(["number", "expver"], errors="ignore").squeeze()
                    
    #                 # Unit Conversions
    #                 if var in ["TMIN", "TEMP", "TMAX", "SST"]:
    #                     combined_ds = combined_ds - 273.15
    #                 elif var == "PRCP":
    #                     combined_ds = combined_ds * 1000 
    #                 elif var in ["DSWR", "DLWR", "NOLR"]:
    #                     combined_ds = combined_ds / 86400
    #                 elif var == "SLP":
    #                     combined_ds = combined_ds / 100
                    
    #                 combined_ds = combined_ds.isel(Y=slice(None, None, -1))
                    
    #                 # NO AGGREGATION HERE either
    #                 combined_ds = combined_ds.sortby('T')
    
    #                 store_file_path[var] = combined_ds
    #                 combined_ds.to_netcdf(combined_output_path)
    #                 print(f"Saved combined dataset to {combined_output_path}")
                    
    #                 # Cleanup yearly files
    #                 for year in range(self.year_start, year_end + 1):
    #                     single_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
    #                     if single_file_path.exists():
    #                         try: os.remove(single_file_path)
    #                         except: pass
    #             else:
    #                 print(f"No data combined for {var}. Check download success.")
    
    #         return store_file_path

    
    def download_reanalysis_(self, force_download=False):
            """
            Download reanalysis data (ERA5 and NOAA ERSST).
            Saves raw MONTHLY data (1-12) without aggregation, preserving seasonality.
            """
            year_end = self.year_forecast
            
            # We download ALL months to preserve the full seasonal cycle for the analog search
            months_int = list(range(1, 13))
            # Naming convention for the season string in filename
            season_str = "JanDec" 
            
            variables = [item['variable'] for item in self.predictor_vars]
            centers = [item['reanalysis_name'] for item in self.predictor_vars]
            areas = [item['area'] for item in self.predictor_vars]
            
            # Variable Mappings
            variables_1 = {
                "PRCP": "total_precipitation", "TEMP": "2m_temperature",
                "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
                "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
                "UGRD10": "10m_u_component_of_wind", "VGRD10": "10m_v_component_of_wind",
                "SST": "sea_surface_temperature", "SLP": "mean_sea_level_pressure",
                "DSWR": "surface_solar_radiation_downwards", "DLWR": "surface_thermal_radiation_downwards",
                "NOLR": "top_net_thermal_radiation",
            }
            variables_2 = {
                "HUSS_1000": "specific_humidity", "HUSS_925": "specific_humidity", "HUSS_850": "specific_humidity",
                "UGRD_1000": "u_component_of_wind", "UGRD_925": "u_component_of_wind", "UGRD_850": "u_component_of_wind",
                "VGRD_1000": "v_component_of_wind", "VGRD_925": "v_component_of_wind", "VGRD_850": "v_component_of_wind",
            }
        
            dir_to_save = Path(self.dir_to_save)
            dir_to_save.mkdir(parents=True, exist_ok=True)
            store_file_path = {}
        
            for center, var, area in zip(centers, variables, areas):
                # Filename includes "JanDec" to indicate full monthly data
                combined_output_path = dir_to_save / f"{center}_{var}_{self.year_start}_{year_end}_{season_str}_{area[0]}{area[1]}{area[2]}{area[3]}.nc" 
        
                if not force_download and combined_output_path.exists():
                    print(f"{combined_output_path} exists. Skipping download.")
                    store_file_path[var] = xr.open_dataset(combined_output_path)
                    continue
        
                # =================================================================
                # CASE 1: NOAA ERSST (Direct NCEI Download)
                # =================================================================
                if f"{center}.{var}" == "NOAA.SST":
                    try:
                        print(f"Starting NCEI Direct Download for NOAA ERSST: {self.year_start}-{year_end}")
                        cache_dir = dir_to_save / "ersst_cache"
                        cache_dir.mkdir(exist_ok=True)
                        
                        downloaded_paths = []
                        # Using v5 as standard, change URL to v6 if specifically required
                        base_url = "https://www.ncei.noaa.gov/data/sea-surface-temperature-extended-reconstructed/v6/access/" 
        
                        for s_year in range(self.year_start, year_end + 1):
                            for month in months_int:
                                fname = f"ersst.v6.{s_year}{month:02d}.nc"
                                local_path = cache_dir / fname
                                
                                if not local_path.exists():
                                    url = f"{base_url}/{fname}"
                                    # print(f"Downloading {fname}...") 
                                    r = requests.get(url, timeout=30)
                                    r.raise_for_status()
                                    with open(local_path, 'wb') as f:
                                        f.write(r.content)
                                downloaded_paths.append(local_path)
        
                        # Process ERSST
                        # decode_times=True ensures we interpret the time correctly immediately
                        ds = xr.open_mfdataset(downloaded_paths, combine='by_coords')
                        ds = ds[["sst"]].drop_vars("lev", errors="ignore").squeeze()
        
                        # Standardize Coordinates (0-360 -> -180/180)
                        if 'lon' in ds.coords:
                            ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
                            ds = ds.sortby('lon')
                            
                        # Rename to standard (X, Y, T, SST)
                        rename_dict = {}
                        if 'lat' in ds.coords: rename_dict['lat'] = 'Y'
                        if 'lon' in ds.coords: rename_dict['lon'] = 'X'
                        if 'time' in ds.coords: rename_dict['time'] = 'T'
                        if 'sst' in ds.variables: rename_dict['sst'] = 'SST'
                        ds = ds.rename(rename_dict)
                            
                        # Slice Area [N, W, S, E]
                        # Ensure dimensions are sorted before slicing
                        ds = ds.sortby(['X', 'Y'])
                        # Note: ERSST Lat is usually Ascending (S->N), so slice(S, N) i.e., slice(area[2], area[0])
                        ds = ds.sel(X=slice(area[1], area[3]), Y=slice(area[2], area[0]))
        
                        # NO AGGREGATION: Just ensure time is sorted
                        ds = ds.sortby('T')
                            
                        # Final Cleanup
                        # The class typically expects 'sst' (lowercase) in the final dataset
                        if 'SST' in ds.variables:
                            ds = ds.rename({"SST": "sst"}) 
                            
                        ds.to_netcdf(combined_output_path)
                        store_file_path[var] = ds
                        print(f"Saved NOAA ERSST data to {combined_output_path}")
                        continue
                    except Exception as e:
                        print(f"Failed to download NOAA.SST via NCEI: {e}")
                        import traceback
                        traceback.print_exc()
                        continue
        
                # =================================================================
                # CASE 2: ERA5 / CDSAPI
                # =================================================================
                combined_datasets = []
                client = cdsapi.Client()
                
                is_pressure = var in variables_2
                dataset_name = "reanalysis-era5-pressure-levels-monthly-means" if is_pressure else "reanalysis-era5-single-levels-monthly-means"
                api_var = variables_2[var] if is_pressure else variables_1.get(var, var)
                press_level = var.split("_")[1] if is_pressure else None
        
                for year in range(self.year_start, year_end + 1):
                    yearly_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                    
                    if not force_download and yearly_file_path.exists():
                        try:
                            ds = xr.open_dataset(yearly_file_path).load()
                            if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                            combined_datasets.append(ds)
                            continue
                        except Exception as e:
                            print(f"Failed to load {yearly_file_path}: {e}")
                    
                    try:
                        request = {
                            "product_type": "monthly_averaged_reanalysis",
                            "variable": api_var,
                            "year": str(year),
                            "month": [f"{m:02d}" for m in months_int],
                            "time": "00:00",
                            "area": area,
                            "format": "netcdf",
                        }
                        if is_pressure:
                            request["pressure_level"] = press_level
        
                        print(f"Downloading {var} for {year} from {center}...")
                        client.retrieve(dataset_name, request).download(str(yearly_file_path))
                        
                        with xr.open_dataset(yearly_file_path) as ds:
                            if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                            ds = ds.load()
                            combined_datasets.append(ds)
                    except Exception as e:
                        print(f"Failed to download/process {var} for {year}: {e}")
                        continue
                
                if combined_datasets:
                    print(f"Concatenating {var} datasets...")
                    combined_ds = xr.concat(combined_datasets, dim="T")
                    combined_ds = combined_ds.drop_vars(["number", "expver"], errors="ignore").squeeze()
                    
                    # Unit Conversions
                    if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                        combined_ds = combined_ds - 273.15
                    elif var == "PRCP":
                        combined_ds = combined_ds * 1000 
                    elif var in ["DSWR", "DLWR", "NOLR"]:
                        combined_ds = combined_ds / 86400
                    elif var == "SLP":
                        combined_ds = combined_ds / 100
                    
                    combined_ds = combined_ds.isel(Y=slice(None, None, -1))
                    
                    # NO AGGREGATION HERE either
                    combined_ds = combined_ds.sortby('T')
        
                    store_file_path[var] = combined_ds
                    combined_ds.to_netcdf(combined_output_path)
                    print(f"Saved combined dataset to {combined_output_path}")
                    
                    # Cleanup yearly files
                    for year in range(self.year_start, year_end + 1):
                        single_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                        if single_file_path.exists():
                            try: os.remove(single_file_path)
                            except: pass
                else:
                    print(f"No data combined for {var}. Check download success.")
        
            return store_file_path

    
    def download_reanalysis(self, force_download=False):
        """Download reanalysis data."""
        year_end = self.year_forecast
        variables = [item['variable'] for item in self.predictor_vars]
        centers = [item['reanalysis_name'] for item in self.predictor_vars]
        areas = [item['area'] for item in self.predictor_vars]
        
        variables_1 = {
            "PRCP": "total_precipitation",
            "TEMP": "2m_temperature",
            "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
            "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
            "UGRD10": "10m_u_component_of_wind",
            "VGRD10": "10m_v_component_of_wind",
            "SST": "sea_surface_temperature",
            "SLP": "mean_sea_level_pressure",
            "DSWR": "surface_solar_radiation_downwards",
            "DLWR": "surface_thermal_radiation_downwards",
            "NOLR": "top_net_thermal_radiation",
        }
        variables_2 = {
            "HUSS_1000": "specific_humidity",
            "HUSS_925": "specific_humidity",
            "HUSS_850": "specific_humidity",
            "UGRD_1000": "u_component_of_wind",
            "UGRD_925": "u_component_of_wind",
            "UGRD_850": "u_component_of_wind",
            "VGRD_1000": "v_component_of_wind",
            "VGRD_925": "v_component_of_wind",
            "VGRD_850": "v_component_of_wind",
        }

        dir_to_save = Path(self.dir_to_save)
        os.makedirs(dir_to_save, exist_ok=True)
    
        months = [f"{m:02d}" for m in range(1, 13)]
        season = "".join([calendar.month_abbr[int(m)] for m in months])
        store_file_path = {}

        for center, var, area in zip(centers, variables, areas):
            combined_output_path = dir_to_save / f"{center}_{var}_{self.year_start}_{year_end}_{season}_{area[0]}{area[1]}{area[2]}{area[3]}.nc" 
            if not force_download and combined_output_path.exists():
                print(f"{combined_output_path} exists. Skipping download.")
                store_file_path[var] = xr.open_dataset(combined_output_path)
                continue

            if f"{center}.{var}" == "NOAA.SST":
                
                print(f"Starting download for NOAA ERSST (v6) for {self.year_start}-{year_end}")
                # Create a cache dir for individual monthly files so we don't redownload them unnecessarily
                cache_dir = dir_to_save / "ersst_cache"
                cache_dir.mkdir(exist_ok=True)
                
                # Identify exact list of year-months required for the seasonal request
                # This ensures we handle cross-year (e.g. NDJ) correctly
                required_files = []
                season_months = [int(m) for m in range(1, 13)]
                pivot = season_months[0]
                
                # Loop through the "Seasonal Years"
                for s_year in range(self.year_start, year_end + 1):
                    for month in season_months:
                        # Calculate actual calendar year for this month
                        # If month is >= pivot, it belongs to s_year (e.g. Nov 2000)
                        # If month is < pivot, it belongs to s_year + 1 (e.g. Jan 2001)
                        if month >= pivot:
                            cal_year = s_year
                        else:
                            cal_year = s_year + 1
                        
                        # ERSST v6 filename format: ersst.v6.YYYYMM.nc
                        fname = f"ersst.v6.{cal_year}{month:02d}.nc"
                        required_files.append((cal_year, month, fname))

                # Download Loop
                # Base URL for ERSST 'v6'.
                base_url = "https://www.ncei.noaa.gov/data/sea-surface-temperature-extended-reconstructed/v6/access/"
                
                downloaded_paths = []
                
                try:
                    for (yr, mn, fname) in required_files:
                        local_path = cache_dir / fname
                        if not local_path.exists():
                            # Construct full URL
                            url = f"{base_url}/{fname}"
                            print(f"Downloading {fname}...")
                            r = requests.get(url, timeout=30)
                            r.raise_for_status()
                            with open(local_path, 'wb') as f:
                                f.write(r.content)
                        downloaded_paths.append(local_path)
                    
                    print("All files downloaded. merging and processing...")

                    # Open all files
                    # use_cftime=True is safer for long historical records
                    ds = xr.open_mfdataset(downloaded_paths)#, combine='by_coords', decode_times=True)
                    ds = ds[["sst"]].drop_vars("lev").squeeze()

                    # --- PRE-PROCESSING SPECIFIC TO ERSST ---
                    
                    # 1. Standardize Coordinates (Lat/Lon/Time)
                    # ERSST uses 'lat', 'lon', 'time' usually.
                    # ERSST Longitude is 0 to 360. We likely want -180 to 180 to match ERA5/bounding box.
                    if 'lon' in ds.coords:
                        ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
                        ds = ds.sortby('lon')
                    
                    # 2. Rename to internal standard (X, Y)
                    rename_dict = {}
                    if 'lat' in ds.coords: rename_dict['lat'] = 'Y'
                    if 'lon' in ds.coords: rename_dict['lon'] = 'X'
                    if 'time' in ds.variables: rename_dict['time'] = 'T' # Rename var to capital SST for consistency
                    if 'SST' in ds.variables: rename_dict['SST'] = 'sst' # Rename var to capital SST for consistency
            
                    
                    ds = ds.rename(rename_dict)

                    print(ds)
                    
                    # 3. Slice the Area
                    # Area format: [N, W, S, E]
                    lat_max, lon_min, lat_min, lon_max = area
                    
                    # Ensure X and Y are sorted for slicing
                    ds = ds.sortby(['X', 'Y'])
                    ds = ds.sel(X=slice(lon_min, lon_max), Y=slice(lat_min, lat_max))

                        
                    ds["T"] = ds["T"].astype("datetime64[ns]")
                    print(ds)
                    # ds = ds.rename({'SST': 'sst'})
                    store_file_path[var] = ds
                    ds.to_netcdf(combined_output_path)
                    print(f"Saved NOAA ERSST data to {combined_output_path}")
                    return store_file_path
                    
                except Exception as e:
                    print(f"Failed to download or process NOAA/SST: {str(e)}")
                    # Optional: clean up partial downloads if needed, or leave for retry
                    import traceback
                    traceback.print_exc()
                continue 


            # if f"{center}.{var}" == "NOAA.SST":
            #     try:
            #         url = build_iridl_url_ersst(
            #             year_start=self.year_start,
            #             year_end=year_end,
            #             bbox=area,
            #             run_avg=None,
            #             month_start="Jan",
            #             month_end="Dec"
            #         )
            #         print(f"Using IRIDL URL: {url}")
        
            #         ds = xr.open_dataset(url, decode_times=False)
            #         ds = decode_cf(ds, "T").rename({"T": "time"}).convert_calendar("proleptic_gregorian", align_on="year").rename({"time": "T"})
            #         ds = ds.assign_coords(T=ds.T - pd.Timedelta(days=15))
            #         ds = ds.rename({'sst': 'SST'})
            #         ds = self._postprocess_ersst(ds, var)
            #         ds['T'] = ds['T'].astype('datetime64[ns]')
            #         ds = ds.rename({'SST': 'sst'})
            #         store_file_path[var] = ds
            #         ds.to_netcdf(combined_output_path)
            #         print(f"Saved NOAA ERSST data to {combined_output_path}")
            #         return store_file_path
            #     except Exception as e:
            #         print(f"Failed to download NOAA.SST: {e}")
            #         continue

            combined_datasets = []
            client = cdsapi.Client()
        
            for year in range(self.year_start, year_end + 1):
                yearly_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
            
                if not force_download and yearly_file_path.exists():
                    print(f"{yearly_file_path} exists. Loading existing file.")
                    try:
                        ds = xr.open_dataset(yearly_file_path).load()
                        if 'latitude' in ds.coords and 'longitude' in ds.coords:
                            ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                        combined_datasets.append(ds)
                        continue
                    except Exception as e:
                        print(f"Failed to load {yearly_file_path}: {e}")
            
                try:
                    if var in variables_2:
                        press_level = var.split("_")[1]
                        dataset = "reanalysis-era5-pressure-levels-monthly-means"
                        request = {
                            "product_type": "monthly_averaged_reanalysis",
                            "variable": variables_2[var],
                            "pressure_level": press_level,
                            "year": str(year),
                            "month": months,
                            "time": "00:00",
                            "area": area,
                            "format": "netcdf",
                        }
                    else:
                        dataset = "reanalysis-era5-single-levels-monthly-means"
                        request = {
                            "product_type": "monthly_averaged_reanalysis",
                            "variable": variables_1.get(var, var),
                            "year": str(year),
                            "month": months,
                            "time": "00:00",
                            "area": area,
                            "format": "netcdf",
                        }
            
                    print(f"Downloading {var} for {year} from {center}...")
                    client.retrieve(dataset, request).download(str(yearly_file_path))
                
                    with xr.open_dataset(yearly_file_path) as ds:
                        if 'latitude' in ds.coords and 'longitude' in ds.coords:
                            ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                        ds = ds.load()
                        combined_datasets.append(ds)
                except Exception as e:
                    print(f"Failed to download/process {var} for {year}: {e}")
                    continue
        
            if combined_datasets:
                print(f"Concatenating {var} datasets...")
                combined_ds = xr.concat(combined_datasets, dim="T")
                combined_ds = combined_ds.drop_vars(["number", "expver"], errors="ignore").squeeze()
                
                if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                    combined_ds = combined_ds - 273.15
                elif var == "PRCP":
                    combined_ds = combined_ds * 1000
                elif var in ["DSWR", "DLWR", "NOLR"]:
                    combined_ds = combined_ds / 86400
                elif var == "SLP":
                    combined_ds = combined_ds / 100
                
                combined_ds = combined_ds.isel(Y=slice(None, None, -1))
                store_file_path[var] = combined_ds
                combined_ds.to_netcdf(combined_output_path)
                print(f"Saved combined dataset to {combined_output_path}")
                
                for year in range(self.year_start, year_end + 1):
                    single_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                    if single_file_path.exists():
                        os.remove(single_file_path)
                        print(f"Deleted yearly file: {single_file_path}")
            else:
                print(f"No data combined for {var}. Check download success.")
        return store_file_path

    def download_models(self, force_download=False):
        """Download and process seasonal forecast and hindcast data."""
        centre_map = {
            "BOM_2": "bom", "ECMWF_51": "ecmwf", "UKMO_604": "ukmo", "UKMO_603": "ukmo",
            "METEOFRANCE_8": "meteo_france", "METEOFRANCE_9": "meteo_france",
            "DWD_21": "dwd", "DWD_22": "dwd", "CMCC_35": "cmcc",
            "NCEP_2": "ncep", "JMA_3": "jma", "ECCC_4": "eccc", "ECCC_5": "eccc"
        }
        system_map = {
            "BOM_2": "2", "ECMWF_51": "51", "UKMO_604": "604", "UKMO_603": "603",
            "METEOFRANCE_8": "8", "METEOFRANCE_9": "9", "DWD_21": "21", "DWD_22": "22",
            "CMCC_35": "35", "NCEP_2": "2", "JMA_3": "3", "ECCC_4": "4", "ECCC_5": "5"
        }
        variables_map = {
            "PRCP": "total_precipitation", "TEMP": "2m_temperature",
            "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
            "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
            "UGRD10": "10m_u_component_of_wind", "VGRD10": "10m_v_component_of_wind",
            "SST": "sea_surface_temperature", "SLP": "mean_sea_level_pressure",
            "DSWR": "surface_solar_radiation_downwards",
            "DLWR": "surface_thermal_radiation_downwards",
            "NOLR": "top_thermal_radiation",
            "HUSS_1000": "specific_humidity", "HUSS_925": "specific_humidity",
            "HUSS_850": "specific_humidity", "UGRD_1000": "u_component_of_wind",
            "UGRD_925": "u_component_of_wind", "UGRD_850": "u_component_of_wind",
            "VGRD_1000": "v_component_of_wind", "VGRD_925": "v_component_of_wind",
            "VGRD_850": "v_component_of_wind",
        }

        centers = [item['model_name'] for item in self.predictor_vars]
        variables = [item['variable'] for item in self.predictor_vars]
        areas = [item['area'] for item in self.predictor_vars]
        selected_centres = [centre_map[k] for k in centers]
        selected_systems = [system_map[k] for k in centers]

        month_of_initialization = (datetime.now() - relativedelta(months=1)).month if self.month_of_initialization is None else self.month_of_initialization
        lead_time = [1, 2, 3, 4, 5] if self.lead_time is None else self.lead_time
        if not isinstance(lead_time, list) or any(l < 1 or l > 12 for l in lead_time):
            raise ValueError("lead_time must be a list of integers between 1 and 12.")
        year_forecast = datetime.now().year if self.year_forecast is None else self.year_forecast
        hindcast_years = [str(year) for year in range(self.clim_year_start, self.clim_year_end + 1)]

        dir_to_save = Path(self.dir_to_save)
        dir_to_save.mkdir(parents=True, exist_ok=True)
        abb_mont_ini = calendar.month_abbr[int(month_of_initialization)]
        season_months = [((month_of_initialization + l - 1) % 12) + 1 for l in lead_time]
        season_str = "".join(calendar.month_abbr[m] for m in season_months)

        store_file_path = {}
        store_hdcst_file_path = {}

        for cent, syst, var, area in zip(selected_centres, selected_systems, variables, areas):
            forecast_file = dir_to_save / f"forecast_{cent}{syst}_{var}_{year_forecast}_{abb_mont_ini}Ic_{season_str}_{lead_time[0]}_{area[0]}{area[1]}{area[2]}{area[3]}.nc"
            hindcast_file = dir_to_save / f"hindcast_{cent}{syst}_{var}_{self.clim_year_start}_{self.clim_year_end}{abb_mont_ini}Ic_{season_str}_{lead_time[0]}_{area[0]}{area[1]}{area[2]}{area[3]}.nc"
               
            if not force_download and forecast_file.exists():
                print(f"Forecast file {forecast_file} exists. Skipping download.")
                ds = xr.open_dataset(forecast_file)
                store_file_path[var] = ds
                ds.close()
            else:
                try:
                    dataset = "seasonal-monthly-pressure-levels" if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]) else "seasonal-monthly-single-levels"
                    forecast_request = {
                        "originating_centre": cent,
                        "system": syst,
                        "variable": variables_map[var],
                        "product_type": ["monthly_mean"],
                        "year": [str(year_forecast)],
                        "month": [f"{month_of_initialization:02d}"],
                        "leadtime_month": lead_time,
                        "data_format": "netcdf",
                        "area": area,
                    }
                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        forecast_request["pressure_level"] = var.split("_")[1]

                    client = cdsapi.Client()
                    client.retrieve(dataset, forecast_request).download(str(forecast_file))
                    print(f"Downloaded forecast: {forecast_file}")

                    ds_forecast = xr.open_dataset(forecast_file)

                    if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                        ds_forecast = ds_forecast - 273.15
                    elif var == "PRCP":
                        ds_forecast = ds_forecast * (1000 * 30 * 24 * 3600)
                    elif var == "SLP":
                        ds_forecast = ds_forecast / 100
                    elif var in ["DSWR", "DLWR", "NOLR"]:
                        ds_forecast = ds_forecast / 86400

                    if self.ensemble_mean and "number" in ds_forecast.dims:
                        ds_forecast = getattr(ds_forecast, self.ensemble_mean)(dim="number", skipna=True)

                    ds_forecast = ds_forecast.isel(latitude=slice(None, None, -1))
                    rename_dict = {
                        "latitude": "Y",
                        "longitude": "X",
                        "time": "T" if "time" in ds_forecast.coords else None,
                        "indexing_time": "T" if "indexing_time" in ds_forecast.coords else None,
                        "forecast_reference_time": "T" if "forecast_reference_time" in ds_forecast.coords else None,
                    }
                    rename_dict = {k: v for k, v in rename_dict.items() if v is not None}
                    ds_forecast = ds_forecast.rename(rename_dict)


                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        ds_forecast = ds_forecast.drop_vars("pressure_level", errors="ignore").squeeze()
                    
    
                    ds_forecast.close()
                    os.remove(str(forecast_file))
                    ds_forecast.to_netcdf(str(forecast_file))
                    print(f"Saved processed hindcast to {forecast_file}")
                    store_file_path[var] = ds_forecast

                except Exception as e:
                    print(f"Failed to process {var} for {cent}{syst}: {e}")
                    continue

            if not force_download and hindcast_file.exists():
                print(f"Hindcast file {hindcast_file} exists. Skipping download.")
                ds = xr.open_dataset(hindcast_file)
                store_hdcst_file_path[var] = ds
                ds.close()
            else:
                try:
                    dataset = "seasonal-monthly-pressure-levels" if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]) else "seasonal-monthly-single-levels"
                    hindcast_request = {
                        "originating_centre": cent,
                        "system": syst,
                        "variable": variables_map[var],
                        "product_type": ["monthly_mean"],
                        "year": hindcast_years,
                        "month": [f"{month_of_initialization:02d}"],
                        "leadtime_month": lead_time,
                        "data_format": "netcdf",
                        "area": area,
                    }
                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        hindcast_request["pressure_level"] = var.split("_")[1]
        
                    client = cdsapi.Client()
                    client.retrieve(dataset, hindcast_request).download(str(hindcast_file))
                    print(f"Downloaded hindcast: {hindcast_file}")

                    ds_hindcast = xr.open_dataset(hindcast_file)
                    if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                        ds_hindcast = ds_hindcast - 273.15
                    elif var == "PRCP":
                        ds_hindcast = ds_hindcast * (1000 * 30 * 24 * 3600)
                    elif var == "SLP":
                        ds_hindcast = ds_hindcast / 100
                    elif var in ["DSWR", "DLWR", "NOLR"]:
                        ds_hindcast = ds_hindcast / 86400

                    if self.ensemble_mean and "number" in ds_hindcast.dims:
                        ds_hindcast = getattr(ds_hindcast, self.ensemble_mean)(dim="number", skipna=True)                        

                    ds_hindcast = ds_hindcast.isel(latitude=slice(None, None, -1))
                    rename_dict = {
                        "latitude": "Y",
                        "longitude": "X",
                        "time": "T" if "time" in ds_hindcast.coords else None,
                        "indexing_time": "T" if "indexing_time" in ds_hindcast.coords else None,
                        "forecast_reference_time": "T" if "forecast_reference_time" in ds_hindcast.coords else None,
                    }
                    rename_dict = {k: v for k, v in rename_dict.items() if v is not None}
                    ds_hindcast = ds_hindcast.rename(rename_dict)

                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        ds_hindcast = ds_hindcast.drop_vars("pressure_level", errors="ignore").squeeze()

                    ds_hindcast = ds_hindcast.sortby("T")
                    ds_hindcast.close()
                    os.remove(hindcast_file)
                    ds_hindcast.to_netcdf(hindcast_file)
                    print(f"Saved processed hindcast to {hindcast_file}")
                    store_hdcst_file_path[var] = ds_hindcast


                except Exception as e:
                    print(f"Failed to process {var} for {cent}{syst}: {e}")
                    continue

        return store_hdcst_file_path, store_file_path

    def anomaly_timeseries(self, ds):
        """Compute anomalies by removing the climatological mean."""
        if self.clim_year_start and self.clim_year_end:
            clim_period = ds.sel(T=slice(f"{self.clim_year_start}-01-01", f"{self.clim_year_end}-12-31"))
        else:
            clim_period = ds

        clim_mean = clim_period.groupby("T.month").mean("T", skipna=True)
        anomaly_slices = []

        for month in range(1, 13):
            month_mask = ds["T"].dt.month == month
            ds_month = ds.where(month_mask, drop=True)
            month_mean = clim_mean.sel(month=month)
            ds_anomaly_month = ds_month - month_mean
            anomaly_slices.append(ds_anomaly_month)

        ds_anomaly = xr.concat(anomaly_slices, dim="T")
        ds_anomaly = ds_anomaly.sortby("T")
        if "month" in ds_anomaly.coords:
            ds_anomaly = ds_anomaly.drop_vars("month")
        if "month" in ds_anomaly.dims:
            ds_anomaly = ds_anomaly.drop_dims("month")
        return ds_anomaly

    def standardize_timeseries(self, ds):
        """Standardize the dataset by removing climatological mean and scaling by standard deviation."""
        if self.clim_year_start and self.clim_year_end:
            clim_period = ds.sel(T=slice(f"{self.clim_year_start}-01-01", f"{self.clim_year_end}-12-31"))
        else:
            clim_period = ds

        clim_mean = clim_period.groupby("T.month").mean("T", skipna=True)
        clim_std = clim_period.groupby("T.month").std("T", skipna=True)
        clim_std = clim_std.where(clim_std != 0, 1e-10)

        standardized_slices = []

        for month in range(1, 13):
            month_mask = ds["T"].dt.month == month
            ds_month = ds.where(month_mask, drop=True)
            month_mean = clim_mean.sel(month=month)
            month_std = clim_std.sel(month=month)
            ds_standardized_month = (ds_month - month_mean) / month_std
            standardized_slices.append(ds_standardized_month)

        ds_standardized = xr.concat(standardized_slices, dim="T")
        ds_standardized = ds_standardized.sortby("T")
        if "month" in ds_standardized.coords:
            ds_standardized = ds_standardized.drop_vars("month")
        if "month" in ds_standardized.dims:
            ds_standardized = ds_standardized.drop_dims("month")
        return ds_standardized

    def compute_eofs(self, data):
        """Compute EOFs and retain modes explaining at least the specified variance."""
        data = data.rename({"X": "lon", "Y": "lat"})
        data = data.fillna(data.mean(dim="T", skipna=True))
        print("Computing EOFs...")
        model = xe.single.EOF(n_modes=100, use_coslat=True, center=False)
        model.fit(data, dim='T')
        explained_var = model.explained_variance_ratio()
        cumulative_var = explained_var.cumsum()
        n_modes = np.where(cumulative_var >= self.eof_explained_var)[0][0] + 1
        print(f"Selected {n_modes} modes explaining {cumulative_var[n_modes-1].values*100:.2f}% variance")
        return model.scores().isel(mode=slice(0, n_modes))

    # def _detrended_da(self, da):
    #     """Detrend a DataArray by removing the linear trend."""
    #     if 'T' not in da.dims:
    #         raise ValueError("DataArray must have a time dimension 'T' for detrending.")
    #     trend = da.polyfit(dim='T', deg=1)
    #     da_detrended = da - (trend.polyval(da['T']) if 'polyval' in dir(trend) else trend)
    #     return da_detrended.isel(degree=0, drop=True).to_array().drop_vars('variable').squeeze(),\
    #          trend.isel(degree=0, drop=True).to_array().drop_vars('variable').squeeze()

    def _detrended_da(self, da, dim='T'):
            """
            Detrend a DataArray by removing the linear trend.
            """
            if dim not in da.dims:
                raise ValueError(f"DataArray must have a time dimension '{dim}' for detrending.")
    
            # 1. Calculate coefficients (Slope + Intercept)
            # Returns a Dataset containing 'polyfit_coefficients'
            coeffs = da.polyfit(dim=dim, deg=1)
    
            # 2. Reconstruct the trend line using standard xarray syntax
            trend = xr.polyval(da[dim], coeffs.polyfit_coefficients)
    
            # 3. Return anomalies and the FULL coefficients 
            # (We need the full coefficients to detrend future data for prediction)
            return da - trend, coeffs

    def _prepare_data_for_clustering(self, unique_years, target_year=None):
        """Helper to prepare data (indices or EOFs) for clustering methods."""
        _, ddd = self.download_and_process()
        data_scaled = None
        index_data = None

        if self.index_compute:
            ddd_sst = ddd.get('SST')
            if self.detrend:
                ddd_sst, _ = self._detrended_da(ddd_sst)
            if ddd_sst is None:
                raise ValueError("SST data not found in downloaded datasets for clustering.")
            indices_dataset = self.calc_index(self.index_compute, ddd_sst)
            indices_dataset_sel_years = indices_dataset.where(indices_dataset['T'].dt.year.isin(unique_years), drop=True)
            data_scaled, index_data = self.arrange_indices_for_som(indices_dataset_sel_years)
        else:
            if self.multivariateEOF:
                data_vars = [self._detrended_da(ddd[var].rename({"X": "lon", "Y": "lat"}))[0] if self.detrend else ddd[var].rename({"X": "lon", "Y": "lat"}) for var in ddd]
                model = xe.single.EOF(n_modes=100, use_coslat=True, center=False)
                model.fit(data_vars, dim='T')
                explained_var = model.explained_variance_ratio()
                n_modes = np.where(explained_var.cumsum() >= self.eof_explained_var)[0][0] + 1
                scores = model.scores().isel(mode=slice(0, n_modes))
                scores = scores.where(scores['T'].dt.year.isin(unique_years), drop=True)
                
                df_final = pd.DataFrame()
                scores = scores.assign_coords(year=('T', scores['T'].dt.year.data), month_name=('T', scores['T'].dt.strftime('%b').data))
                df = scores.to_dataframe(name='value').reset_index()
                df['mode_month'] = df['mode'].apply(lambda m: f"mode{m:02d}") + "_" + df['month_name']
                df_pivot = df.pivot(index='year', columns='mode_month', values='value')
                month_order = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
                df_pivot = df_pivot.reindex(columns=sorted(df_pivot.columns, key=lambda x: (month_order.index(x.split('_')[1]), x)))
                df_final = df_pivot
                data_scaled = np.array(df_final, dtype=np.float64)
                index_data = df_final.index
            else:
                df_final = pd.DataFrame()
                for var in ddd:
                    if self.detrend:
                        ddd[var], _ = self._detrended_da(ddd[var])
                    scores = self.compute_eofs(ddd[var])
                    scores = scores.where(scores['T'].dt.year.isin(unique_years), drop=True)
                    scores = scores.assign_coords(year=('T', scores['T'].dt.year.data), month_name=('T', scores['T'].dt.strftime('%b').data))
                    df = scores.to_dataframe(name='value').reset_index()
                    df['mode_month'] = df['mode'].apply(lambda m: f"mode{m:02d}") + "_" + df['month_name'] + f"_{var}"
                    df_pivot = df.pivot(index='year', columns='mode_month', values='value')
                    month_order = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
                    df_pivot = df_pivot.reindex(columns=sorted(df_pivot.columns, key=lambda x: (month_order.index(x.split('_')[1]), x)))
                    df_final = pd.concat([df_final, df_pivot], axis=1)
                data_scaled = np.array(df_final, dtype=np.float64)
                index_data = df_final.index
        
        if np.isnan(data_scaled).any():
            print("NaN values found in data for clustering. Imputing with column mean.")
            col_means = np.nanmean(data_scaled, axis=0)
            nan_indices = np.where(np.isnan(data_scaled))
            data_scaled[nan_indices] = col_means[nan_indices[1]]
            
        return data_scaled, index_data

    def arrange_indices_for_som(self, ds):
        """Arrange climate indices for SOM training."""
        ds = ds.assign(year=('T', ds['T'].dt.year.data), month_name=('T', ds['T'].dt.strftime('%b').data))
        
        def pivot_var(var_name):
            df = ds[[var_name, 'year', 'month_name']].to_dataframe().reset_index()
            df_pivot = df.pivot(index='year', columns='month_name', values=var_name)
            df_pivot = df_pivot.reindex(columns=['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'])
            df_pivot.columns = [f"{var_name}_{month}" for month in df_pivot.columns]
            return df_pivot
        
        vars_order = self.index_compute or []
        df_vars = {v: pivot_var(v) for v in vars_order}
        months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
        ordered_columns = [f"{v}_{m}" for m in months for v in vars_order]
        
        df_final = pd.concat(df_vars.values(), axis=1)[ordered_columns]
        return np.array(df_final, dtype=np.float64), df_final.index

    def identify_analogs_bmu(self, target_year, bmu_dict):
        """Identify analog years based on SOM Best Matching Units (BMUs)."""
        radius = self.radius or 1.0
        neuron_years = defaultdict(list)
        for year, coords in bmu_dict.items():
            neuron_years[coords].append(year)
        
        target_coords = bmu_dict.get(target_year)
        if not target_coords:
            return []
        
        analogs = [year for year in neuron_years[target_coords] if year != target_year]
        if not analogs:
            for year, coords in bmu_dict.items():
                if year == target_year:
                    continue
                dist = np.sqrt((coords[0] - target_coords[0])**2 + (coords[1] - target_coords[1])**2)
                if dist <= radius:
                    analogs.append(year)
        
        return analogs

    def SOM(self, predictant, itrain, ireference_year):
        """Identify similar years using Self-Organizing Maps (SOM)."""

        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)
        som_grid_size = (4, 4) if set(self.some_grid_size) == {None} else self.some_grid_size
        som = MiniSom(
            x=som_grid_size[0],
            y=som_grid_size[1],
            input_len=data_scaled.shape[1],
            sigma=self.some_sigma,
            learning_rate=self.some_learning_rate,
            neighborhood_function=self.some_neighborhood_function,
            random_seed=42
        )
        som.random_weights_init(data_scaled)
        som.train_random(data=data_scaled, num_iteration=self.some_num_iteration)
        
        bmu_coords = [som.winner(x) for x in data_scaled]
        bmu_dict = {year: coords for year, coords in zip(index_data, bmu_coords)}
        similar_years = self.identify_analogs_bmu(reference_year, bmu_dict)
        print(f"Similar years for {reference_year}: {similar_years}")
        return np.array(similar_years)


    def Bias_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using bias-based analog method."""
        _, ddd = self.download_and_process()
        ddd = ddd.get(self.predictor_vars[0]['variable'])
        if ddd is None:
            raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
       
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
        sst_reference = ddd.sel(T=str(reference_year))
       
        biases = []
        for year in unique_years:
            tmp = ddd.sel(T=str(year))
            tmp['T'] = sst_reference['T'] 
            bias = np.abs(tmp - sst_reference).mean(dim="T")
            biases.append(bias)
       
        similar = xr.concat(biases, dim='T').assign_coords(T=unique_years).compute()
        similar = xr.where(similar < 0.5, 1, 0).sum(dim=["X", "Y"])
        similar = similar.sortby(similar, ascending=False)
        top_2 = similar.isel(T=slice(1, 3))
        similar_years = top_2['T'].to_numpy()
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years
    
    def Corr_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using correlation-based analog method."""

        _, ddd = self.download_and_process()

        ddd = ddd.get(self.predictor_vars[0]['variable'])

        if ddd is None:
            raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
        
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
        sst_reference = ddd.sel(T=str(reference_year))
        
        correlations = []
        for year in unique_years:
            tmp = ddd.sel(T=str(year))
            tmp['T'] = sst_reference['T']
            correlation = xr.corr(tmp, sst_reference, dim="T").compute()
            correlations.append(correlation)
        
        similar = xr.concat(correlations, dim='T').assign_coords(T=unique_years)
        similar = xr.where(similar > 0.3, 1, 0).sum(dim=["X", "Y"])
        similar = similar.sortby(similar, ascending=False)
        top_2 = similar.isel(T=slice(1, 3))
        similar_years = top_2['T'].to_numpy()
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years


    # def Pca_Based(self, predictant, itrain, ireference_year):
    #     """Identify similar years using PCA-based analog method."""
    #     _, ddd = self.download_and_process()
    #     ddd = ddd.get(self.predictor_vars[0]['variable'])
    #     if ddd is None:
    #         raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
        
    #     # predictor_ = ddd.fillna(ddd.groupby("T.month").mean("T", skipna=True))
    #     # predictor_detrend = sig.detrend(predictor_, axis=0)
    #     # ddd = xr.DataArray(predictor_detrend, dims=predictor_.dims, coords=predictor_.coords)

    #     if self.detrend:
    #         ddd, _ = self._detrended_da(ddd)
    #     # ddd = ddd.rename({"X": "lon", "Y": "lat"})
    #     # eof = xe.single.EOF(n_modes=50, use_coslat=True, center=False)
    #     # eof.fit(ddd.fillna(ddd.mean(dim="T", skipna=True)), dim="T")
    #     # scores = eof.scores()
    #     scores = self.compute_eofs(ddd)
        
    #     predictant = predictant.copy()
    #     predictant['T'] = predictant['T'].astype('datetime64[ns]')

    #     if ireference_year==2100:
    #         reference_year = self.year_forecast
    #         predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
    #     else:
    #         reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
    #         predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")

    #     sst_ref = scores.sel(T=str(reference_year)).stack(score=('mode', 'T'))       
    #     unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

    #     correlations = []
    #     for year in unique_years:
    #         tmp = scores.sel(T=str(year)).stack(score=('mode', 'T'))
    #         correlation = xr.corr(tmp, sst_ref, dim="score").compute()
    #         correlations.append(correlation)
        
    #     similar = xr.concat(correlations, dim='T').assign_coords(T=unique_years)
    #     similar = similar.sortby(similar, ascending=False)
    #     top_2 = similar.isel(T=slice(1,3))
    #     similar_years = top_2['T'].to_numpy()
    #     print(f"Similar years for {reference_year}: {similar_years}")
    #     return similar_years

    def Pca_Based(self, predictant, itrain, ireference_year):
            """Identify similar years using PCA-based analog method."""
            _, ddd = self.download_and_process()
            ddd = ddd.get(self.predictor_vars[0]['variable'])
            if ddd is None:
                raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
           
            # predictor_ = ddd.fillna(ddd.groupby("T.month").mean("T", skipna=True))
            # predictor_detrend = sig.detrend(predictor_, axis=0)
            # ddd = xr.DataArray(predictor_detrend, dims=predictor_.dims, coords=predictor_.coords)
            if self.detrend:
                ddd, _ = self._detrended_da(ddd)
            # ddd = ddd.rename({"X": "lon", "Y": "lat"})
            # eof = xe.single.EOF(n_modes=50, use_coslat=True, center=False)
            # eof.fit(ddd.fillna(ddd.mean(dim="T", skipna=True)), dim="T")
            # scores = eof.scores()
            scores = self.compute_eofs(ddd)
           
            predictant = predictant.copy()
            predictant['T'] = predictant['T'].astype('datetime64[ns]')
            if ireference_year==2100:
                reference_year = self.year_forecast
                predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
            else:
                reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
                predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
            unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
            
            ref_scores = scores.sel(T=slice(f"{reference_year}-01-01", f"{reference_year}-12-31"))
            ref_scores = ref_scores.assign_coords(T=np.arange(ref_scores.sizes['T']))
            sst_ref = ref_scores.stack(score=('mode', 'T'))
           
            correlations = []
            for year in unique_years:
                tmp_scores = scores.sel(T=slice(f"{year}-01-01", f"{year}-12-31"))
                tmp_scores = tmp_scores.assign_coords(T=np.arange(tmp_scores.sizes['T']))
                tmp = tmp_scores.stack(score=('mode', 'T'))
                correlation = xr.corr(tmp, sst_ref, dim="score").compute()
                correlations.append(correlation)
           
            similar = xr.concat(correlations, dim='T').assign_coords(T=unique_years)
            similar = similar.sortby(similar, ascending=False)
            top_2 = similar.isel(T=slice(1,3))
            similar_years = top_2['T'].to_numpy()
            print(f"Similar years for {reference_year}: {similar_years}")
            return similar_years

    def KMeans_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using K-Means clustering."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")

        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)

        kmeans = KMeans(n_clusters=self.n_clusters, random_state=42, n_init='auto')
        cluster_labels = kmeans.fit_predict(data_scaled)

        target_year_idx = np.where(index_data == reference_year)[0]
        if len(target_year_idx) == 0:
            print(f"Target year {reference_year} not found in index data.")
            return np.array([])
        target_cluster_label = cluster_labels[target_year_idx[0]]

        similar_years_indices = np.where(cluster_labels == target_cluster_label)[0]
        similar_years = [index_data[i] for i in similar_years_indices if index_data[i] != reference_year]
        
        print(f"Similar years for {reference_year} (K-Means): {similar_years}")
        return np.array(similar_years)

    def Agglomerative_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using Agglomerative Clustering."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")

        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)

        agg_clustering = AgglomerativeClustering(n_clusters=self.n_clusters, metric=self.affinity, linkage=self.linkage)
        cluster_labels = agg_clustering.fit_predict(data_scaled)

        target_year_idx = np.where(index_data == reference_year)[0]
        if len(target_year_idx) == 0:
            print(f"Target year {reference_year} not found in index data.")
            return np.array([])
        target_cluster_label = cluster_labels[target_year_idx[0]]

        similar_years_indices = np.where(cluster_labels == target_cluster_label)[0]
        similar_years = [index_data[i] for i in similar_years_indices if index_data[i] != reference_year]

        print(f"Similar years for {reference_year} (Agglomerative): {similar_years}")
        return np.array(similar_years)

    # def download_and_process(self):
    #     """Download and process reanalysis and forecast data.

    #     Combines reanalysis and forecast data, applies standardization or anomaly calculation,
    #     and optionally applies a rolling mean.

    #     Returns
    #     -------
    #     data_var_concatenated : dict
    #         Dictionary of concatenated, processed datasets.
    #     data_var_shifted : dict
    #         Dictionary of time-shifted, processed datasets.
    #     """
    #     lead_time = [1, 2, 3, 4, 5] if self.lead_time is None else self.lead_time
    #     month_of_initialization = (datetime.now() - relativedelta(months=1)).month if self.month_of_initialization is None else self.month_of_initialization
        
    #     sst_hist = self.download_reanalysis(force_download=False)
    #     sst_hdcst, sst_for = self.download_models(force_download=False)

    #     variables = [item['variable'] for item in self.predictor_vars]
    #     print(f"Processing variables: {variables}")
    #     data_var_concatenated = {}
    #     data_var_shifted = {}

    #     for var in variables:
    #         if var not in sst_hist or var not in sst_for:
    #             print(f"Skipping {var}: data not available.")
    #             continue  

    #         sst_hist_ = sst_hist[var]
    #         sst_hdcst_ = sst_hdcst[var]
    #         sst_for_ = sst_for[var]

    #         # process hindcast data
    #         hindcast = []
    #         for i in sst_hdcst_['T']:
    #             sst_hdcst_i = sst_hdcst_.sel(T=i)
    #             base_times = pd.Timestamp(sst_hdcst_i['T'].values)
    #             new_times = [base_times + pd.DateOffset(months=int(m)) for m in sst_hdcst_i['forecastMonth'].values]
    #             sst_hdcst_i = sst_hdcst_i.assign_coords(forecastMonth=("forecastMonth", new_times))
    #             sst_hdcst_i = sst_hdcst_i.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})
    #             hindcast.append(sst_hdcst_i)
    #         sst_hdcst_ = xr.concat(hindcast, dim='T')          
    #         sst_hdcst_ = sst_hdcst_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})

            
    #         # process forecast data
    #         sst_for_ = sst_for_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})
    #         if (isinstance(sst_for_['T'].values, np.ndarray)):
    #             base_time = pd.Timestamp(sst_for_['T'].values[-1])
    #         else:
    #             base_time = pd.Timestamp(sst_for_['T'].values)
    #         new_times = [base_time + pd.DateOffset(months=int(m)) for m in sst_for_['forecastMonth'].values]
    #         sst_for_ = sst_for_.assign_coords(forecastMonth=("forecastMonth", new_times))
    #         sst_for_ = sst_for_.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})

    #         # Correct forecast systematic bias
    #         biasr = WAS_bias_correction()

    #         sst_for_bias_corrected = []
    #         for m in lead_time:
    #             target_time = base_time + pd.DateOffset(months=m)
    #             current_target_month = target_time.month

    #             sst_hist_mean = sst_hist_.where(sst_hist_['T'].dt.month.isin(current_target_month), drop=True)
    #             sst_hist_mean = sst_hist_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()

    #             sst_hdcst_mean = sst_hdcst_.where(sst_hdcst_['T'].dt.month.isin(current_target_month), drop=True)
    #             sst_hdcst_mean = sst_hdcst_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    #             sst_hist_mean, sst_hdcst_mean  = xr.align(sst_hist_mean, sst_hdcst_mean, join="inner")
    #             sst_for_mean = sst_for_.where(sst_for_['T'].dt.month.isin(current_target_month), drop=True)
    #             sst_for_mean = sst_for_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze('variable', drop=True)
                
    #             if var in ["TMIN", "TEMP", "TMAX", "SST", "HUSS_1000", "HUSS_925", "HUSS_850", "SLP", "UGRD10", "VGRD10", "UGRD_1000", "UGRD_925", "UGRD_850", "VGRD_1000", "VGRD_925", "VGRD_850"]:
    #                 fobj_quant_da = biasr.fitBC(sst_hist_mean, sst_hdcst_mean, method='QUANT', qstep=0.01, nboot=10)
    #                 sst_for_bias_corrected.append(biasr.doBC(sst_for_mean, fobj_quant_da, type='linear'))
    #                 # sst_for_bias_corrected.append(sst_for_.where(sst_for_['T'].dt.month.isin(m), drop=True) - sst_hdcst_mean + sst_hist_mean)
    #             # if var in ["PRCP", "DSWR", "DLWR", "N.OLR"]:
    #                 # fobj_quant = qmap.fitQmap(sst_hist_mean, sst_hdcst_mean, method='QUANT', wet_day=0.1,  qstep=0.001)
    #                 # sst_for_bias_corrected.append(sst_for_.where(sst_for_['T'].dt.month.isin(m), drop=True) * sst_hist_mean / sst_hdcst_mean)
    #         sst_for_ = xr.concat(sst_for_bias_corrected, dim='T').sortby("T")

    #         # Concatenate hist and forecast data
    #         sst_hist_ = sst_hist_.sel(T=slice(f"{self.year_start}-01", f"{base_time.year}-{base_time.month:02d}"))
    #         sst_hist_ = sst_hist_.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    #         concatenated_ds = xr.concat([sst_hist_, sst_for_], dim='T')
    #         concatenated_ds = concatenated_ds.sortby("T")
    #         concatenated_ds = concatenated_ds.to_dataset(name=var)

    #         if isinstance(self.rolling, int) and self.rolling > 1:
    #            print(f"Applied rolling mean with window size {self.rolling} for {var}.")
    #            concatenated_ds = concatenated_ds.rolling(T=self.rolling, center=False, min_periods=self.rolling).mean()

    #         if self.standardize:
    #             print(f"Standardizing timeseries for {var}.")
    #             concatenated_ds_st = self.standardize_timeseries(concatenated_ds)
    #         else:
    #             print(f"Computing anomalies for {var}.")
    #             concatenated_ds_st = self.anomaly_timeseries(concatenated_ds)

    #         # first_year = concatenated_ds_st.isel(T=0).T.dt.year.item()
    #         # last_month = concatenated_ds_st.isel(T=-1).T.dt.month.item()

    #         select_parts = concatenated_ds_st['T'].size - (self.year_forecast - self.year_start) * 12 

    #         # if last_month == 12:
    #         #     start_month = 1
    #         #     first_year += 1
    #         # else:
    #         #     start_month = last_month + 1
            
    #         concatenated_ds_st = concatenated_ds_st.isel(T=slice(select_parts, None))
            
    #         # concatenated_ds_st = concatenated_ds_st.sel(T=slice(f"{first_year}-{start_month:02d}", f"{self.year_forecast}-{last_month:02d}"))
    #         # if last_month == 12:
    #         #     new_time = pd.DatetimeIndex([pd.to_datetime(f"{first_year}-01-01") + pd.DateOffset(months=t) for t in range(len(concatenated_ds_st['T']))])
    #         # else:
    #         #     new_time = pd.DatetimeIndex([pd.to_datetime(f"{first_year+1}-01-01") + pd.DateOffset(months=t) for t in range(len(concatenated_ds_st['T']))])    

    #         new_time = pd.DatetimeIndex([pd.to_datetime(f"{self.year_start+1}-01-01") + pd.DateOffset(months=t) for t in range(len(concatenated_ds_st['T']))]) 

    #         ds_shifted = concatenated_ds_st.assign_coords(T=new_time)
    
    #         data_var_shifted[var] = ds_shifted.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    #         data_var_concatenated[var] = concatenated_ds_st.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    #     return data_var_concatenated, data_var_shifted


    def download_and_process(self):
            """Download and process reanalysis and forecast data."""
            
            lead_time = [1, 2, 3, 4, 5] if self.lead_time is None else self.lead_time
            month_of_initialization = (datetime.now() - relativedelta(months=1)).month if self.month_of_initialization is None else self.month_of_initialization
            
            sst_hist = self.download_reanalysis(force_download=False)
            sst_hdcst, sst_for = self.download_models(force_download=False)
    
            variables = [item['variable'] for item in self.predictor_vars]
            print(f"Processing variables: {variables}")
            data_var_concatenated = {}
            data_var_shifted = {}
    
            for var in variables:
                if var not in sst_hist or var not in sst_for:
                    print(f"Skipping {var}: data not available.")
                    continue  
    
                sst_hist_ = sst_hist[var]
                sst_hdcst_ = sst_hdcst[var]
                sst_for_ = sst_for[var]
    
                # --- Process Hindcast Data ---
                hindcast = []
                for i in sst_hdcst_['T']:
                    sst_hdcst_i = sst_hdcst_.sel(T=i)
                    base_times = pd.Timestamp(sst_hdcst_i['T'].values)
                    new_times = [base_times + pd.DateOffset(months=int(m)) for m in sst_hdcst_i['forecastMonth'].values]
                    sst_hdcst_i = sst_hdcst_i.assign_coords(forecastMonth=("forecastMonth", new_times))
                    sst_hdcst_i = sst_hdcst_i.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})
                    hindcast.append(sst_hdcst_i)
                
                sst_hdcst_ = xr.concat(hindcast, dim='T')           
                sst_hdcst_ = sst_hdcst_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})
    
                # --- Process Forecast Data ---
                sst_for_ = sst_for_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})
                
                if isinstance(sst_for_['T'].values, np.ndarray):
                    base_time = pd.Timestamp(sst_for_['T'].values[-1])
                else:
                    base_time = pd.Timestamp(sst_for_['T'].values)
                    
                new_times = [base_time + pd.DateOffset(months=int(m)) for m in sst_for_['forecastMonth'].values]
                sst_for_ = sst_for_.assign_coords(forecastMonth=("forecastMonth", new_times))
                sst_for_ = sst_for_.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})
    
                # --- Bias Correction ---
                biasr = WAS_bias_correction()
                sst_for_bias_corrected = []
                
                for m in lead_time:
                    # 1. Calcul correct du mois cible (gestion du cycle 12 -> 1)
                    target_time = base_time + pd.DateOffset(months=m)
                    current_target_month = target_time.month
    
                    # 2. Filtrage
                    sst_hist_mean = sst_hist_.where(sst_hist_['T'].dt.month == current_target_month, drop=True)
                    sst_hist_mean = sst_hist_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    
                    sst_hdcst_mean = sst_hdcst_.where(sst_hdcst_['T'].dt.month == current_target_month, drop=True)
                    sst_hdcst_mean = sst_hdcst_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                    
                    # 3. Alignement
                    # Create a set of unique years from the hindcast data
                    years_to_keep = sst_hdcst_mean['T'].dt.year.to_numpy()
                    
                    # Filter the historical data based on this set of years
                    sst_hist_mean = sst_hist_mean.where(sst_hist_mean['T'].dt.year.isin(years_to_keep), drop=True)
                    sst_hdcst_mean['T'] = sst_hist_mean['T']
                    sst_hist_aligned, sst_hdcst_aligned  = xr.align(sst_hist_mean, sst_hdcst_mean, join="inner")
                    
                    # 4. SÉCURITÉ CRITIQUE : Vérifier si les données sont vides après alignement
                    if sst_hist_aligned.sizes['T'] == 0 or sst_hdcst_aligned.sizes['T'] == 0:
                        print(f"WARNING: No overlapping data found for month {current_target_month} (lead {m}). Skipping bias correction for this step.")
                        continue # Passe à l'itération suivante pour éviter le crash
    
                    sst_for_mean = sst_for_.where(sst_for_['T'].dt.month == current_target_month, drop=True)
                    sst_for_mean = sst_for_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze('variable', drop=True)
                    
                    if var in ["TMIN", "TEMP", "TMAX", "SST", "HUSS_1000", "HUSS_925", "HUSS_850", "SLP", "UGRD10", "VGRD10", "UGRD_1000", "UGRD_925", "UGRD_850", "VGRD_1000", "VGRD_925", "VGRD_850"]:
                        # Utilisation des variables alignées (_aligned)
                        fobj_quant_da = biasr.fitBC(sst_hist_aligned, sst_hdcst_aligned, method='QUANT', qstep=0.01, nboot=10)
                        sst_for_bias_corrected.append(biasr.doBC(sst_for_mean, fobj_quant_da, type='linear'))
                
                # --- Concatenation & Post-Processing ---
                if not sst_for_bias_corrected:
                    print(f"ERROR: Bias correction failed for all lead times for variable {var}.")
                    continue
    
                sst_for_ = xr.concat(sst_for_bias_corrected, dim='T').sortby("T")
    
                sst_hist_ = sst_hist_.sel(T=slice(f"{self.year_start}-01", f"{base_time.year}-{base_time.month:02d}"))
                sst_hist_ = sst_hist_.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                
                concatenated_ds = xr.concat([sst_hist_, sst_for_], dim='T')
                concatenated_ds = concatenated_ds.sortby("T")
                concatenated_ds = concatenated_ds.to_dataset(name=var)
    
                if isinstance(self.rolling, int) and self.rolling > 1:
                    print(f"Applied rolling mean with window size {self.rolling} for {var}.")
                    concatenated_ds = concatenated_ds.rolling(T=self.rolling, center=False, min_periods=self.rolling).mean()
    
                if self.standardize:
                    print(f"Standardizing timeseries for {var}.")
                    concatenated_ds_st = self.standardize_timeseries(concatenated_ds)
                else:
                    print(f"Computing anomalies for {var}.")
                    concatenated_ds_st = self.anomaly_timeseries(concatenated_ds)
    
                select_parts = concatenated_ds_st['T'].size - (self.year_forecast - self.year_start) * 12 
                concatenated_ds_st = concatenated_ds_st.isel(T=slice(select_parts, None))
                
                new_time = pd.DatetimeIndex([pd.to_datetime(f"{self.year_start+1}-01-01") + pd.DateOffset(months=t) for t in range(len(concatenated_ds_st['T']))]) 
    
                ds_shifted = concatenated_ds_st.assign_coords(T=new_time)
        
                data_var_shifted[var] = ds_shifted.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                data_var_concatenated[var] = concatenated_ds_st.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                
            return data_var_concatenated, data_var_shifted

    def compute_model(self, predictant, itrain, itest):
        """Compute deterministic hindcast using the specified analog method."""
        method_map = {
            "som": self.SOM,
            "bias_based": self.Bias_Based,
            "cor_based": self.Corr_Based,
            "pca_based": self.Pca_Based,
            "kmeans_based": self.KMeans_Based,
            "agglomerative_based": self.Agglomerative_Based
        }
        if self.method_analog not in method_map:
            raise ValueError(f"Invalid analog method: {self.method_analog}. Choose 'som', 'bias_based', 'cor_based', 'pca_based', 'kmeans_based', or 'agglomerative_based'.")

        similar_years = method_map[self.method_analog](predictant, itrain, itest)
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        
        sim_obs = [predictant.sel(T=str(year)) for year in similar_years if year in predictant['T'].dt.year.values]
        if not sim_obs:
            raise ValueError("No valid similar years found for hindcast.")
        
        hindcast_det = xr.concat(sim_obs, dim="T").mean(dim="T").expand_dims({'T': predictant.isel(T=itest)['T'].values})
        return hindcast_det

    # ------------------ 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.3, loc=loc, scale=scale),
                    norm.ppf(0.68, loc=loc, scale=scale),
                )
            elif code == 2:
                return (
                    lognorm.ppf(0.3, s=shape, loc=loc, scale=scale),
                    lognorm.ppf(0.68, s=shape, loc=loc, scale=scale),
                )
            elif code == 3:
                return (
                    expon.ppf(0.3, loc=loc, scale=scale),
                    expon.ppf(0.68, loc=loc, scale=scale),
                )
            elif code == 4:
                return (
                    gamma.ppf(0.3, a=shape, loc=loc, scale=scale),
                    gamma.ppf(0.68, a=shape, loc=loc, scale=scale),
                )
            elif code == 5:
                return (
                    weibull_min.ppf(0.3, c=shape, loc=loc, scale=scale),
                    weibull_min.ppf(0.68, c=shape, loc=loc, scale=scale),
                )
            elif code == 6:
                # Note: Renamed 't_dist' to 't' for standard scipy.stats
                return (
                    t.ppf(0.3, df=shape, loc=loc, scale=scale),
                    t.ppf(0.68, 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.3, mu=shape, loc=loc),
                    poisson.ppf(0.68, 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.33, n=shape, p=scale, loc=loc),
                    nbinom.ppf(0.68, 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.33, 0.68], 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, hindcast_det, best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """Generate deterministic and probabilistic forecasts for the target year."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        reference_year = self.year_forecast
        unique_years = np.append(np.unique(predictant['T'].dt.year), reference_year)
        
        method_map = {
            "som": self.SOM,
            "bias_based": self.Bias_Based,
            "cor_based": self.Corr_Based,
            "pca_based": self.Pca_Based,
            "kmeans_based": self.KMeans_Based,
            "agglomerative_based": self.Agglomerative_Based
        }
        if self.method_analog not in method_map:
            raise ValueError(f"Invalid analog method: {self.method_analog}. Choose 'som','bias_based', 'cor_based', 'pca_based', 'kmeans_based', or 'agglomerative_based'.")

        itrain = [i for i, year in enumerate(predictant['T'].dt.year.values) if year in unique_years and year != reference_year]
        ireference_year = 2100

        similar_years = method_map[self.method_analog](predictant, itrain, ireference_year)
        if len(similar_years)>2:
            similar_years = similar_years[:2]
        else:
            similar_years = similar_years[:len(similar_years)]

        sim_obs = [predictant.sel(T=str(year)) for year in similar_years if year in predictant['T'].dt.year.values]

        if not sim_obs:
            raise ValueError("No valid similar years found for forecast.")
        
        forecast_det = xr.concat(sim_obs, dim="T").mean(dim="T")
        T_value = predictant.isel(T=0).coords['T'].values
        month = T_value.astype('datetime64[M]').astype(int) % 12 + 1
        new_T_value = np.datetime64(f"{reference_year}-{month:02d}-01")
        forecast_det = forecast_det.expand_dims({'T': [new_T_value]})
        forecast_det['T'] = forecast_det['T'].astype('datetime64[ns]')
        
        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.3, 0.68], 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')
        
        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_det,
                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_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}")
        
        forecast_prob = forecast_prob.assign_coords(probability=('probability', ['PB', 'PN', 'PA'])).transpose('probability', 'T', 'Y', 'X')
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years, forecast_det, forecast_prob

    def composite_plot(self, predictant, clim_year_start, clim_year_end, hindcast_det, plot_predictor=True, variable="SST",best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """Create composite plots of predictors or predictands."""
        clim_slice = slice(str(clim_year_start), str(clim_year_end))
        clim_mean = predictant.sel(T=clim_slice).mean(dim='T', skipna=True)
        similar_years, result_, _ = self.forecast(predictant, clim_year_start, clim_year_end, hindcast_det, best_code_da, best_shape_da, best_loc_da, best_scale_da)
        result_ = result_.drop_vars('T', errors='ignore').squeeze()
        reference_year = self.year_forecast
        _, ddd = self.download_and_process()

        ddd = ddd[variable]
        
        sim_all = []
        tmp = ddd.sel(T=str(reference_year))
        months = list(tmp['T'].dt.month.values)
        tmp['T'] = months
        sim_all.append(tmp)
        
        sim_ = []
        pred_rain = []
        for year in np.array([str(i) for i in similar_years]):
            tmp = ddd.sel(T=year)
            months = list(tmp['T'].dt.month.values)
            tmp['T'] = months
            sim_.append(tmp)
            pred_rain.append(100 * predictant.sel(T=year) / clim_mean)
        
        sim__ = xr.concat(sim_, dim="year").assign_coords(year=('year', similar_years)).mean(dim="year", skipna=True)
        sim_all.append(sim__)
        sim_all = xr.concat(sim_all, dim="output").assign_coords(output=('output', ['forecast year', 'composite analog']))
        
        pred_rain = xr.concat(pred_rain, dim='T')
        
        if plot_predictor:
            sim_all.plot(
                x="X", y="Y", row="T", col="output",
                figsize=(12, len(months) * 4),
                cbar_kwargs={"shrink": 0.3, "aspect": 50, "pad": 0.05, "label": f"{variable} anomaly", "orientation": "horizontal"},
                robust=True,
                subplot_kws={'projection': ccrs.PlateCarree()}
            )
            for ax in plt.gcf().axes:
                if isinstance(ax, plt.Axes) and hasattr(ax, 'coastlines'):
                    ax.coastlines()
                    ax.gridlines(draw_labels=True)
                    ax.add_feature(cfeature.LAND, edgecolor="black")
                    ax.add_feature(cfeature.OCEAN, facecolor="lightblue")
            plt.suptitle(f"SST Composites for {reference_year}", fontsize=14)
            plt.tight_layout()
            plt.show()
        else:
            colors_list = ['#d7191c', '#fdae61', '#ffffbf', '#a6d96a', '#1a9641']
            bounds = [0, 50, 90, 110, 150, 200]
            cmap = ListedColormap(colors_list)
            norm = BoundaryNorm(bounds, cmap.N)
            
            data_var = pred_rain.sel(Y=slice(None, 19.5))
            n_times = len(data_var['T'])
            n_cols = 2
            n_rows = (n_times + n_cols - 1) // n_cols
            
            fig, axes = plt.subplots(
                n_rows, n_cols, figsize=(n_cols * 6, n_rows * 4),
                subplot_kw={'projection': ccrs.PlateCarree()}
            )
            axes = np.ravel(axes)
            
            for i, t in enumerate(data_var['T'].values):
                ax = axes[i]
                data = data_var.sel(T=t).squeeze()
                im = ax.pcolormesh(
                    data['X'], data['Y'], data, cmap=cmap, norm=norm,
                    transform=ccrs.PlateCarree()
                )
                ax.coastlines()
                ax.gridlines(draw_labels=False)
                ax.add_feature(cfeature.LAND, edgecolor="black")
                ax.add_feature(cfeature.OCEAN, facecolor="lightblue")
                ax.set_title(f"Season: {str(t)[:10]}")
            
            for j in range(n_times, len(axes)):
                fig.delaxes(axes[j])
            
            cbar = fig.colorbar(im, ax=axes[:n_times], orientation="horizontal", fraction=0.05, pad=0.1)
            cbar.set_label('Precipitation Ratio to Normal (%)')
            cbar.set_ticks([0, 50, 90, 110, 150, 200])
            cbar.ax.set_xticklabels(['0', '50', '90', '110', '150', '200'])
            fig.suptitle("Ratio to Normal Precipitation [%]", fontsize=14)
            plt.tight_layout()
            fig.subplots_adjust(top=0.9, bottom=0.15)
            plt.show()        
        return similar_years
########################################################################################################################################################################################################################################

# Core Python Libraries
import os
import requests  # Added for NCEI download
from pathlib import Path
import calendar
from datetime import datetime
from dateutil.relativedelta import relativedelta
from collections import defaultdict

# Numerical and Data Manipulation Libraries
import numpy as np
import pandas as pd
import xarray as xr
import dask.array as da

# Visualization Libraries
import matplotlib.pyplot as plt
import seaborn as sns
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap, BoundaryNorm

# Climate Data API
import cdsapi

# Machine Learning and Statistical Analysis Libraries
from minisom import MiniSom
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.cluster import KMeans, AgglomerativeClustering
from scipy.stats import pearsonr, gamma, lognorm, stats, norm, expon, weibull_min, t, poisson, nbinom
from scipy.optimize import fsolve
from scipy.special import gamma as gamma_function
import scipy.signal as sig

# EOF Analysis and Verification Libraries
import xeofs as xe
import xskillscore as xs
from wass2s.utils import *
from wass2s.was_bias_correction import *

class WAS_Analog__:
    """Analog-based forecasting toolkit for seasonal climate applications.

    This class orchestrates the end-to-end workflow required to build analog ensembles
    for Sea-Surface Temperature (SST) predictors and to translate them into deterministic
    and probabilistic rainfall forecasts over West Africa or any user-defined domain.
    Supports five analog-selection strategies: Self-Organizing Maps (SOM),
    correlation ranking, Principal-Component (EOF) similarity, K-Means clustering,
    and Agglomerative clustering.

    Parameters
    ----------
    dir_to_save : str
        Directory path to save downloaded and processed data files.
    year_start : int
        Starting year for historical data.
    year_forecast : int
        Target forecast year.
    predictor_vars : list of dict, optional
        List of dictionaries specifying predictor variables, each containing
        'reanalysis_name', 'model_name', 'variable', and 'area'.
        Default is a list with NOAA SST and ERA5 SP, VGRD_850 variables.
    method_analog : str, optional
        Analog method to use ("som", "cor_based", "bias_based, "pca_based", "kmeans_based", "agglomerative_based"). Default is "som".
    month_of_initialization : int, optional
        Month of initialization for forecasts (1-12). If None, uses previous month.
    lead_time : list, optional
        List of lead times in months. If None, defaults to [1, 2, 3, 4, 5].
    ensemble_mean : str, optional
        Method for ensemble mean ("mean" or "median"). Default is "mean".
    rolling : int, optional
        Window size for rolling mean. Default is 3.
    standardize : bool, optional
        If True, standardize data; otherwise, compute anomalies. Default is True.
    multivariateEOF : bool, optional
        If True, use multivariate EOF analysis. Default is False.
    eof_explained_var : float, optional
        Fraction of variance to explain with EOF modes. Default is 0.95.
    clim_year_start : int, optional
        Start year for climatology period.
    clim_year_end : int, optional
        End year for climatology period.
    index_compute : list, optional
        List of climate indices to compute (e.g., ['NINO34', 'DMI']).
    some_grid_size : tuple, optional
        Grid size for SOM (rows, columns). Default is (None, None) for automatic sizing.
    some_learning_rate : float, optional
        Learning rate for SOM training. Default is 0.5.
    radius : float, optional
        Neighborhood radius for SOM analog search. Default is 1.0.
    some_neighborhood_function : str, optional
        Neighborhood function for SOM ("gaussian", "mexican_hat"). Default is "gaussian".
    some_sigma : float, optional
        Initial neighborhood radius for SOM. Default is 1.0.
    some_num_iteration : int, optional
        Number of iterations for SOM training. Default is 2000.
    dist_method : str, optional
        Method for probability calculation ("gamma", "t", "normal", "lognormal", "nonparam").
        Default is "gamma".
    n_clusters : int, optional
        Number of clusters for K-Means and Agglomerative clustering. Default is 4.
    affinity : str, optional
        Distance metric for Agglomerative clustering (e.g., 'euclidean', 'manhattan'). Default is 'euclidean'.
    linkage : str, optional
        Linkage criterion for Agglomerative clustering (e.g., 'ward', 'complete', 'average'). Default is 'ward'.
    """

    def __init__(self, dir_to_save, year_start, year_forecast,
                 predictor_vars=[{'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'SST', 'area': [60, -180, -60, 180]},
                                 {'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'SP', 'area': [60, -180, -60, 180]},
                                 {'reanalysis_name': 'ERA5', 'model_name': 'ECMWF_51', 'variable': 'VGRD_850', 'area': [60, -180, -60, 180]}],
                 method_analog="som", month_of_initialization=None,
                 lead_time=None, ensemble_mean="mean", rolling=3, standardize=True, detrend=False, multivariateEOF=False,
                 eof_explained_var=0.95, clim_year_start=None, clim_year_end=None,
                 index_compute=None, some_grid_size=(None, None), some_learning_rate=0.5, radius=1.0,
                 some_neighborhood_function='gaussian', some_sigma=1.0, some_num_iteration=2000, dist_method="gamma",
                 n_clusters=4, affinity='euclidean', linkage='ward'):
        
        self.dir_to_save = dir_to_save
        self.year_start = year_start
        self.year_forecast = year_forecast
        self.predictor_vars = predictor_vars
        self.method_analog = method_analog
        self.month_of_initialization = month_of_initialization
        self.lead_time = lead_time
        self.ensemble_mean = ensemble_mean
        self.eof_explained_var = eof_explained_var
        self.multivariateEOF = multivariateEOF
        self.rolling = rolling
        self.standardize = standardize
        self.detrend = False
        self.clim_year_start = clim_year_start
        self.clim_year_end = clim_year_end
        self.index_compute = index_compute
        self.some_grid_size = some_grid_size
        self.some_learning_rate = some_learning_rate
        self.radius = radius
        self.some_neighborhood_function = some_neighborhood_function
        self.some_sigma = some_sigma
        self.some_num_iteration = some_num_iteration
        self.dist_method = dist_method
        self.n_clusters = n_clusters
        self.affinity = affinity
        self.linkage = linkage

    def calc_index(self, indices, sst):
        """Calculate climate indices from SST data."""
        sst_indices_name = {
            "NINO34": ("Nino3.4", -170, -120, -5, 5),
            "NINO12": ("Niño1+2", -90, -80, -10, 0),
            "NINO3": ("Nino3", -150, -90, -5, 5),
            "NINO4": ("Nino4", -150, 160, -5, 5),
            "NINO_Global": ("ALL NINO Zone", -80, 160, -10, 5),
            "TNA": ("Tropical Northern Atlantic Index", -55, -15, 5, 25),
            "TSA": ("Tropical Southern Atlantic Index", -30, 10, -20, 0),
            "NAT": ("North Atlantic Tropical", -40, -20, 5, 20),
            "SAT": ("South Atlantic Tropical", -15, 5, -20, 5),
            "TASI": ("NAT-SAT", None, None, None, None),
            "WTIO": ("Western Tropical Indian Ocean (WTIO)", 50, 70, -10, 10),
            "SETIO": ("Southeastern Tropical Indian Ocean (SETIO)", 90, 110, -10, 0),
            "DMI": ("WTIO - SETIO", None, None, None, None),
            "MB": ("Mediterranean Basin", 0, 50, 30, 42),
            "M1": ("M1", -50, 5, -50, -25),
            "M2": ("M2", -75, -10, 25, 50),
            "M3": ("M3", -175, -125, 25, 50),
            "M4": ("M4", -175, -125, -50, -25),
        }
        
        predictor = {}
        for idx in sst_indices_name:
            if idx in ["TASI", "DMI"]:
                continue
            _, lon_min, lon_max, lat_min, lat_max = sst_indices_name[idx]
            sst_region = sst.sel(X=slice(lon_min, lon_max), Y=slice(lat_min, lat_max)).mean(dim=["X", "Y"], skipna=True)
            predictor[idx] = sst_region

        predictor["TASI"] = predictor["NAT"] - predictor["SAT"]
        predictor["DMI"] = predictor["WTIO"] - predictor["SETIO"]
        
        selected_indices = {i: predictor[i] for i in indices if i in predictor}
        data_vars = {key: ds.rename(key) for key, ds in selected_indices.items()}
        return xr.Dataset(data_vars)

    def _postprocess_ersst(self, ds, var_name):
        """Post-process ERSST dataset."""
        ds = ds.drop_vars('zlev', errors='ignore').squeeze()
        keep_vars = [var_name, 'T', 'X', 'Y']
        drop_vars = [v for v in ds.variables if v not in keep_vars]
        return ds.drop_vars(drop_vars, errors="ignore")

    def download_reanalysis(self, force_download=False):
        """
        Download reanalysis data (ERA5 and NOAA ERSST).
        Saves raw MONTHLY data (1-12) without aggregation, preserving seasonality.
        """
        year_end = self.year_forecast
        
        # We download ALL months to preserve the full seasonal cycle for the analog search
        months_int = list(range(1, 13))
        season_str = "JanDec" 
        
        variables = [item['variable'] for item in self.predictor_vars]
        centers = [item['reanalysis_name'] for item in self.predictor_vars]
        areas = [item['area'] for item in self.predictor_vars]
        
        # Variable Mappings
        variables_1 = {
            "PRCP": "total_precipitation", "TEMP": "2m_temperature",
            "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
            "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
            "UGRD10": "10m_u_component_of_wind", "VGRD10": "10m_v_component_of_wind",
            "SST": "sea_surface_temperature", "SLP": "mean_sea_level_pressure",
            "DSWR": "surface_solar_radiation_downwards", "DLWR": "surface_thermal_radiation_downwards",
            "NOLR": "top_net_thermal_radiation",
        }
        variables_2 = {
            "HUSS_1000": "specific_humidity", "HUSS_925": "specific_humidity", "HUSS_850": "specific_humidity",
            "UGRD_1000": "u_component_of_wind", "UGRD_925": "u_component_of_wind", "UGRD_850": "u_component_of_wind",
            "VGRD_1000": "v_component_of_wind", "VGRD_925": "v_component_of_wind", "VGRD_850": "v_component_of_wind",
        }
    
        dir_to_save = Path(self.dir_to_save)
        dir_to_save.mkdir(parents=True, exist_ok=True)
        store_file_path = {}
    
        for center, var, area in zip(centers, variables, areas):
            # Filename includes "JanDec" to indicate full monthly data
            combined_output_path = dir_to_save / f"{center}_{var}_{self.year_start}_{year_end}_{season_str}_{area[0]}{area[1]}{area[2]}{area[3]}.nc" 
    
            if not force_download and combined_output_path.exists():
                print(f"{combined_output_path} exists. Skipping download.")
                store_file_path[var] = xr.open_dataset(combined_output_path)
                continue
    
            # =================================================================
            # CASE 1: NOAA ERSST (Direct NCEI Download)
            # =================================================================
            if f"{center}.{var}" == "NOAA.SST":
                try:
                    print(f"Starting NCEI Direct Download for NOAA ERSST: {self.year_start}-{year_end}")
                    cache_dir = dir_to_save / "ersst_cache"
                    cache_dir.mkdir(exist_ok=True)
                    
                    downloaded_paths = []
                    # v5 is standard. Change to v6 if required.
                    base_url = "https://www.ncei.noaa.gov/data/sea-surface-temperature-extended-reconstructed/v5/access/" 
    
                    for s_year in range(self.year_start, year_end + 1):
                        for month in months_int:
                            fname = f"ersst.v5.{s_year}{month:02d}.nc"
                            local_path = cache_dir / fname
                            
                            if not local_path.exists():
                                url = f"{base_url}/{fname}"
                                r = requests.get(url, timeout=30)
                                r.raise_for_status()
                                with open(local_path, 'wb') as f:
                                    f.write(r.content)
                            downloaded_paths.append(local_path)
    
                    # Process ERSST
                    ds = xr.open_mfdataset(downloaded_paths, combine='by_coords')
                    ds = ds[["sst"]].drop_vars("lev", errors="ignore").squeeze()
    
                    # Standardize Coordinates (0-360 -> -180/180)
                    if 'lon' in ds.coords:
                        ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
                        ds = ds.sortby('lon')
                        
                    # Rename to standard (X, Y, T, SST)
                    rename_dict = {}
                    if 'lat' in ds.coords: rename_dict['lat'] = 'Y'
                    if 'lon' in ds.coords: rename_dict['lon'] = 'X'
                    if 'time' in ds.coords: rename_dict['time'] = 'T'
                    if 'sst' in ds.variables: rename_dict['sst'] = 'SST'
                    ds = ds.rename(rename_dict)
                        
                    # Slice Area [N, W, S, E]
                    ds = ds.sortby(['X', 'Y'])
                    ds = ds.sel(X=slice(area[1], area[3]), Y=slice(area[2], area[0]))
    
                    # NO AGGREGATION: Just ensure time is sorted
                    ds = ds.sortby('T')
                        
                    # Final Cleanup
                    if 'SST' in ds.variables:
                        ds = ds.rename({"SST": "sst"}) 
                        
                    ds.to_netcdf(combined_output_path)
                    store_file_path[var] = ds
                    print(f"Saved NOAA ERSST data to {combined_output_path}")
                    continue
                except Exception as e:
                    print(f"Failed to download NOAA.SST via NCEI: {e}")
                    import traceback
                    traceback.print_exc()
                    continue
    
            # =================================================================
            # CASE 2: ERA5 / CDSAPI
            # =================================================================
            combined_datasets = []
            client = cdsapi.Client()
            
            is_pressure = var in variables_2
            dataset_name = "reanalysis-era5-pressure-levels-monthly-means" if is_pressure else "reanalysis-era5-single-levels-monthly-means"
            api_var = variables_2[var] if is_pressure else variables_1.get(var, var)
            press_level = var.split("_")[1] if is_pressure else None
    
            for year in range(self.year_start, year_end + 1):
                yearly_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                
                if not force_download and yearly_file_path.exists():
                    try:
                        ds = xr.open_dataset(yearly_file_path).load()
                        if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                        combined_datasets.append(ds)
                        continue
                    except Exception as e:
                        print(f"Failed to load {yearly_file_path}: {e}")
                
                try:
                    request = {
                        "product_type": "monthly_averaged_reanalysis",
                        "variable": api_var,
                        "year": str(year),
                        "month": [f"{m:02d}" for m in months_int],
                        "time": "00:00",
                        "area": area,
                        "format": "netcdf",
                    }
                    if is_pressure:
                        request["pressure_level"] = press_level
    
                    print(f"Downloading {var} for {year} from {center}...")
                    client.retrieve(dataset_name, request).download(str(yearly_file_path))
                    
                    with xr.open_dataset(yearly_file_path) as ds:
                        if 'latitude' in ds.coords: ds = ds.rename({"latitude": "Y", "longitude": "X", "valid_time": "T"})
                        ds = ds.load()
                        combined_datasets.append(ds)
                except Exception as e:
                    print(f"Failed to download/process {var} for {year}: {e}")
                    continue
            
            if combined_datasets:
                print(f"Concatenating {var} datasets...")
                combined_ds = xr.concat(combined_datasets, dim="T")
                combined_ds = combined_ds.drop_vars(["number", "expver"], errors="ignore").squeeze()
                
                # Unit Conversions
                if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                    combined_ds = combined_ds - 273.15
                elif var == "PRCP":
                    combined_ds = combined_ds * 1000 
                elif var in ["DSWR", "DLWR", "NOLR"]:
                    combined_ds = combined_ds / 86400
                elif var == "SLP":
                    combined_ds = combined_ds / 100
                
                combined_ds = combined_ds.isel(Y=slice(None, None, -1))
                combined_ds = combined_ds.sortby('T')
    
                store_file_path[var] = combined_ds
                combined_ds.to_netcdf(combined_output_path)
                print(f"Saved combined dataset to {combined_output_path}")
                
                # Cleanup yearly files
                for year in range(self.year_start, year_end + 1):
                    single_file_path = dir_to_save / f"{center}_{var}_{year}.nc"
                    if single_file_path.exists():
                        try: os.remove(single_file_path)
                        except: pass
            else:
                print(f"No data combined for {var}. Check download success.")
    
        return store_file_path

    def download_models(self, force_download=False):
        """Download and process seasonal forecast and hindcast data."""
        centre_map = {
            "BOM_2": "bom", "ECMWF_51": "ecmwf", "UKMO_604": "ukmo", "UKMO_603": "ukmo",
            "METEOFRANCE_8": "meteo_france", "METEOFRANCE_9": "meteo_france",
            "DWD_21": "dwd", "DWD_22": "dwd", "CMCC_35": "cmcc",
            "NCEP_2": "ncep", "JMA_3": "jma", "ECCC_4": "eccc", "ECCC_5": "eccc"
        }
        system_map = {
            "BOM_2": "2", "ECMWF_51": "51", "UKMO_604": "604", "UKMO_603": "603",
            "METEOFRANCE_8": "8", "METEOFRANCE_9": "9", "DWD_21": "21", "DWD_22": "22",
            "CMCC_35": "35", "NCEP_2": "2", "JMA_3": "3", "ECCC_4": "4", "ECCC_5": "5"
        }
        variables_map = {
            "PRCP": "total_precipitation", "TEMP": "2m_temperature",
            "TMAX": "maximum_2m_temperature_in_the_last_24_hours",
            "TMIN": "minimum_2m_temperature_in_the_last_24_hours",
            "UGRD10": "10m_u_component_of_wind", "VGRD10": "10m_v_component_of_wind",
            "SST": "sea_surface_temperature", "SLP": "mean_sea_level_pressure",
            "DSWR": "surface_solar_radiation_downwards",
            "DLWR": "surface_thermal_radiation_downwards",
            "NOLR": "top_thermal_radiation",
            "HUSS_1000": "specific_humidity", "HUSS_925": "specific_humidity",
            "HUSS_850": "specific_humidity", "UGRD_1000": "u_component_of_wind",
            "UGRD_925": "u_component_of_wind", "UGRD_850": "u_component_of_wind",
            "VGRD_1000": "v_component_of_wind", "VGRD_925": "v_component_of_wind",
            "VGRD_850": "v_component_of_wind",
        }

        centers = [item['model_name'] for item in self.predictor_vars]
        variables = [item['variable'] for item in self.predictor_vars]
        areas = [item['area'] for item in self.predictor_vars]
        selected_centres = [centre_map[k] for k in centers]
        selected_systems = [system_map[k] for k in centers]

        month_of_initialization = (datetime.now() - relativedelta(months=1)).month if self.month_of_initialization is None else self.month_of_initialization
        lead_time = [1, 2, 3, 4, 5] if self.lead_time is None else self.lead_time
        if not isinstance(lead_time, list) or any(l < 1 or l > 12 for l in lead_time):
            raise ValueError("lead_time must be a list of integers between 1 and 12.")
        year_forecast = datetime.now().year if self.year_forecast is None else self.year_forecast
        hindcast_years = [str(year) for year in range(self.clim_year_start, self.clim_year_end + 1)]

        dir_to_save = Path(self.dir_to_save)
        dir_to_save.mkdir(parents=True, exist_ok=True)
        abb_mont_ini = calendar.month_abbr[int(month_of_initialization)]
        season_months = [((month_of_initialization + l - 1) % 12) + 1 for l in lead_time]
        season_str = "".join(calendar.month_abbr[m] for m in season_months)

        store_file_path = {}
        store_hdcst_file_path = {}

        for cent, syst, var, area in zip(selected_centres, selected_systems, variables, areas):
            forecast_file = dir_to_save / f"forecast_{cent}{syst}_{var}_{year_forecast}_{abb_mont_ini}Ic_{season_str}_{lead_time[0]}_{area[0]}{area[1]}{area[2]}{area[3]}.nc"
            hindcast_file = dir_to_save / f"hindcast_{cent}{syst}_{var}_{self.clim_year_start}_{self.clim_year_end}{abb_mont_ini}Ic_{season_str}_{lead_time[0]}_{area[0]}{area[1]}{area[2]}{area[3]}.nc"
                
            if not force_download and forecast_file.exists():
                print(f"Forecast file {forecast_file} exists. Skipping download.")
                ds = xr.open_dataset(forecast_file)
                store_file_path[var] = ds
                ds.close()
            else:
                try:
                    dataset = "seasonal-monthly-pressure-levels" if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]) else "seasonal-monthly-single-levels"
                    forecast_request = {
                        "originating_centre": cent,
                        "system": syst,
                        "variable": variables_map[var],
                        "product_type": ["monthly_mean"],
                        "year": [str(year_forecast)],
                        "month": [f"{month_of_initialization:02d}"],
                        "leadtime_month": lead_time,
                        "data_format": "netcdf",
                        "area": area,
                    }
                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        forecast_request["pressure_level"] = var.split("_")[1]

                    client = cdsapi.Client()
                    client.retrieve(dataset, forecast_request).download(str(forecast_file))
                    print(f"Downloaded forecast: {forecast_file}")

                    ds_forecast = xr.open_dataset(forecast_file)

                    if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                        ds_forecast = ds_forecast - 273.15
                    elif var == "PRCP":
                        ds_forecast = ds_forecast * (1000 * 30 * 24 * 3600)
                    elif var == "SLP":
                        ds_forecast = ds_forecast / 100
                    elif var in ["DSWR", "DLWR", "NOLR"]:
                        ds_forecast = ds_forecast / 86400

                    if self.ensemble_mean and "number" in ds_forecast.dims:
                        ds_forecast = getattr(ds_forecast, self.ensemble_mean)(dim="number", skipna=True)

                    ds_forecast = ds_forecast.isel(latitude=slice(None, None, -1))
                    rename_dict = {
                        "latitude": "Y",
                        "longitude": "X",
                        "time": "T" if "time" in ds_forecast.coords else None,
                        "indexing_time": "T" if "indexing_time" in ds_forecast.coords else None,
                        "forecast_reference_time": "T" if "forecast_reference_time" in ds_forecast.coords else None,
                    }
                    rename_dict = {k: v for k, v in rename_dict.items() if v is not None}
                    ds_forecast = ds_forecast.rename(rename_dict)


                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        ds_forecast = ds_forecast.drop_vars("pressure_level", errors="ignore").squeeze()
                    
    
                    ds_forecast.close()
                    os.remove(str(forecast_file))
                    ds_forecast.to_netcdf(str(forecast_file))
                    print(f"Saved processed hindcast to {forecast_file}")
                    store_file_path[var] = ds_forecast

                except Exception as e:
                    print(f"Failed to process {var} for {cent}{syst}: {e}")
                    continue

            if not force_download and hindcast_file.exists():
                print(f"Hindcast file {hindcast_file} exists. Skipping download.")
                ds = xr.open_dataset(hindcast_file)
                store_hdcst_file_path[var] = ds
                ds.close()
            else:
                try:
                    dataset = "seasonal-monthly-pressure-levels" if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]) else "seasonal-monthly-single-levels"
                    hindcast_request = {
                        "originating_centre": cent,
                        "system": syst,
                        "variable": variables_map[var],
                        "product_type": ["monthly_mean"],
                        "year": hindcast_years,
                        "month": [f"{month_of_initialization:02d}"],
                        "leadtime_month": lead_time,
                        "data_format": "netcdf",
                        "area": area,
                    }
                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        hindcast_request["pressure_level"] = var.split("_")[1]
        
                    client = cdsapi.Client()
                    client.retrieve(dataset, hindcast_request).download(str(hindcast_file))
                    print(f"Downloaded hindcast: {hindcast_file}")

                    ds_hindcast = xr.open_dataset(hindcast_file)
                    if var in ["TMIN", "TEMP", "TMAX", "SST"]:
                        ds_hindcast = ds_hindcast - 273.15
                    elif var == "PRCP":
                        ds_hindcast = ds_hindcast * (1000 * 30 * 24 * 3600)
                    elif var == "SLP":
                        ds_hindcast = ds_hindcast / 100
                    elif var in ["DSWR", "DLWR", "NOLR"]:
                        ds_hindcast = ds_hindcast / 86400

                    if self.ensemble_mean and "number" in ds_hindcast.dims:
                        ds_hindcast = getattr(ds_hindcast, self.ensemble_mean)(dim="number", skipna=True)                        

                    ds_hindcast = ds_hindcast.isel(latitude=slice(None, None, -1))
                    rename_dict = {
                        "latitude": "Y",
                        "longitude": "X",
                        "time": "T" if "time" in ds_hindcast.coords else None,
                        "indexing_time": "T" if "indexing_time" in ds_hindcast.coords else None,
                        "forecast_reference_time": "T" if "forecast_reference_time" in ds_hindcast.coords else None,
                    }
                    rename_dict = {k: v for k, v in rename_dict.items() if v is not None}
                    ds_hindcast = ds_hindcast.rename(rename_dict)

                    if var in variables_map and any(v in var for v in ["HUSS", "UGRD", "VGRD"]):
                        ds_hindcast = ds_hindcast.drop_vars("pressure_level", errors="ignore").squeeze()

                    ds_hindcast = ds_hindcast.sortby("T")
                    ds_hindcast.close()
                    os.remove(hindcast_file)
                    ds_hindcast.to_netcdf(hindcast_file)
                    print(f"Saved processed hindcast to {hindcast_file}")
                    store_hdcst_file_path[var] = ds_hindcast


                except Exception as e:
                    print(f"Failed to process {var} for {cent}{syst}: {e}")
                    continue

        return store_hdcst_file_path, store_file_path

    def anomaly_timeseries(self, ds):
        """Compute anomalies by removing the climatological mean."""
        if self.clim_year_start and self.clim_year_end:
            clim_period = ds.sel(T=slice(f"{self.clim_year_start}-01-01", f"{self.clim_year_end}-12-31"))
        else:
            clim_period = ds

        clim_mean = clim_period.groupby("T.month").mean("T", skipna=True)
        anomaly_slices = []

        for month in range(1, 13):
            month_mask = ds["T"].dt.month == month
            ds_month = ds.where(month_mask, drop=True)
            month_mean = clim_mean.sel(month=month)
            ds_anomaly_month = ds_month - month_mean
            anomaly_slices.append(ds_anomaly_month)

        ds_anomaly = xr.concat(anomaly_slices, dim="T")
        ds_anomaly = ds_anomaly.sortby("T")
        if "month" in ds_anomaly.coords:
            ds_anomaly = ds_anomaly.drop_vars("month")
        if "month" in ds_anomaly.dims:
            ds_anomaly = ds_anomaly.drop_dims("month")
        return ds_anomaly

    def standardize_timeseries(self, ds):
        """Standardize the dataset by removing climatological mean and scaling by standard deviation."""
        if self.clim_year_start and self.clim_year_end:
            clim_period = ds.sel(T=slice(f"{self.clim_year_start}-01-01", f"{self.clim_year_end}-12-31"))
        else:
            clim_period = ds

        clim_mean = clim_period.groupby("T.month").mean("T", skipna=True)
        clim_std = clim_period.groupby("T.month").std("T", skipna=True)
        clim_std = clim_std.where(clim_std != 0, 1e-10)

        standardized_slices = []

        for month in range(1, 13):
            month_mask = ds["T"].dt.month == month
            ds_month = ds.where(month_mask, drop=True)
            month_mean = clim_mean.sel(month=month)
            month_std = clim_std.sel(month=month)
            ds_standardized_month = (ds_month - month_mean) / month_std
            standardized_slices.append(ds_standardized_month)

        ds_standardized = xr.concat(standardized_slices, dim="T")
        ds_standardized = ds_standardized.sortby("T")
        if "month" in ds_standardized.coords:
            ds_standardized = ds_standardized.drop_vars("month")
        if "month" in ds_standardized.dims:
            ds_standardized = ds_standardized.drop_dims("month")
        return ds_standardized

    def compute_eofs(self, data):
        """Compute EOFs and retain modes explaining at least the specified variance."""
        data = data.rename({"X": "lon", "Y": "lat"})
        data = data.fillna(data.mean(dim="T", skipna=True))
        print("Computing EOFs...")
        model = xe.single.EOF(n_modes=100, use_coslat=True, center=False)
        model.fit(data, dim='T')
        explained_var = model.explained_variance_ratio()
        cumulative_var = explained_var.cumsum()
        n_modes = np.where(cumulative_var >= self.eof_explained_var)[0][0] + 1
        print(f"Selected {n_modes} modes explaining {cumulative_var[n_modes-1].values*100:.2f}% variance")
        return model.scores().isel(mode=slice(0, n_modes))

    def _detrended_da(self, da, dim='T'):
            """
            Detrend a DataArray by removing the linear trend.
            """
            if dim not in da.dims:
                raise ValueError(f"DataArray must have a time dimension '{dim}' for detrending.")
    
            # 1. Calculate coefficients (Slope + Intercept)
            # Returns a Dataset containing 'polyfit_coefficients'
            coeffs = da.polyfit(dim=dim, deg=1)
    
            # 2. Reconstruct the trend line using standard xarray syntax
            trend = xr.polyval(da[dim], coeffs.polyfit_coefficients)
    
            # 3. Return anomalies and the FULL coefficients 
            # (We need the full coefficients to detrend future data for prediction)
            return da - trend, coeffs

    def _prepare_data_for_clustering(self, unique_years, target_year=None):
        """Helper to prepare data (indices or EOFs) for clustering methods."""
        _, ddd = self.download_and_process()
        data_scaled = None
        index_data = None

        if self.index_compute:
            ddd_sst = ddd.get('SST')
            if self.detrend:
                ddd_sst, _ = self._detrended_da(ddd_sst)
            if ddd_sst is None:
                raise ValueError("SST data not found in downloaded datasets for clustering.")
            indices_dataset = self.calc_index(self.index_compute, ddd_sst)
            indices_dataset_sel_years = indices_dataset.where(indices_dataset['T'].dt.year.isin(unique_years), drop=True)
            data_scaled, index_data = self.arrange_indices_for_som(indices_dataset_sel_years)
        else:
            if self.multivariateEOF:
                data_vars = [self._detrended_da(ddd[var].rename({"X": "lon", "Y": "lat"}))[0] if self.detrend else ddd[var].rename({"X": "lon", "Y": "lat"}) for var in ddd]
                model = xe.single.EOF(n_modes=100, use_coslat=True, center=False)
                model.fit(data_vars, dim='T')
                explained_var = model.explained_variance_ratio()
                n_modes = np.where(explained_var.cumsum() >= self.eof_explained_var)[0][0] + 1
                scores = model.scores().isel(mode=slice(0, n_modes))
                scores = scores.where(scores['T'].dt.year.isin(unique_years), drop=True)
                
                df_final = pd.DataFrame()
                scores = scores.assign_coords(year=('T', scores['T'].dt.year.data), month_name=('T', scores['T'].dt.strftime('%b').data))
                df = scores.to_dataframe(name='value').reset_index()
                df['mode_month'] = df['mode'].apply(lambda m: f"mode{m:02d}") + "_" + df['month_name']
                df_pivot = df.pivot(index='year', columns='mode_month', values='value')
                month_order = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
                df_pivot = df_pivot.reindex(columns=sorted(df_pivot.columns, key=lambda x: (month_order.index(x.split('_')[1]), x)))
                df_final = df_pivot
                data_scaled = np.array(df_final, dtype=np.float64)
                index_data = df_final.index
            else:
                df_final = pd.DataFrame()
                for var in ddd:
                    if self.detrend:
                        ddd[var], _ = self._detrended_da(ddd[var])
                    scores = self.compute_eofs(ddd[var])
                    scores = scores.where(scores['T'].dt.year.isin(unique_years), drop=True)
                    scores = scores.assign_coords(year=('T', scores['T'].dt.year.data), month_name=('T', scores['T'].dt.strftime('%b').data))
                    df = scores.to_dataframe(name='value').reset_index()
                    df['mode_month'] = df['mode'].apply(lambda m: f"mode{m:02d}") + "_" + df['month_name'] + f"_{var}"
                    df_pivot = df.pivot(index='year', columns='mode_month', values='value')
                    month_order = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
                    df_pivot = df_pivot.reindex(columns=sorted(df_pivot.columns, key=lambda x: (month_order.index(x.split('_')[1]), x)))
                    df_final = pd.concat([df_final, df_pivot], axis=1)
                data_scaled = np.array(df_final, dtype=np.float64)
                index_data = df_final.index
        
        if np.isnan(data_scaled).any():
            print("NaN values found in data for clustering. Imputing with column mean.")
            col_means = np.nanmean(data_scaled, axis=0)
            nan_indices = np.where(np.isnan(data_scaled))
            data_scaled[nan_indices] = col_means[nan_indices[1]]
            
        return data_scaled, index_data

    def arrange_indices_for_som(self, ds):
        """Arrange climate indices for SOM training."""
        ds = ds.assign(year=('T', ds['T'].dt.year.data), month_name=('T', ds['T'].dt.strftime('%b').data))
        
        def pivot_var(var_name):
            df = ds[[var_name, 'year', 'month_name']].to_dataframe().reset_index()
            df_pivot = df.pivot(index='year', columns='month_name', values=var_name)
            df_pivot = df_pivot.reindex(columns=['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'])
            df_pivot.columns = [f"{var_name}_{month}" for month in df_pivot.columns]
            return df_pivot
        
        vars_order = self.index_compute or []
        df_vars = {v: pivot_var(v) for v in vars_order}
        months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
        ordered_columns = [f"{v}_{m}" for m in months for v in vars_order]
        
        df_final = pd.concat(df_vars.values(), axis=1)[ordered_columns]
        return np.array(df_final, dtype=np.float64), df_final.index

    def identify_analogs_bmu(self, target_year, bmu_dict):
        """Identify analog years based on SOM Best Matching Units (BMUs)."""
        radius = self.radius or 1.0
        neuron_years = defaultdict(list)
        for year, coords in bmu_dict.items():
            neuron_years[coords].append(year)
        
        target_coords = bmu_dict.get(target_year)
        if not target_coords:
            return []
        
        analogs = [year for year in neuron_years[target_coords] if year != target_year]
        if not analogs:
            for year, coords in bmu_dict.items():
                if year == target_year:
                    continue
                dist = np.sqrt((coords[0] - target_coords[0])**2 + (coords[1] - target_coords[1])**2)
                if dist <= radius:
                    analogs.append(year)
        
        return analogs

    def SOM(self, predictant, itrain, ireference_year):
        """Identify similar years using Self-Organizing Maps (SOM)."""

        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)
        som_grid_size = (4, 4) if set(self.some_grid_size) == {None} else self.some_grid_size
        som = MiniSom(
            x=som_grid_size[0],
            y=som_grid_size[1],
            input_len=data_scaled.shape[1],
            sigma=self.some_sigma,
            learning_rate=self.some_learning_rate,
            neighborhood_function=self.some_neighborhood_function,
            random_seed=42
        )
        som.random_weights_init(data_scaled)
        som.train_random(data=data_scaled, num_iteration=self.some_num_iteration)
        
        bmu_coords = [som.winner(x) for x in data_scaled]
        bmu_dict = {year: coords for year, coords in zip(index_data, bmu_coords)}
        similar_years = self.identify_analogs_bmu(reference_year, bmu_dict)
        print(f"Similar years for {reference_year}: {similar_years}")
        return np.array(similar_years)


    def Bias_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using bias-based analog method."""
        _, ddd = self.download_and_process()
        ddd = ddd.get(self.predictor_vars[0]['variable'])
        if ddd is None:
            raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
        
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
        sst_reference = ddd.sel(T=str(reference_year))
        
        biases = []
        for year in unique_years:
            tmp = ddd.sel(T=str(year))
            tmp['T'] = sst_reference['T'] 
            bias = np.abs(tmp - sst_reference).mean(dim="T")
            biases.append(bias)
        
        similar = xr.concat(biases, dim='T').assign_coords(T=unique_years).compute()
        similar = xr.where(similar < 0.5, 1, 0).sum(dim=["X", "Y"])
        similar = similar.sortby(similar, ascending=False)
        top_2 = similar.isel(T=slice(1, 3))
        similar_years = top_2['T'].to_numpy()
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years
    
    def Corr_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using correlation-based analog method."""

        _, ddd = self.download_and_process()

        ddd = ddd.get(self.predictor_vars[0]['variable'])

        if ddd is None:
            raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
        
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
        sst_reference = ddd.sel(T=str(reference_year))
        
        correlations = []
        for year in unique_years:
            tmp = ddd.sel(T=str(year))
            tmp['T'] = sst_reference['T']
            correlation = xr.corr(tmp, sst_reference, dim="T").compute()
            correlations.append(correlation)
        
        similar = xr.concat(correlations, dim='T').assign_coords(T=unique_years)
        similar = xr.where(similar > 0.3, 1, 0).sum(dim=["X", "Y"])
        similar = similar.sortby(similar, ascending=False)
        top_2 = similar.isel(T=slice(1, 3))
        similar_years = top_2['T'].to_numpy()
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years

    def Pca_Based(self, predictant, itrain, ireference_year):
            """Identify similar years using PCA-based analog method."""
            _, ddd = self.download_and_process()
            ddd = ddd.get(self.predictor_vars[0]['variable'])
            if ddd is None:
                raise ValueError(f"Variable {self.predictor_vars[0]['variable']} not found in downloaded data.")
            
            # predictor_ = ddd.fillna(ddd.groupby("T.month").mean("T", skipna=True))
            # predictor_detrend = sig.detrend(predictor_, axis=0)
            # ddd = xr.DataArray(predictor_detrend, dims=predictor_.dims, coords=predictor_.coords)
            if self.detrend:
                ddd, _ = self._detrended_da(ddd)
            # ddd = ddd.rename({"X": "lon", "Y": "lat"})
            # eof = xe.single.EOF(n_modes=50, use_coslat=True, center=False)
            # eof.fit(ddd.fillna(ddd.mean(dim="T", skipna=True)), dim="T")
            # scores = eof.scores()
            scores = self.compute_eofs(ddd)
            
            predictant = predictant.copy()
            predictant['T'] = predictant['T'].astype('datetime64[ns]')
            if ireference_year==2100:
                reference_year = self.year_forecast
                predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
            else:
                reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
                predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
            unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)
            
            ref_scores = scores.sel(T=slice(f"{reference_year}-01-01", f"{reference_year}-12-31"))
            ref_scores = ref_scores.assign_coords(T=np.arange(ref_scores.sizes['T']))
            sst_ref = ref_scores.stack(score=('mode', 'T'))
            
            correlations = []
            for year in unique_years:
                tmp_scores = scores.sel(T=slice(f"{year}-01-01", f"{year}-12-31"))
                tmp_scores = tmp_scores.assign_coords(T=np.arange(tmp_scores.sizes['T']))
                tmp = tmp_scores.stack(score=('mode', 'T'))
                correlation = xr.corr(tmp, sst_ref, dim="score").compute()
                correlations.append(correlation)
            
            similar = xr.concat(correlations, dim='T').assign_coords(T=unique_years)
            similar = similar.sortby(similar, ascending=False)
            top_2 = similar.isel(T=slice(1,3))
            similar_years = top_2['T'].to_numpy()
            print(f"Similar years for {reference_year}: {similar_years}")
            return similar_years

    def KMeans_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using K-Means clustering."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")

        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)

        kmeans = KMeans(n_clusters=self.n_clusters, random_state=42, n_init='auto')
        cluster_labels = kmeans.fit_predict(data_scaled)

        target_year_idx = np.where(index_data == reference_year)[0]
        if len(target_year_idx) == 0:
            print(f"Target year {reference_year} not found in index data.")
            return np.array([])
        target_cluster_label = cluster_labels[target_year_idx[0]]

        similar_years_indices = np.where(cluster_labels == target_cluster_label)[0]
        similar_years = [index_data[i] for i in similar_years_indices if index_data[i] != reference_year]
        
        print(f"Similar years for {reference_year} (K-Means): {similar_years}")
        return np.array(similar_years)

    def Agglomerative_Based(self, predictant, itrain, ireference_year):
        """Identify similar years using Agglomerative Clustering."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')

        if ireference_year==2100:
            reference_year = self.year_forecast
            predictant_ = xr.concat([predictant.isel(T=itrain)], dim="T")
        else:
            reference_year = int(np.unique(predictant.isel(T=ireference_year)['T'].dt.year))
            predictant_ = xr.concat([predictant.isel(T=itrain), predictant.isel(T=ireference_year)], dim="T")

        unique_years = np.append(np.unique(predictant_['T'].dt.year),reference_year)

        data_scaled, index_data = self._prepare_data_for_clustering(unique_years, reference_year)

        agg_clustering = AgglomerativeClustering(n_clusters=self.n_clusters, metric=self.affinity, linkage=self.linkage)
        cluster_labels = agg_clustering.fit_predict(data_scaled)

        target_year_idx = np.where(index_data == reference_year)[0]
        if len(target_year_idx) == 0:
            print(f"Target year {reference_year} not found in index data.")
            return np.array([])
        target_cluster_label = cluster_labels[target_year_idx[0]]

        similar_years_indices = np.where(cluster_labels == target_cluster_label)[0]
        similar_years = [index_data[i] for i in similar_years_indices if index_data[i] != reference_year]

        print(f"Similar years for {reference_year} (Agglomerative): {similar_years}")
        return np.array(similar_years)

    def download_and_process(self):
            """Download and process reanalysis and forecast data."""
            
            lead_time = [1, 2, 3, 4, 5] if self.lead_time is None else self.lead_time
            month_of_initialization = (datetime.now() - relativedelta(months=1)).month if self.month_of_initialization is None else self.month_of_initialization
            
            sst_hist = self.download_reanalysis(force_download=False)
            sst_hdcst, sst_for = self.download_models(force_download=False)
    
            variables = [item['variable'] for item in self.predictor_vars]
            print(f"Processing variables: {variables}")
            data_var_concatenated = {}
            data_var_shifted = {}
    
            for var in variables:
                if var not in sst_hist or var not in sst_for:
                    print(f"Skipping {var}: data not available.")
                    continue  
    
                sst_hist_ = sst_hist[var]
                sst_hdcst_ = sst_hdcst[var]
                sst_for_ = sst_for[var]
    
                # --- Process Hindcast Data ---
                hindcast = []
                for i in sst_hdcst_['T']:
                    sst_hdcst_i = sst_hdcst_.sel(T=i)
                    base_times = pd.Timestamp(sst_hdcst_i['T'].values)
                    new_times = [base_times + pd.DateOffset(months=int(m)) for m in sst_hdcst_i['forecastMonth'].values]
                    sst_hdcst_i = sst_hdcst_i.assign_coords(forecastMonth=("forecastMonth", new_times))
                    sst_hdcst_i = sst_hdcst_i.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})
                    hindcast.append(sst_hdcst_i)
                
                sst_hdcst_ = xr.concat(hindcast, dim='T')            
                sst_hdcst_ = sst_hdcst_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})
    
                # --- Process Forecast Data ---
                sst_for_ = sst_for_.interp(Y=sst_hist_.Y, X=sst_hist_.X, method="linear", kwargs={"fill_value": "extrapolate"})
                
                if isinstance(sst_for_['T'].values, np.ndarray):
                    base_time = pd.Timestamp(sst_for_['T'].values[-1])
                else:
                    base_time = pd.Timestamp(sst_for_['T'].values)
                    
                new_times = [base_time + pd.DateOffset(months=int(m)) for m in sst_for_['forecastMonth'].values]
                sst_for_ = sst_for_.assign_coords(forecastMonth=("forecastMonth", new_times))
                sst_for_ = sst_for_.drop_vars('T', errors='ignore').squeeze().rename({'forecastMonth': 'T'})
    
                # --- Bias Correction ---
                biasr = WAS_bias_correction()
                sst_for_bias_corrected = []
                
                for m in lead_time:
                    # 1. Calculate target month
                    target_time = base_time + pd.DateOffset(months=m)
                    current_target_month = target_time.month
    
                    # 2. Filtering
                    sst_hist_mean = sst_hist_.where(sst_hist_['T'].dt.month == current_target_month, drop=True)
                    sst_hist_mean = sst_hist_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()
    
                    sst_hdcst_mean = sst_hdcst_.where(sst_hdcst_['T'].dt.month == current_target_month, drop=True)
                    sst_hdcst_mean = sst_hdcst_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                    
                    # 3. Alignment
                    # Create a set of unique years from the hindcast data
                    years_to_keep = sst_hdcst_mean['T'].dt.year.to_numpy()
                    
                    # Filter the historical data based on this set of years
                    sst_hist_mean = sst_hist_mean.where(sst_hist_mean['T'].dt.year.isin(years_to_keep), drop=True)
                    sst_hdcst_mean['T'] = sst_hist_mean['T']
                    sst_hist_aligned, sst_hdcst_aligned  = xr.align(sst_hist_mean, sst_hdcst_mean, join="inner")
                    
                    # 4. CRITICAL SAFETY CHECK: Check if data is empty after alignment
                    if sst_hist_aligned.sizes['T'] == 0 or sst_hdcst_aligned.sizes['T'] == 0:
                        print(f"WARNING: No overlapping data found for month {current_target_month} (lead {m}). Skipping bias correction for this step.")
                        continue # Skip to avoid crash
    
                    sst_for_mean = sst_for_.where(sst_for_['T'].dt.month == current_target_month, drop=True)
                    sst_for_mean = sst_for_mean.to_array().drop_vars(['variable'], errors='ignore').squeeze('variable', drop=True)
                    
                    if var in ["TMIN", "TEMP", "TMAX", "SST", "HUSS_1000", "HUSS_925", "HUSS_850", "SLP", "UGRD10", "VGRD10", "UGRD_1000", "UGRD_925", "UGRD_850", "VGRD_1000", "VGRD_925", "VGRD_850"]:
                        # Use aligned variables
                        fobj_quant_da = biasr.fitBC(sst_hist_aligned, sst_hdcst_aligned, method='QUANT', qstep=0.01, nboot=10)
                        sst_for_bias_corrected.append(biasr.doBC(sst_for_mean, fobj_quant_da, type='linear'))
                
                # --- Concatenation & Post-Processing ---
                if not sst_for_bias_corrected:
                    print(f"ERROR: Bias correction failed for all lead times for variable {var}.")
                    continue
    
                sst_for_ = xr.concat(sst_for_bias_corrected, dim='T').sortby("T")
    
                sst_hist_ = sst_hist_.sel(T=slice(f"{self.year_start}-01", f"{base_time.year}-{base_time.month:02d}"))
                sst_hist_ = sst_hist_.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                
                concatenated_ds = xr.concat([sst_hist_, sst_for_], dim='T')
                concatenated_ds = concatenated_ds.sortby("T")
                concatenated_ds = concatenated_ds.to_dataset(name=var)
    
                if isinstance(self.rolling, int) and self.rolling > 1:
                    print(f"Applied rolling mean with window size {self.rolling} for {var}.")
                    concatenated_ds = concatenated_ds.rolling(T=self.rolling, center=False, min_periods=self.rolling).mean()
    
                if self.standardize:
                    print(f"Standardizing timeseries for {var}.")
                    concatenated_ds_st = self.standardize_timeseries(concatenated_ds)
                else:
                    print(f"Computing anomalies for {var}.")
                    concatenated_ds_st = self.anomaly_timeseries(concatenated_ds)
    
                select_parts = concatenated_ds_st['T'].size - (self.year_forecast - self.year_start) * 12 
                concatenated_ds_st = concatenated_ds_st.isel(T=slice(select_parts, None))
                
                new_time = pd.DatetimeIndex([pd.to_datetime(f"{self.year_start+1}-01-01") + pd.DateOffset(months=t) for t in range(len(concatenated_ds_st['T']))]) 
    
                ds_shifted = concatenated_ds_st.assign_coords(T=new_time)
        
                data_var_shifted[var] = ds_shifted.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                data_var_concatenated[var] = concatenated_ds_st.to_array().drop_vars(['variable'], errors='ignore').squeeze()
                
            return data_var_concatenated, data_var_shifted

    def compute_model(self, predictant, itrain, itest):
        """Compute deterministic hindcast using the specified analog method."""
        method_map = {
            "som": self.SOM,
            "bias_based": self.Bias_Based,
            "cor_based": self.Corr_Based,
            "pca_based": self.Pca_Based,
            "kmeans_based": self.KMeans_Based,
            "agglomerative_based": self.Agglomerative_Based
        }
        if self.method_analog not in method_map:
            raise ValueError(f"Invalid analog method: {self.method_analog}. Choose 'som', 'bias_based', 'cor_based', 'pca_based', 'kmeans_based', or 'agglomerative_based'.")

        similar_years = method_map[self.method_analog](predictant, itrain, itest)
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        
        sim_obs = [predictant.sel(T=str(year)) for year in similar_years if year in predictant['T'].dt.year.values]
        if not sim_obs:
            raise ValueError("No valid similar years found for hindcast.")
        
        hindcast_det = xr.concat(sim_obs, dim="T").mean(dim="T").expand_dims({'T': predictant.isel(T=itest)['T'].values})
        return hindcast_det

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

    @staticmethod
    def _ppf_terciles_from_code(dist_code, shape, loc, scale):
        """
        Return tercile thresholds (T1, T2) 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.3, loc=loc, scale=scale),
                    norm.ppf(0.68, loc=loc, scale=scale),
                )
            elif code == 2:
                return (
                    lognorm.ppf(0.3, s=shape, loc=loc, scale=scale),
                    lognorm.ppf(0.68, s=shape, loc=loc, scale=scale),
                )
            elif code == 3:
                return (
                    expon.ppf(0.3, loc=loc, scale=scale),
                    expon.ppf(0.68, loc=loc, scale=scale),
                )
            elif code == 4:
                return (
                    gamma.ppf(0.3, a=shape, loc=loc, scale=scale),
                    gamma.ppf(0.68, a=shape, loc=loc, scale=scale),
                )
            elif code == 5:
                return (
                    weibull_min.ppf(0.3, c=shape, loc=loc, scale=scale),
                    weibull_min.ppf(0.68, c=shape, loc=loc, scale=scale),
                )
            elif code == 6:
                # Note: Renamed 't_dist' to 't' for standard scipy.stats
                return (
                    t.ppf(0.3, df=shape, loc=loc, scale=scale),
                    t.ppf(0.68, 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.3, mu=shape, loc=loc),
                    poisson.ppf(0.68, 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.33, n=shape, p=scale, loc=loc),
                    nbinom.ppf(0.68, 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.
        """
        
        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)

        # 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=best_guess, 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]
                V = error_variance
                
                # 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(WAS_Analog_.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.33, 0.68], 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, hindcast_det, best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """Generate deterministic and probabilistic forecasts for the target year."""
        predictant = predictant.copy()
        predictant['T'] = predictant['T'].astype('datetime64[ns]')
        reference_year = self.year_forecast
        unique_years = np.append(np.unique(predictant['T'].dt.year), reference_year)
        
        method_map = {
            "som": self.SOM,
            "bias_based": self.Bias_Based,
            "cor_based": self.Corr_Based,
            "pca_based": self.Pca_Based,
            "kmeans_based": self.KMeans_Based,
            "agglomerative_based": self.Agglomerative_Based
        }
        if self.method_analog not in method_map:
            raise ValueError(f"Invalid analog method: {self.method_analog}. Choose 'som','bias_based', 'cor_based', 'pca_based', 'kmeans_based', or 'agglomerative_based'.")

        itrain = [i for i, year in enumerate(predictant['T'].dt.year.values) if year in unique_years and year != reference_year]
        ireference_year = 2100

        similar_years = method_map[self.method_analog](predictant, itrain, ireference_year)
        if len(similar_years)>2:
            similar_years = similar_years[:2]
        else:
            similar_years = similar_years[:len(similar_years)]

        sim_obs = [predictant.sel(T=str(year)) for year in similar_years if year in predictant['T'].dt.year.values]

        if not sim_obs:
            raise ValueError("No valid similar years found for forecast.")
        
        forecast_det = xr.concat(sim_obs, dim="T").mean(dim="T")
        T_value = predictant.isel(T=0).coords['T'].values
        month = T_value.astype('datetime64[M]').astype(int) % 12 + 1
        new_T_value = np.datetime64(f"{reference_year}-{month:02d}-01")
        forecast_det = forecast_det.expand_dims({'T': [new_T_value]})
        forecast_det['T'] = forecast_det['T'].astype('datetime64[ns]')
        
        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.3, 0.68], 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')
        
        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_det,
                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_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}")
        
        forecast_prob = forecast_prob.assign_coords(probability=('probability', ['PB', 'PN', 'PA'])).transpose('probability', 'T', 'Y', 'X')
        print(f"Similar years for {reference_year}: {similar_years}")
        return similar_years, forecast_det, forecast_prob

    def composite_plot(self, predictant, clim_year_start, clim_year_end, hindcast_det, plot_predictor=True, variable="SST",best_code_da=None, best_shape_da=None, best_loc_da=None, best_scale_da=None):
        """Create composite plots of predictors or predictands."""
        clim_slice = slice(str(clim_year_start), str(clim_year_end))
        clim_mean = predictant.sel(T=clim_slice).mean(dim='T', skipna=True)
        similar_years, result_, _ = self.forecast(predictant, clim_year_start, clim_year_end, hindcast_det, best_code_da, best_shape_da, best_loc_da, best_scale_da)
        result_ = result_.drop_vars('T', errors='ignore').squeeze()
        reference_year = self.year_forecast
        _, ddd = self.download_and_process()

        ddd = ddd[variable]
        
        sim_all = []
        tmp = ddd.sel(T=str(reference_year))
        months = list(tmp['T'].dt.month.values)
        tmp['T'] = months
        sim_all.append(tmp)
        
        sim_ = []
        pred_rain = []
        for year in np.array([str(i) for i in similar_years]):
            tmp = ddd.sel(T=year)
            months = list(tmp['T'].dt.month.values)
            tmp['T'] = months
            sim_.append(tmp)
            pred_rain.append(100 * predictant.sel(T=year) / clim_mean)
        
        sim__ = xr.concat(sim_, dim="year").assign_coords(year=('year', similar_years)).mean(dim="year", skipna=True)
        sim_all.append(sim__)
        sim_all = xr.concat(sim_all, dim="output").assign_coords(output=('output', ['forecast year', 'composite analog']))
        
        pred_rain = xr.concat(pred_rain, dim='T')
        
        if plot_predictor:
            sim_all.plot(
                x="X", y="Y", row="T", col="output",
                figsize=(12, len(months) * 4),
                cbar_kwargs={"shrink": 0.3, "aspect": 50, "pad": 0.05, "label": f"{variable} anomaly", "orientation": "horizontal"},
                robust=True,
                subplot_kws={'projection': ccrs.PlateCarree()}
            )
            for ax in plt.gcf().axes:
                if isinstance(ax, plt.Axes) and hasattr(ax, 'coastlines'):
                    ax.coastlines()
                    ax.gridlines(draw_labels=True)
                    ax.add_feature(cfeature.LAND, edgecolor="black")
                    ax.add_feature(cfeature.OCEAN, facecolor="lightblue")
            plt.suptitle(f"SST Composites for {reference_year}", fontsize=14)
            plt.tight_layout()
            plt.show()
        else:
            colors_list = ['#d7191c', '#fdae61', '#ffffbf', '#a6d96a', '#1a9641']
            bounds = [0, 50, 90, 110, 150, 200]
            cmap = ListedColormap(colors_list)
            norm = BoundaryNorm(bounds, cmap.N)
            
            data_var = pred_rain.sel(Y=slice(None, 19.5))
            n_times = len(data_var['T'])
            n_cols = 2
            n_rows = (n_times + n_cols - 1) // n_cols
            
            fig, axes = plt.subplots(
                n_rows, n_cols, figsize=(n_cols * 6, n_rows * 4),
                subplot_kw={'projection': ccrs.PlateCarree()}
            )
            axes = np.ravel(axes)
            
            for i, t in enumerate(data_var['T'].values):
                ax = axes[i]
                data = data_var.sel(T=t).squeeze()
                im = ax.pcolormesh(
                    data['X'], data['Y'], data, cmap=cmap, norm=norm,
                    transform=ccrs.PlateCarree()
                )
                ax.coastlines()
                ax.gridlines(draw_labels=False)
                ax.add_feature(cfeature.LAND, edgecolor="black")
                ax.add_feature(cfeature.OCEAN, facecolor="lightblue")
                ax.set_title(f"Season: {str(t)[:10]}")
            
            for j in range(n_times, len(axes)):
                fig.delaxes(axes[j])
            
            cbar = fig.colorbar(im, ax=axes[:n_times], orientation="horizontal", fraction=0.05, pad=0.1)
            cbar.set_label('Precipitation Ratio to Normal (%)')
            cbar.set_ticks([0, 50, 90, 110, 150, 200])
            cbar.ax.set_xticklabels(['0', '50', '90', '110', '150', '200'])
            fig.suptitle("Ratio to Normal Precipitation [%]", fontsize=14)
            plt.tight_layout()
            fig.subplots_adjust(top=0.9, bottom=0.15)
            plt.show()        
        return similar_years