# import xarray as xr
# import numpy as np
# from scipy.stats import skew, boxcox
# from sklearn.cluster import KMeans
# from sklearn.preprocessing import PowerTransformer
# from fitter import Fitter
# import pandas as pd
# import matplotlib.pyplot as plt
# import cartopy.crs as ccrs
# import cartopy.feature as cfeature
# from matplotlib.colors import ListedColormap, BoundaryNorm
# from scipy.stats import (norm, lognorm, expon, gamma as gamma_dist, weibull_min, t as t_dist, poisson, nbinom)


# def inv_boxcox(y, lmbda):
#     """
#     Inverse Box-Cox transformation for SciPy 1.11.3 compatibility.

#     Parameters
#     ----------
#     y : array_like
#         Transformed data.
#     lmbda : float
#         Box-Cox lambda parameter.

#     Returns
#     -------
#     x : ndarray
#         Original data before Box-Cox transformation.

#     Notes
#     -----
#     Implements the inverse of the Box-Cox transformation manually.
#     For lmbda near 0, uses exponential transformation.
#     """
#     if abs(lmbda) < 1e-6:
#         return np.exp(y)
#     return (y * lmbda + 1) ** (1 / lmbda)


# class WAS_TransformData:
#     """
#     Analyze skewness, apply transformations, fit distributions, and visualize 
#     geospatial time-series data.

#     This class provides methods to:
#     1. Detect and classify skewness in geospatial time-series data
#     2. Recommend transformations based on skewness characteristics
#     3. Apply various transformations to reduce skewness
#     4. Fit statistical distributions to data using clustering or grid-based methods
#     5. Visualize results as geospatial maps

#     Parameters
#     ----------
#     data : xarray.DataArray
#         Input data with dimensions (T, Y, X) for time, latitude, and longitude.
#     distribution_map : dict, optional
#         Mapping of distribution names to numeric codes. Default includes:
#         {'norm': 1, 'lognorm': 2, 'expon': 3, 'gamma': 4, 'weibull_min': 5,
#          't_dist': 6, 'poisson': 7, 'nbinom': 8}.
#     n_clusters : int, optional
#         Number of clusters for KMeans in distribution fitting. Default is 1000.

#     Attributes
#     ----------
#     data : xarray.DataArray
#         Input geospatial time-series data.
#     distribution_map : dict
#         Mapping of distribution names to codes.
#     n_clusters : int
#         Number of clusters for KMeans.
#     transformed_data : xarray.DataArray or None
#         Transformed data after applying transformations.
#     transform_methods : xarray.DataArray or None
#         Transformation methods applied per grid cell.
#     transform_params : xarray.DataArray or None
#         Parameters for parametric transformations (e.g., Box-Cox lambda).
#     skewness_ds : xarray.Dataset or None
#         Skewness analysis results.
#     handle_ds : xarray.Dataset or None
#         Skewness handling recommendations.

#     Raises
#     ------
#     ValueError
#         If `data` is not an xarray.DataArray or lacks required dimensions.
#     """

#     def __init__(self, data, distribution_map=None, n_clusters=1000):
#         """
#         Initialize the WAS_TransformData class.

#         Parameters
#         ----------
#         data : xarray.DataArray
#             Input data with dimensions (T, Y, X).
#         distribution_map : dict, optional
#             Mapping of distribution names to numeric codes.
#         n_clusters : int, optional
#             Number of clusters for KMeans.

#         Raises
#         ------
#         ValueError
#             If data is not xarray.DataArray or lacks required dimensions.
#         """
#         if not isinstance(data, xr.DataArray):
#             raise ValueError("`data` must be an xarray.DataArray")
#         if not all(dim in data.dims for dim in ('T', 'Y', 'X')):
#             raise ValueError("`data` must have dimensions ('T', 'Y', 'X')")

#         self.data = data
#         self.distribution_map = distribution_map or {
#             'norm': 1,
#             'lognorm': 2,
#             'expon': 3,
#             'gamma': 4,
#             'weibull_min': 5,
#             "t_dist": 6,
#             "poisson": 7,
#             "nbinom": 8
#         }
#         self.n_clusters = n_clusters
#         self.transformed_data = None
#         self.transform_methods = None
#         self.transform_params = None
#         self.skewness_ds = None
#         self.handle_ds = None

#     @staticmethod
#     def _safe_boxcox(arr1d):
#         """
#         Apply Box-Cox transformation while handling NaNs.

#         Parameters
#         ----------
#         arr1d : array_like
#             1D array of data to transform.

#         Returns
#         -------
#         transformed : ndarray
#             Transformed array, same shape as input, with NaNs preserved.
#         lmbda : float
#             Box-Cox lambda parameter.

#         Raises
#         ------
#         ValueError
#             If fewer than 2 non-NaN values or if data is not strictly positive.
#         """
#         out = arr1d.copy()
#         valid = ~np.isnan(arr1d)
#         if valid.sum() < 2:
#             raise ValueError("Need at least two non-NaN values for Box-Cox")
#         if not np.all(arr1d[valid] > 0):
#             raise ValueError("Box-Cox requires strictly positive data")
#         out[valid], lmbda = boxcox(arr1d[valid])
#         return out, lmbda

#     def detect_skewness(self):
#         """
#         Compute and classify skewness for each grid cell.

#         Returns
#         -------
#         skewness_ds : xarray.Dataset
#             Dataset with variables:
#             - 'skewness' (float): Skewness coefficient
#             - 'skewness_class' (str): Classification of skewness
#         summary : dict
#             Dictionary with 'class_counts' mapping skewness classes to grid cell counts.

#         Notes
#         -----
#         Skewness classification thresholds:
#         - Symmetric: -0.5 ≤ skewness ≤ 0.5
#         - Moderate positive: 0.5 < skewness ≤ 1
#         - Moderate negative: -1 ≤ skewness < -0.5
#         - High positive: skewness > 1
#         - High negative: skewness < -1
#         - Invalid: insufficient data (< 3 non-NaN values)

#         Skewness is computed using `scipy.stats.skew` with `nan_policy='omit'`.
#         """
#         def _compute(precip):
#             precip = np.asarray(precip)
#             valid = ~np.isnan(precip)
#             if valid.sum() < 3:
#                 return np.nan, 'invalid'
#             sk = skew(precip[valid], axis=0, nan_policy='omit')
#             if np.isnan(sk):
#                 cls = 'invalid'
#             elif -0.5 <= sk <= 0.5:
#                 cls = 'symmetric'
#             elif 0.5 < sk <= 1:
#                 cls = 'moderate_positive'
#             elif -1 <= sk < -0.5:
#                 cls = 'moderate_negative'
#             elif sk > 1:
#                 cls = 'high_positive'
#             else:
#                 cls = 'high_negative'
#             return sk, cls

#         res = xr.apply_ufunc(
#             _compute,
#             self.data,
#             input_core_dims=[['T']],
#             output_core_dims=[[], []],
#             vectorize=True,
#             dask='parallelized',
#             output_dtypes=[float, str]
#         )

#         self.skewness_ds = xr.Dataset(
#             {
#                 'skewness': (('Y', 'X'), res[0].data),
#                 'skewness_class': (('Y', 'X'), res[1].data)
#             },
#             coords={'Y': self.data.Y, 'X': self.data.X}
#         )

#         counts = pd.Series(self.skewness_ds['skewness_class'].values.ravel()).value_counts().to_dict()
#         return self.skewness_ds, {'class_counts': counts}

#     def handle_skewness(self):
#         """
#         Recommend transformations based on skewness and data properties.

#         Returns
#         -------
#         handle_ds : xarray.Dataset
#             Dataset with variables:
#             - 'skewness': Skewness coefficient
#             - 'skewness_class': Skewness classification
#             - 'recommended_methods': Semicolon-separated transformation methods
#         summary : dict
#             Dictionary with 'general_recommendations' mapping skewness classes to advice.

#         Raises
#         ------
#         ValueError
#             If `detect_skewness` has not been called.

#         Notes
#         -----
#         Recommendations consider data properties (zeros, negatives) and skewness class.
#         Available methods: 'log', 'square_root', 'box_cox', 'yeo_johnson',
#         'clipping', 'binning', 'reflect_log', 'reflect_square_root', 'reflect_yeo_johnson'.
#         """
#         if self.skewness_ds is None:
#             raise ValueError("Run detect_skewness() first")

#         def _suggest(precip, sk_class):
#             if sk_class == 'invalid':
#                 return 'none'
#             precip = np.asarray(precip)
#             valid = precip[~np.isnan(precip)]
#             all_pos = np.all(valid > 0)
#             has_zeros = np.any(valid == 0)
#             methods = []
#             if sk_class in ('moderate_positive', 'high_positive'):
#                 if all_pos and not has_zeros:
#                     methods += ['log', 'square_root', 'box_cox']
#                 elif all_pos:
#                     methods += ['square_root', 'box_cox']
#                 methods += ['yeo_johnson', 'clipping', 'binning']
#             elif sk_class in ('moderate_negative', 'high_negative'):
#                 if all_pos and not has_zeros:
#                     methods += ['reflect_log']
#                 elif all_pos:
#                     methods += ['reflect_square_root']
#                 methods += ['reflect_yeo_johnson', 'clipping', 'binning']
#             else:
#                 methods.append('none')
#             return ';'.join(methods)

#         recommended = xr.apply_ufunc(
#             _suggest,
#             self.data,
#             self.skewness_ds['skewness_class'],
#             input_core_dims=[['T'], []],
#             output_core_dims=[[]],
#             vectorize=True,
#             dask='parallelized',
#             output_dtypes=[str]
#         )

#         self.handle_ds = xr.Dataset(
#             {
#                 'skewness': self.skewness_ds['skewness'],
#                 'skewness_class': self.skewness_ds['skewness_class'],
#                 'recommended_methods': (('Y', 'X'), recommended.data)
#             },
#             coords={'Y': self.data.Y, 'X': self.data.X}
#         )

#         general = {
#             'symmetric': 'No transformation needed.',
#             'moderate_positive': (
#                 'Consider square root or Yeo-Johnson; log or Box-Cox if no zeros; '
#                 'clip or bin outliers.'
#             ),
#             'high_positive': (
#                 'Strongly consider log (no zeros), Box-Cox (positive), or Yeo-Johnson; '
#                 'clip or bin extremes.'
#             ),
#             'moderate_negative': (
#                 'Reflect and apply square root or Yeo-Johnson; clip or bin outliers.'
#             ),
#             'high_negative': (
#                 'Reflect and apply log (no zeros), Box-Cox, or Yeo-Johnson; '
#                 'clip or bin extremes.'
#             ),
#             'invalid': 'Insufficient valid data for skewness calculation.'
#         }

#         return self.handle_ds, {'general_recommendations': general}

#     def apply_transformation(self, method=None):
#         """
#         Apply transformations to reduce skewness in the data.

#         Parameters
#         ----------
#         method : str or xarray.DataArray, optional
#             Transformation method to apply:
#             - None: Use first recommended method per grid cell from `handle_skewness`
#             - str: Apply the same method to all grid cells (e.g., 'log', 'box_cox')
#             - xarray.DataArray: Specify method per grid cell with dimensions (Y, X)

#         Returns
#         -------
#         transformed_data : xarray.DataArray
#             Transformed data with same shape as input.

#         Raises
#         ------
#         ValueError
#             If `method` is None and `handle_skewness` has not been called.

#         Notes
#         -----
#         Supported methods: 'log', 'square_root', 'box_cox', 'yeo_johnson',
#         'reflect_log', 'reflect_square_root', 'reflect_yeo_johnson', 'clipping', 'binning'.
#         Transformations are skipped for invalid data or methods, with warnings printed.
#         """
#         if method is None and self.handle_ds is None:
#             raise ValueError("Run handle_skewness() first or specify `method`")

#         if method is None:
#             def extract_first_method(x):
#                 if isinstance(x, str) and x and x != 'none':
#                     return x.split(';')[0]
#                 return 'none'
#             method = xr.apply_ufunc(
#                 extract_first_method,
#                 self.handle_ds['recommended_methods'],
#                 vectorize=True,
#                 output_dtypes=[str]
#             )

#         self.transformed_data = self.data.copy()
#         self.transform_methods = method if isinstance(method, xr.DataArray) else xr.DataArray(
#             np.full((self.data.sizes['Y'], self.data.sizes['X']), method),
#             coords={'Y': self.data.Y, 'X': self.data.X},
#             dims=('Y', 'X')
#         )
#         self.transform_params = xr.DataArray(
#             np.empty((self.data.sizes['Y'], self.data.sizes['X']), dtype=object),
#             coords={'Y': self.data.Y, 'X': self.data.X},
#             dims=('Y', 'X')
#         )

#         for iy in range(self.data.sizes['Y']):
#             for ix in range(self.data.sizes['X']):
#                 m = self.transform_methods[iy, ix].item()
#                 if m == 'none' or np.all(np.isnan(self.data[:, iy, ix])):
#                     continue
#                 cell = self.data[:, iy, ix].values
#                 valid = cell[~np.isnan(cell)]
#                 if len(valid) < 2:
#                     continue

#                 if m == 'log':
#                     if np.any(valid <= 0):
#                         print(f"Skip log at Y={iy}, X={ix}: non-positive values")
#                         continue
#                     self.transformed_data[:, iy, ix] = np.log(cell)
#                 elif m == 'square_root':
#                     if np.any(valid < 0):
#                         print(f"Skip square_root at Y={iy}, X={ix}: negative values")
#                         continue
#                     self.transformed_data[:, iy, ix] = np.sqrt(cell)
#                 elif m == 'box_cox':
#                     try:
#                         transformed, lam = self._safe_boxcox(cell)
#                         self.transformed_data[:, iy, ix] = transformed
#                         self.transform_params[iy, ix] = {'lambda': lam}
#                     except ValueError as err:
#                         print(f"Skip Box-Cox at Y={iy}, X={ix}: {err}")
#                         continue
#                 elif m == 'yeo_johnson':
#                     pt = PowerTransformer(method='yeo-johnson')
#                     transformed = pt.fit_transform(cell.reshape(-1, 1)).ravel()
#                     self.transformed_data[:, iy, ix] = transformed
#                     self.transform_params[iy, ix] = {'transformer': pt}
#                 elif m == 'reflect_log':
#                     cell_ref = -cell
#                     if np.any(cell_ref <= 0):
#                         print(f"Skip reflect_log at Y={iy}, X={ix}: non-positive values")
#                         continue
#                     self.transformed_data[:, iy, ix] = np.log(cell_ref)
#                 elif m == 'reflect_square_root':
#                     cell_ref = -cell
#                     if np.any(cell_ref < 0):
#                         print(f"Skip reflect_square_root at Y={iy}, X={ix}: negative values")
#                         continue
#                     self.transformed_data[:, iy, ix] = np.sqrt(cell_ref)
#                 elif m == 'reflect_yeo_johnson':
#                     pt = PowerTransformer(method='yeo-johnson')
#                     transformed = pt.fit_transform((-cell).reshape(-1, 1)).ravel()
#                     self.transformed_data[:, iy, ix] = transformed
#                     self.transform_params[iy, ix] = {'transformer': pt}
#                 elif m in ('clipping', 'binning'):
#                     self.transformed_data[:, iy, ix] = cell
#                 else:
#                     pass

#         return self.transformed_data

#     def inverse_transform(self):
#         """
#         Reverse transformations to recover original data scale.

#         Returns
#         -------
#         inverse_data : xarray.DataArray
#             Data in original scale with same shape as input.

#         Raises
#         ------
#         ValueError
#             If no transformation has been applied or required parameters are missing.

#         Notes
#         -----
#         Non-invertible methods ('clipping', 'binning') return unchanged data with a warning.
#         """
#         if self.transformed_data is None or self.transform_methods is None:
#             raise ValueError("No transformation applied. Run apply_transformation() first")

#         def _inv(vec, method, params):
#             if method in ('none', None) or (isinstance(method, float) and np.isnan(method)):
#                 return vec
#             if method in ('clipping', 'binning'):
#                 print(f"Warning: '{method}' is not invertible")
#                 return vec
#             if method == 'log':
#                 return np.exp(vec)
#             if method == 'square_root':
#                 return vec ** 2
#             if method == 'box_cox':
#                 lam = params.get('lambda') if params else None
#                 if lam is None:
#                     raise ValueError("Missing lambda for Box-Cox inversion")
#                 return inv_boxcox(vec, lam)
#             if method == 'yeo_johnson':
#                 tr = params.get('transformer') if params else None
#                 if tr is None:
#                     raise ValueError("Missing transformer for Yeo-Johnson inversion")
#                 return tr.inverse_transform(vec.reshape(-1, 1)).ravel()
#             if method.startswith('reflect_'):
#                 if method == 'reflect_log':
#                     temp = np.exp(vec)
#                 elif method == 'reflect_square_root':
#                     temp = vec ** 2
#                 else:  # reflect_yeo_johnson
#                     tr = params.get('transformer') if params else None
#                     if tr is None:
#                         raise ValueError("Missing transformer for reflect_yeo_johnson")
#                     temp = tr.inverse_transform(vec.reshape(-1, 1)).ravel()
#                 return -temp
#             raise ValueError(f"Unknown method '{method}'")

#         return xr.apply_ufunc(
#             _inv,
#             self.transformed_data,
#             self.transform_methods,
#             self.transform_params,
#             input_core_dims=[['T'], [], []],
#             output_core_dims=[['T']],
#             vectorize=True,
#             dask='parallelized',
#             output_dtypes=[float]
#         )

#     def fit_best_distribution_grid_onlycluster(self, use_transformed=False):
#         """
#         Fit best distributions by homogeneous zones (clusters) and map to each grid cell.

#         Parameters
#         ----------
#         use_transformed : bool, optional
#             If True, use self.transformed_data; otherwise use original self.data.

#         Returns
#         -------
#         best_code : xarray.DataArray, shape (Y, X)
#             Code of best distribution per grid cell.
#         best_shape : xarray.DataArray, shape (Y, X)
#             Primary shape parameter (e.g., 'a' in gamma, 'mu' in poisson, 'n' in nbinom).
#         best_loc : xarray.DataArray, shape (Y, X)
#             Location parameter.
#         best_scale : xarray.DataArray, shape (Y, X)
#             Scale parameter (e.g., 'p' in nbinom).
#         cluster_da : xarray.DataArray, shape (Y, X)
#             Cluster labels for each grid cell.

#         Notes
#         -----
#         Method:
#         1. Cluster grid cells based on mean and standard deviation using KMeans
#         2. For each cluster, pool all time series values
#         3. Fit candidate distributions to pooled data using Maximum Likelihood Estimation
#         4. Select distribution with minimum Akaike Information Criterion (AIC)
#         5. Assign distribution parameters to all cells in the cluster

#         Supported distributions (if present in distribution_map):
#         - Continuous (precip-like): 'lognorm', 'gamma', 'weibull_min', 'expon' (loc=0)
#         - Continuous (unbounded): 'norm', 't' (no non-negativity constraint)
#         - Discrete (counts): 'poisson', 'nbinom' (non-negative INTEGER values, loc=0)

#         Requires sufficient valid data per cluster; returns NaNs for insufficient data.
#         """
#         import numpy as np
#         import xarray as xr
#         from sklearn.cluster import KMeans
#         from scipy.stats import (
#             norm,
#             lognorm,
#             expon,
#             gamma as gamma_dist,
#             weibull_min,
#             t as t_dist,
#             poisson,
#             nbinom,
#         )

#         # -------- Select working data --------
#         data = self.transformed_data if (use_transformed and self.transformed_data is not None) else self.data

#         if not isinstance(data, xr.DataArray):
#             raise ValueError("Internal error: data must be an xarray.DataArray")

#         if not all(dim in data.dims for dim in ("T", "Y", "X")):
#             raise ValueError("`data` must have dimensions ('T', 'Y', 'X')")

#         Y = data.sizes["Y"]
#         X = data.sizes["X"]
#         coords = {"Y": data.Y, "X": data.X}

#         # Map distribution names to scipy objects
#         dist_objs = {
#             "norm": norm,
#             "lognorm": lognorm,
#             "expon": expon,
#             "gamma": gamma_dist,
#             "weibull_min": weibull_min,
#             "t": t_dist,
#             "poisson": poisson,
#             "nbinom": nbinom,
#         }

#         # -------- 1. Build clustering features (mean, std) per grid cell --------
#         mean_da = data.mean("T", skipna=True)
#         std_da = data.std("T", skipna=True)

#         feat_ds = xr.Dataset({"mean": mean_da, "std": std_da})
#         feat_df = feat_ds.to_dataframe().dropna()  # index: (Y, X), cols: mean, std

#         if feat_df.shape[0] < self.n_clusters:
#             print("Warning: insufficient valid grid cells for clustering; returning NaNs.")
#             nan_arr = np.full((Y, X), np.nan, float)
#             return (
#                 xr.DataArray(nan_arr, coords=coords, dims=("Y", "X")),
#                 xr.DataArray(nan_arr, coords=coords, dims=("Y", "X")),
#                 xr.DataArray(nan_arr, coords=coords, dims=("Y", "X")),
#                 xr.DataArray(nan_arr, coords=coords, dims=("Y", "X")),
#                 xr.DataArray(nan_arr, coords=coords, dims=("Y", "X")),
#             )

#         # Run KMeans on [mean, std]
#         features = feat_df[["mean", "std"]].values
#         kmeans = KMeans(n_clusters=self.n_clusters, random_state=42, n_init=10)
#         feat_df["cluster"] = kmeans.fit_predict(features)

#         # -------- 2. Map cluster labels back to full (Y, X) grid --------
#         cluster_da = (
#             feat_df["cluster"]
#             .to_xarray()
#             .reindex(Y=data.Y, X=data.X)
#         )

#         # -------- 3. For each cluster, fit best distribution on pooled values --------
#         best_params_by_cluster = {}  # cl -> (code, shape, loc, scale)

#         unique_clusters = np.unique(cluster_da.values)
#         unique_clusters = unique_clusters[np.isfinite(unique_clusters)]

#         for cl_val in unique_clusters:
#             cl = int(cl_val)

#             # Extract pooled values for this cluster
#             mask_cl = (cluster_da == cl)
#             cl_vals = data.where(mask_cl).values
#             cl_vals = cl_vals[np.isfinite(cl_vals)]

#             if cl_vals.size < 30:
#                 best_params_by_cluster[cl] = None
#                 continue

#             # Create specific subsets for different distribution families
#             cl_vals_pos = cl_vals[cl_vals >= 0]
#             cl_vals_int = cl_vals_pos[(cl_vals_pos == np.floor(cl_vals_pos))]

#             best_aic = np.inf
#             best_choice = None  # Will store (code, (shape, loc, scale))

#             for name, code in self.distribution_map.items():
#                 if name not in dist_objs:
#                     continue

#                 dist = dist_objs[name]
#                 is_discrete = False

#                 # Choose appropriate sample for this distribution
#                 if name in ("poisson", "nbinom"):
#                     sample = cl_vals_int
#                     is_discrete = True
#                 elif name in ("lognorm", "gamma", "weibull_min", "expon"):
#                     sample = cl_vals_pos
#                 else:
#                     # 'norm' and 't' work on full sample
#                     sample = cl_vals

#                 if sample.size < 30:
#                     continue

#                 try:
#                     # --- Fit parameters and get k (number of estimated params) ---
#                     if name == "poisson":
#                         # Manual MLE fit for Poisson (mu=mean), loc is fixed
#                         mu = sample.mean()
#                         loc = 0
#                         params = (mu, loc)
#                         k = 1  # Only mu is estimated

#                     elif name in ("lognorm", "gamma", "weibull_min", "expon", "nbinom"):
#                         # Fit with fixed location at 0
#                         params = dist.fit(sample, floc=0)
#                         k = len(params) - 1  # loc is fixed, not estimated

#                     else:
#                         # 'norm' and 't' (and any others with free loc)
#                         params = dist.fit(sample)
#                         k = len(params)  # All params (loc, scale, shapes) are estimated

#                     # --- Calculate Log-Likelihood ---
#                     if is_discrete:
#                         if name == "poisson":
#                             logL = np.sum(dist.logpmf(sample, mu=params[0], loc=params[1]))
#                         else:  # nbinom
#                             logL = np.sum(dist.logpmf(sample, *params))
#                     else:
#                         logL = np.sum(dist.logpdf(sample, *params))

#                     if not np.isfinite(logL):
#                         continue

#                     aic = 2 * k - 2 * logL

#                     if aic < best_aic:
#                         best_aic = aic
#                         # --- Standardize parameters to (shape, loc, scale) ---
#                         if name == "poisson":
#                             # params = (mu, loc)
#                             best_choice = (code, (params[0], params[1], np.nan))
#                         elif name == "nbinom":
#                             # params = (n, p, loc)
#                             # Store as (shape=n, loc=loc, scale=p)
#                             best_choice = (code, (params[0], params[2], params[1]))
#                         elif name in ("norm", "expon"):
#                             # params = (loc, scale)
#                             # Store as (shape=nan, loc=loc, scale=scale)
#                             best_choice = (code, (np.nan, params[0], params[1]))
#                         else:
#                             # 't', 'gamma', 'lognorm', 'weibull_min'
#                             # params = (shape, loc, scale)
#                             best_choice = (code, (params[0], params[1], params[2]))

#                 except Exception:
#                     continue  # Fitting failed

#             if best_choice is None:
#                 best_params_by_cluster[cl] = None
#             else:
#                 code, params_tuple = best_choice
#                 # params_tuple is already (shape, loc, scale)
#                 best_params_by_cluster[cl] = (code, params_tuple[0], params_tuple[1], params_tuple[2])

#         # -------- 4. Broadcast cluster-level params to each grid cell (Vectorized) --------

#         # Create empty arrays with the correct coordinates and NaNs
#         best_code_da = xr.DataArray(np.full((Y, X), np.nan, dtype=float), coords=coords, dims=("Y", "X"))
#         best_shape_da = best_code_da.copy()
#         best_loc_da = best_code_da.copy()
#         best_scale_da = best_code_da.copy()

#         # Loop over the few clusters, not the millions of pixels
#         for cl, params in best_params_by_cluster.items():
#             if params is None:
#                 continue  # Leave these cells as NaN

#             # Create a boolean mask for all cells belonging to this cluster
#             mask = (cluster_da == cl)

#             # Unpack the parameters
#             code, shape, loc, scale = params

#             # "Paint" the values onto the grid where the mask is True
#             best_code_da = best_code_da.where(~mask, code)
#             best_shape_da = best_shape_da.where(~mask, shape)
#             best_loc_da = best_loc_da.where(~mask, loc)
#             best_scale_da = best_scale_da.where(~mask, scale)

#         return best_code_da, best_shape_da, best_loc_da, best_scale_da, cluster_da

#     def fit_best_distribution_grid_onlygrid(self, use_transformed=False):
#         """
#         Fit candidate distributions per grid cell and select best by AIC.

#         Parameters
#         ----------
#         use_transformed : bool, optional
#             If True, use self.transformed_data; otherwise use original data.

#         Returns
#         -------
#         best_code : xarray.DataArray (Y, X)
#             Code of best distribution per grid cell.
#         best_shape : xarray.DataArray (Y, X)
#             Shape parameter of best distribution.
#         best_loc : xarray.DataArray (Y, X)
#             Location parameter of best distribution.
#         best_scale : xarray.DataArray (Y, X)
#             Scale parameter of best distribution.

#         Notes
#         -----
#         - Fits are done independently per grid cell (no clustering)
#         - For precip-like variables:
#             * 'lognorm', 'gamma', 'weibull_min', 'expon' are fitted on positive values with loc=0
#             * 'norm' and 't' are fitted on all finite values
#         - AIC = 2k - 2 ln(L) is used for model selection
#         - Requires minimum of 10 valid data points per grid cell
#         """
#         import numpy as np
#         import xarray as xr
#         from scipy.stats import norm, lognorm, expon, gamma as gamma_dist, weibull_min, t as t_dist

#         data = self.transformed_data if (use_transformed and self.transformed_data is not None) else self.data

#         if not isinstance(data, xr.DataArray):
#             raise ValueError("`data` must be an xarray.DataArray")
#         if not all(dim in data.dims for dim in ("T", "Y", "X")):
#             raise ValueError("`data` must have dimensions ('T', 'Y', 'X')")

#         # Map distribution names to scipy.stats objects
#         dist_objs = {
#             "norm": norm,
#             "lognorm": lognorm,
#             "expon": expon,
#             "gamma": gamma_dist,
#             "weibull_min": weibull_min,
#             "t": t_dist,
#         }

#         Y = data.sizes["Y"]
#         X = data.sizes["X"]

#         best_code = np.full((Y, X), np.nan, dtype=float)
#         best_shape = np.full((Y, X), np.nan, dtype=float)
#         best_loc = np.full((Y, X), np.nan, dtype=float)
#         best_scale = np.full((Y, X), np.nan, dtype=float)

#         for iy in range(Y):
#             for ix in range(X):
#                 vals = data[:, iy, ix].values
#                 vals = vals[np.isfinite(vals)]
#                 if vals.size < 10:
#                     continue

#                 # Positive subset for positive-support distributions
#                 vals_pos = vals[vals > 0]

#                 best_aic = np.inf
#                 best = None

#                 for name, code in self.distribution_map.items():
#                     if name not in dist_objs:
#                         continue

#                     dist = dist_objs[name]

#                     # Choose sample depending on support
#                     if name in ("lognorm", "gamma", "weibull_min", "expon"):
#                         sample = vals_pos
#                         if sample.size < 10:
#                             continue
#                     else:  # 'norm', 't', or any other real-line distribution
#                         sample = vals
#                         if sample.size < 10:
#                             continue

#                     try:
#                         # Fit parameters
#                         if name in ("lognorm", "gamma", "weibull_min", "expon"):
#                             params = dist.fit(sample, floc=0)
#                         else:
#                             # norm, t: free loc/scale; t has (df, loc, scale)
#                             params = dist.fit(sample)

#                         # Compute AIC
#                         k = len(params)
#                         logL = np.sum(dist.logpdf(sample, *params))
#                         aic = 2 * k - 2.0 * logL

#                         if np.isfinite(aic) and aic < best_aic:
#                             best_aic = aic
#                             best = (code, params)
#                     except Exception:
#                         # Skip distributions that fail to fit at this grid
#                         continue

#                 if best is not None:
#                     code, params = best
#                     # Normalise to (shape, loc, scale)
#                     if len(params) == 2:
#                         # e.g. norm, expon when not forcing shape
#                         shape, loc, scale = (np.nan, params[0], params[1])
#                     else:
#                         # e.g. gamma (k, loc, scale), t (df, loc, scale), etc.
#                         shape, loc, scale = params[0], params[1], params[2]

#                     best_code[iy, ix] = code
#                     best_shape[iy, ix] = shape
#                     best_loc[iy, ix] = loc
#                     best_scale[iy, ix] = scale

#         coords = {"Y": data.Y, "X": data.X}
#         best_code_da = xr.DataArray(best_code, coords=coords, dims=("Y", "X"))
#         best_shape_da = xr.DataArray(best_shape, coords=coords, dims=("Y", "X"))
#         best_loc_da = xr.DataArray(best_loc, coords=coords, dims=("Y", "X"))
#         best_scale_da = xr.DataArray(best_scale, coords=coords, dims=("Y", "X"))

#         return best_code_da, best_shape_da, best_loc_da, best_scale_da

#     def fit_best_distribution_grid_two_options(self, use_transformed=False, mode="cluster"):
#         """
#         Fit best distributions using either clustering or per-grid methods.

#         Parameters
#         ----------
#         use_transformed : bool, optional
#             If True, use self.transformed_data, else self.data.
#         mode : {'cluster', 'grid'}, optional
#             Method for fitting:
#             - 'cluster': KMeans on (mean,std) -> one distribution per cluster
#             - 'grid': independent fit at each (Y,X) via xr.apply_ufunc

#         Returns
#         -------
#         best_code : xarray.DataArray
#             (Y, X) Array of distribution codes.
#         best_shape : xarray.DataArray
#             (Y, X) Fitted shape parameters.
#         best_loc : xarray.DataArray
#             (Y, X) Fitted location parameters.
#         best_scale : xarray.DataArray
#             (Y, X) Fitted scale parameters.
#         cluster_da : xarray.DataArray
#             (Y, X) Cluster labels for mode='cluster'; all-NaN for mode='grid'.

#         Raises
#         ------
#         ValueError
#             If mode is not 'cluster' or 'grid', or if self.n_clusters is not defined for cluster mode.

#         Notes
#         -----
#         Supported distributions:
#         'norm', 'lognorm', 'expon', 'gamma', 'weibull_min', 't', 'poisson', 'nbinom'
#         (must be present in self.distribution_map).

#         Uses AIC for model selection with different minimum sample sizes:
#         - Grid mode: min 15 samples per cell
#         - Cluster mode: min 30 samples per cluster
#         """
#         # ------------------------------------------------------------------
#         # Select data
#         # ------------------------------------------------------------------
#         data = self.transformed_data if (use_transformed and self.transformed_data is not None) else self.data

#         if not isinstance(data, xr.DataArray):
#             raise ValueError("`data` must be an xarray.DataArray")
#         if not all(dim in data.dims for dim in ("T", "Y", "X")):
#             raise ValueError("`data` must have dimensions ('T', 'Y', 'X')")

#         if not hasattr(self, "distribution_map") or not isinstance(self.distribution_map, dict):
#             raise ValueError("`self.distribution_map` must be a dict of {name: code}")

#         Y = data.sizes["Y"]
#         X = data.sizes["X"]
#         coords = {"Y": data.Y, "X": data.X}

#         # ------------------------------------------------------------------
#         # Map distribution names in distribution_map to scipy.stats objects
#         # (allow aliases like 't_dist' → t)
#         # ------------------------------------------------------------------
#         name_to_dist = {
#             "norm": norm,
#             "lognorm": lognorm,
#             "expon": expon,
#             "gamma": gamma_dist,
#             "weibull_min": weibull_min,
#             "t": t_dist,
#             "t_dist": t_dist,
#             "poisson": poisson,
#             "nbinom": nbinom,
#         }

#         # Helper: filter only supported distributions
#         dist_candidates = {
#             name: (name_to_dist[name], code)
#             for name, code in self.distribution_map.items()
#             if name in name_to_dist
#         }
#         if not dist_candidates:
#             raise ValueError("No valid distributions found in distribution_map")

#         # ------------------------------------------------------------------
#         # Core 1D fitter using AIC
#         # ------------------------------------------------------------------
#         def _fit_best(sample_1d, min_n):
#             """
#             Fit all candidate distributions to a 1D sample; return (code, shape, loc, scale)
#             for the best by AIC. If no fit, return NaNs.
#             """
#             vals = np.asarray(sample_1d, float)
#             vals = vals[np.isfinite(vals)]
#             if vals.size < min_n:
#                 return np.nan, np.nan, np.nan, np.nan

#             vals_pos = vals[vals > 0]
#             vals_int = vals_pos[(vals_pos == np.floor(vals_pos))]

#             best_aic = np.inf
#             best = None  # (code, (shape, loc, scale))

#             for name, (dist, code) in dist_candidates.items():
#                 is_discrete = name in ("poisson", "nbinom")

#                 # respect support
#                 if name in ("poisson", "nbinom"):
#                     sample = vals_int
#                 elif name in ("lognorm", "gamma", "weibull_min", "expon"):
#                     sample = vals_pos
#                 else:  # norm, t, etc.
#                     sample = vals

#                 if sample.size < min_n:
#                     continue

#                 try:
#                     # ---- fit parameters ----
#                     if name == "poisson":
#                         # MLE: mu = mean, loc=0
#                         mu = sample.mean()
#                         if mu <= 0 or not np.isfinite(mu):
#                             continue
#                         params = (mu, 0.0)
#                         k = 1  # only mu

#                     elif name == "nbinom":
#                         # fit (n, p, loc=0)
#                         params = dist.fit(sample, floc=0)
#                         # (n, p, loc) with loc fixed
#                         k = len(params) - 1

#                     elif name in ("lognorm", "gamma", "weibull_min", "expon"):
#                         # positive-support, loc=0 fixed
#                         params = dist.fit(sample, floc=0)
#                         k = len(params) - 1

#                     else:
#                         # norm, t: all params free
#                         params = dist.fit(sample)
#                         k = len(params)

#                     # ---- log-likelihood ----
#                     if is_discrete:
#                         if name == "poisson":
#                             mu, loc = params
#                             logL = np.sum(dist.logpmf(sample, mu, loc=loc))
#                         else:  # nbinom
#                             logL = np.sum(dist.logpmf(sample, *params))
#                     else:
#                         logL = np.sum(dist.logpdf(sample, *params))

#                     if not np.isfinite(logL):
#                         continue

#                     aic = 2.0 * k - 2.0 * logL
#                     if aic < best_aic:
#                         best_aic = aic

#                         # normalize to (shape, loc, scale)
#                         if name == "poisson":
#                             mu, loc = params
#                             best = (code, (mu, loc, np.nan))
#                         elif name == "nbinom":
#                             n_, p_, loc_ = params  # (n, p, loc)
#                             best = (code, (n_, loc_, p_))  # shape=n, loc=loc, scale=p
#                         elif name in ("norm", "expon"):
#                             loc_, scale_ = params
#                             best = (code, (np.nan, loc_, scale_))
#                         else:
#                             # (shape, loc, scale): t, gamma, lognorm, weibull_min
#                             shape_, loc_, scale_ = params
#                             best = (code, (shape_, loc_, scale_))

#                 except Exception:
#                     continue

#             if best is None:
#                 return np.nan, np.nan, np.nan, np.nan

#             code, (shape, loc, scale) = best
#             return float(code), float(shape), float(loc), float(scale)

#         # ------------------------------------------------------------------
#         # Mode: per-grid using xr.apply_ufunc
#         # ------------------------------------------------------------------
#         if mode == "grid":
#             min_n_grid = 15

#             def _fit_best_grid(cell_ts):
#                 return _fit_best(cell_ts, min_n_grid)

#             best_code_da, best_shape_da, best_loc_da, best_scale_da = xr.apply_ufunc(
#                 _fit_best_grid,
#                 data,
#                 input_core_dims=[["T"]],
#                 output_core_dims=[[], [], [], []],
#                 vectorize=True,
#                 dask="parallelized",
#                 output_dtypes=[float, float, float, float],
#             )

#             cluster_da = xr.full_like(best_code_da, np.nan)
#             return (
#                 best_code_da,
#                 best_shape_da,
#                 best_loc_da,
#                 best_scale_da,
#                 cluster_da,
#             )

#         # ------------------------------------------------------------------
#         # Mode: homogeneous zones via KMeans clustering
#         # ------------------------------------------------------------------
#         if mode == "cluster":
#             if not hasattr(self, "n_clusters"):
#                 raise ValueError("For mode='cluster', self.n_clusters must be defined.")

#             min_n_cluster = 30

#             kmeans = KMeans(n_clusters=self.n_clusters, random_state=42, n_init=10)
#             data_dropna = data.to_dataframe().reset_index().dropna().drop(columns=['T'])
#             variable_column = data_dropna.columns[2]
#             data_dropna['cluster'] = kmeans.fit_predict(
#                 data_dropna[[variable_column]]
#             )

#             # Convert cluster assignments back into an xarray structure
#             df_unique = data_dropna.drop_duplicates(subset=['Y', 'X'])
#             dataset = df_unique.set_index(['Y', 'X']).to_xarray()
#             mask = xr.where(~np.isnan(data.isel(T=0)), 1, np.nan)
#             Cluster = (dataset['cluster'] * mask)

#             # Align cluster array with the predictand array
#             xarray1, xarray2 = xr.align(data, Cluster, join="outer")

#             # Identify unique cluster labels
#             clusters = np.unique(xarray2)
#             unique_clusters = clusters[~np.isnan(clusters)]
#             cluster_da = xarray2

#             # 4) Fit best distribution for each cluster (pooled values)
#             best_params_by_cluster = {}

#             for cl_val in unique_clusters:
#                 cl = int(cl_val)
#                 mask_cl = (cluster_da == cl).expand_dims({'T': data['T']})
#                 cl_vals = data.where(mask_cl).stack(sample=('T', 'Y', 'X')).values.ravel()
#                 cl_vals = cl_vals[np.isfinite(cl_vals)]
#                 code, shape, loc, scale = _fit_best(cl_vals, min_n_cluster)
#                 if np.isnan(code):
#                     best_params_by_cluster[cl] = None
#                 else:
#                     best_params_by_cluster[cl] = (code, shape, loc, scale)

#             # 5) Broadcast cluster-level params back to (Y,X) using xarray masks
#             cluster_da_ = cluster_da.drop_vars("T")
#             best_code_da = xr.full_like(cluster_da_, np.nan, dtype=float)
#             best_shape_da = xr.full_like(cluster_da_, np.nan, dtype=float)
#             best_loc_da = xr.full_like(cluster_da_, np.nan, dtype=float)
#             best_scale_da = xr.full_like(cluster_da_, np.nan, dtype=float)

#             for cl, params in best_params_by_cluster.items():
#                 if params is None:
#                     continue
#                 code, shape, loc, scale = params
#                 mask_cl_ = (cluster_da_ == cl)

#                 best_code_da = best_code_da.where(~mask_cl_, code)
#                 best_shape_da = best_shape_da.where(~mask_cl_, shape)
#                 best_loc_da = best_loc_da.where(~mask_cl_, loc)
#                 best_scale_da = best_scale_da.where(~mask_cl_, scale)

#             return best_code_da, best_shape_da, best_loc_da, best_scale_da, cluster_da_

#         # ------------------------------------------------------------------
#         # Invalid mode
#         # ------------------------------------------------------------------
#         raise ValueError("mode must be 'cluster' or 'grid'")

#     def find_best_distribution_grid___(self, use_transformed=False):
#         """
#         Fit distributions to data using KMeans clustering (legacy method).

#         Parameters
#         ----------
#         use_transformed : bool, optional
#             If True, use transformed data; otherwise, use original data.

#         Returns
#         -------
#         dist_codes : xarray.DataArray
#             Numeric codes for best-fitting distributions per grid cell.

#         Notes
#         -----
#         Legacy method using `fitter.Fitter` to fit distributions.
#         Clusters are determined by mean values using KMeans.
#         Returns NaN for insufficient data.
#         """
#         data = self.transformed_data if use_transformed and self.transformed_data is not None else self.data
#         dist_names = tuple(self.distribution_map.keys())
#         df_mean = data.mean('T', skipna=True).to_dataframe(name='value').dropna()
#         if len(df_mean) < self.n_clusters:
#             print("Warning: Insufficient data for clustering, returning NaN array")
#             return xr.DataArray(
#                 np.full((self.data.sizes['Y'], self.data.sizes['X']), np.nan),
#                 coords={'Y': self.data.Y, 'X': self.data.X},
#                 dims=('Y', 'X')
#             )
#         kmeans = KMeans(n_clusters=self.n_clusters, random_state=42)
#         df_mean['cluster'] = kmeans.fit_predict(df_mean[['value']])
#         clusters_da = df_mean['cluster'].to_xarray()
#         valid_mask = ~np.isnan(data.isel(T=0))
#         clusters_da = clusters_da * xr.where(valid_mask, 1, np.nan)
#         _, clusters_aligned = xr.align(data, clusters_da, join='inner')
#         dist_codes = {}
#         for cl in np.unique(clusters_aligned):
#             if np.isnan(cl):
#                 continue
#             cl = int(cl)
#             cl_data = data.where(clusters_aligned == cl).values
#             cl_data = cl_data[~np.isnan(cl_data)]
#             if cl_data.size < 2:
#                 dist_codes[cl] = np.nan
#                 continue
#             try:
#                 ftr = Fitter(cl_data, distributions=dist_names, timeout=120)
#                 ftr.fit()
#                 best_name = next(iter(ftr.get_best(method='sumsquare_error')))
#                 dist_codes[cl] = self.distribution_map[best_name]
#             except (RuntimeError, ValueError):
#                 dist_codes[cl] = np.nan
#         return xr.apply_ufunc(
#             lambda x: dist_codes.get(int(x), np.nan) if not np.isnan(x) else np.nan,
#             clusters_aligned,
#             vectorize=True,
#             output_dtypes=[np.float32]
#         )

#     def plot_best_fit_map(
#         self,
#         data_array,
#         map_dict,
#         output_file='map.png',
#         title='Categorical Map',
#         colors=None,
#         figsize=(10, 6),
#         extent=None,
#         show_plot=False
#     ):
#         """
#         Plot a categorical map of distributions or skewness classes.

#         Parameters
#         ----------
#         data_array : xarray.DataArray
#             Data to plot (e.g., distribution codes or skewness classes) with dimensions (Y, X).
#         map_dict : dict
#             Mapping of category names to numeric codes.
#         output_file : str, optional
#             Path to save the plot. Default is 'map.png'.
#         title : str, optional
#             Plot title. Default is 'Categorical Map'.
#         colors : list, optional
#             Colors for each code. Default is ['blue', 'green', 'red', 'purple', 'orange'].
#         figsize : tuple, optional
#             Figure size (width, height) in inches. Default is (10, 6).
#         extent : tuple, optional
#             Map extent (lon_min, lon_max, lat_min, lat_max). Default is data bounds.
#         show_plot : bool, optional
#             If True, display the plot interactively. Default is False.

#         Raises
#         ------
#         ValueError
#             If insufficient colors are provided for the number of categories.

#         Notes
#         -----
#         Uses `cartopy` for geospatial visualization with PlateCarree projection.
#         Saves the plot as a PNG file with 300 DPI.
#         """
#         if colors is None:
#             colors = ['blue', 'green', 'red', 'purple', 'orange']
#         code2name = {v: k for k, v in map_dict.items()}
#         codes = np.unique(data_array.values[~np.isnan(data_array.values)]).astype(int)
#         if len(colors) < len(codes):
#             raise ValueError(f"Need at least {len(codes)} colors, got {len(colors)}")
#         cmap = ListedColormap([colors[i % len(colors)] for i in range(len(codes))])
#         bounds = np.concatenate([codes - 0.5, [codes[-1] + 0.5]])
#         norm = BoundaryNorm(bounds, cmap.N)
#         fig = plt.figure(figsize=figsize)
#         ax = plt.axes(projection=ccrs.PlateCarree())
#         if extent is None:
#             extent = [
#                 float(data_array.X.min()),
#                 float(data_array.X.max()),
#                 float(data_array.Y.min()),
#                 float(data_array.Y.max())
#             ]
#         ax.set_extent(extent, crs=ccrs.PlateCarree())
#         ax.add_feature(cfeature.COASTLINE, linewidth=0.4)
#         ax.add_feature(cfeature.BORDERS, linewidth=0.4)
#         ax.gridlines(draw_labels=True, linestyle='--', alpha=0.5)
#         mesh = data_array.plot.pcolormesh(
#             ax=ax,
#             transform=ccrs.PlateCarree(),
#             cmap=cmap,
#             norm=norm,
#             add_colorbar=False
#         )
#         cbar = plt.colorbar(mesh, ax=ax, ticks=codes, pad=0.05)
#         cbar.set_ticklabels([code2name.get(c, 'unknown') for c in codes])
#         cbar.set_label('Category')
#         ax.set_title(title)
#         ax.set_xlabel('Longitude')
#         ax.set_ylabel('Latitude')
#         plt.savefig(output_file, dpi=300, bbox_inches='tight')
#         if show_plot:
#             plt.show()
#         plt.close()



import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap, BoundaryNorm
from sklearn.cluster import KMeans
from scipy.stats import (
    skew, boxcox, yeojohnson,
    norm, lognorm, expon, gamma as gamma_dist, 
    weibull_min, t as t_dist, poisson, nbinom
)
import warnings

# Suppress warnings that occur during fitting (convergence warnings)
warnings.filterwarnings("ignore", category=RuntimeWarning)

def inv_boxcox(y, lmbda):
    """Inverse Box-Cox transformation."""
    if abs(lmbda) < 1e-6:
        return np.exp(y)
    return (y * lmbda + 1) ** (1 / lmbda)

def inv_yeojohnson(y, lmbda):
    """
    Inverse Yeo-Johnson transformation.
    """
    out = np.zeros_like(y)
    pos = y >= 0
    neg = ~pos

    # Case 1: y >= 0
    if abs(lmbda) < 1e-6:
        out[pos] = np.exp(y[pos]) - 1
    else:
        out[pos] = (y[pos] * lmbda + 1) ** (1 / lmbda) - 1

    # Case 2: y < 0
    if abs(lmbda - 2) < 1e-6:
        out[neg] = -np.exp(-y[neg]) + 1
    else:
        out[neg] = -((-y[neg] * (2 - lmbda) + 1) ** (1 / (2 - lmbda))) + 1
        
    return out

class WAS_TransformData:
    """
    Optimized class for geospatial Time-Series analysis: 
    Skewness detection, Transformation, and Distribution Fitting.
    """

    def __init__(self, data, distribution_map=None, n_clusters=1000):
        if not isinstance(data, xr.DataArray):
            raise ValueError("`data` must be an xarray.DataArray")
        if not all(dim in data.dims for dim in ('T', 'Y', 'X')):
            raise ValueError("`data` must have dimensions ('T', 'Y', 'X')")

        self.data = data
        self.distribution_map = distribution_map or {
            'norm': 1, 'lognorm': 2, 'expon': 3, 'gamma': 4, 
            'weibull_min': 5, "t_dist": 6, "poisson": 7, "nbinom": 8
        }
        self.n_clusters = n_clusters
        
        # Results containers
        self.transformed_data = None
        self.transform_methods = None
        self.transform_params = None # Now stores floats (lambdas), not objects
        self.skewness_ds = None
        self.handle_ds = None

    # ----------------------------------------------------------------------
    # 1. Skewness Detection (Unchanged - already efficient)
    # ----------------------------------------------------------------------
    def detect_skewness(self):
        """Compute and classify skewness for each grid cell."""
        def _compute(precip):
            precip = np.asarray(precip)
            valid = ~np.isnan(precip)
            if valid.sum() < 3:
                return np.nan, 'invalid'
            
            sk = skew(precip[valid], axis=0, nan_policy='omit')
            
            if np.isnan(sk): cls = 'invalid'
            elif -0.5 <= sk <= 0.5: cls = 'symmetric'
            elif 0.5 < sk <= 1: cls = 'moderate_positive'
            elif -1 <= sk < -0.5: cls = 'moderate_negative'
            elif sk > 1: cls = 'high_positive'
            else: cls = 'high_negative'
            
            return sk, cls

        res = xr.apply_ufunc(
            _compute, self.data,
            input_core_dims=[['T']], output_core_dims=[[], []],
            vectorize=True, dask='parallelized',
            output_dtypes=[float, str]
        )

        self.skewness_ds = xr.Dataset(
            {'skewness': res[0], 'skewness_class': res[1]}
        )
        counts = pd.Series(res[1].values.ravel()).value_counts().to_dict()
        return self.skewness_ds, {'class_counts': counts}

    # ----------------------------------------------------------------------
    # 2. Handle Skewness 
    # ----------------------------------------------------------------------
    def handle_skewness(self):
        """Recommend transformations based on skewness."""
        if self.skewness_ds is None:
            raise ValueError("Run detect_skewness() first")

        def _suggest(precip, sk_class):
            if sk_class == 'invalid': return 'none'
            valid = precip[~np.isnan(precip)]
            if len(valid) == 0: return 'none'
            
            all_pos = np.all(valid > 0)
            has_zeros = np.any(valid == 0)
            
            methods = []
            if sk_class in ('moderate_positive', 'high_positive'):
                if all_pos: methods += ['log', 'box_cox']
                methods += ['yeo_johnson', 'square_root'] # Yeo works on zeros
            elif sk_class in ('moderate_negative', 'high_negative'):
                methods += ['reflect_log'] if all_pos else []
                methods += ['reflect_yeo_johnson']
            else:
                methods.append('none')
            
            return ';'.join(methods) if methods else 'none'

        recommended = xr.apply_ufunc(
            _suggest, self.data, self.skewness_ds['skewness_class'],
            input_core_dims=[['T'], []], output_core_dims=[[]],
            vectorize=True, dask='parallelized', output_dtypes=[str]
        )

        self.handle_ds = xr.Dataset({
            'skewness': self.skewness_ds['skewness'],
            'skewness_class': self.skewness_ds['skewness_class'],
            'recommended_methods': recommended
        })
        return self.handle_ds

    # ----------------------------------------------------------------------
    # 3. Apply Transformation (Vectorized)
    # ----------------------------------------------------------------------
    def apply_transformation(self, method=None):
        """
        Apply transformations using vectorized ufuncs for speed.
        Stores parameters (lambdas) in self.transform_params.
        """
        if method is None and self.handle_ds is None:
            raise ValueError("Run handle_skewness() first or specify `method`")

        # 1. Determine method map
        if method is None:
            def extract_first(x):
                return x.split(';')[0] if (isinstance(x, str) and x) else 'none'
            method_da = xr.apply_ufunc(extract_first, self.handle_ds['recommended_methods'], vectorize=True)
        elif isinstance(method, str):
            method_da = xr.full_like(self.data.isel(T=0), method, dtype=object)
        else:
            method_da = method

        self.transform_methods = method_da

        # 2. Define the Core Transformation Logic (1D)
        def _transform_core(arr, method_str):
            # Returns: (transformed_array, parameter)
            # parameter is lambda for boxcox/yeo, or NaN for others
            
            out = arr.copy()
            param = np.nan
            
            mask = ~np.isnan(arr)
            vals = arr[mask]
            
            if len(vals) < 2 or method_str == 'none':
                return out, param

            try:
                if method_str == 'log':
                    if np.all(vals > 0): out[mask] = np.log(vals)
                
                elif method_str == 'square_root':
                    if np.all(vals >= 0): out[mask] = np.sqrt(vals)
                
                elif method_str == 'box_cox':
                    if np.all(vals > 0):
                        out[mask], param = boxcox(vals)
                
                elif method_str == 'yeo_johnson':
                    out[mask], param = yeojohnson(vals)
                
                elif method_str == 'reflect_log':
                    # Reflect: new = log(max(vals) + 1 - vals) or simple reflection -vals
                    # Assuming simple reflection based on previous code context:
                    ref = -vals
                    if np.all(ref > 0): out[mask] = np.log(ref)

                elif method_str == 'reflect_yeo_johnson':
                    out[mask], param = yeojohnson(-vals)
                    
                elif method_str == 'clipping':
                    # Example: Clip to 1st/99th percentile
                    lower, upper = np.percentile(vals, [1, 99])
                    out[mask] = np.clip(vals, lower, upper)
            
            except Exception:
                pass # Return original if fail
            
            return out, param

        # 3. Apply using xarray wrapper
        transformed, params = xr.apply_ufunc(
            _transform_core,
            self.data,
            self.transform_methods,
            input_core_dims=[['T'], []],
            output_core_dims=[['T'], []],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[float, float] 
        )

        self.transformed_data = transformed
        self.transform_params = params # Now a DataArray of floats (lambdas)
        return self.transformed_data

    # ----------------------------------------------------------------------
    # 4. Inverse Transform 
    # ----------------------------------------------------------------------

    def inverse_transform(self, data=None):
        """
        Reverse transformations to recover original scale.
        
        Parameters
        ----------
        data : xarray.DataArray, optional
            If provided (e.g., model predictions), this data will be inverted 
            using the parameters learned during the `apply_transformation` step.
            If None, the internal `self.transformed_data` is inverted.
            
            IMPORTANT: `data` must have the same Y and X coordinates/dimensions 
            as the original data so the correct lambda matches the correct pixel.
        """
        # 1. Decide what to invert
        if data is not None:
            target_data = data
        elif self.transformed_data is not None:
            target_data = self.transformed_data
        else:
            raise ValueError("No data found to inverse. Pass `data` or run apply_transformation() first.")

        if self.transform_params is None:
             raise ValueError("Transform parameters missing. Run apply_transformation() first.")

        # 2. Define Inverse Logic
        def _inv_core(vec, method, param):
            # Check for missing/invalid methods
            if method == 'none' or pd.isnull(method): 
                return vec
            
            try:
                if method == 'log': return np.exp(vec)
                if method == 'square_root': return vec ** 2
                
                if method == 'box_cox':
                    if np.isnan(param): return vec
                    return inv_boxcox(vec, param)
                
                if method == 'yeo_johnson':
                    if np.isnan(param): return vec
                    return inv_yeojohnson(vec, param)
                
                if method == 'reflect_log':
                    return -np.exp(vec)
                
                if method == 'reflect_yeo_johnson':
                    if np.isnan(param): return -vec
                    return -inv_yeojohnson(vec, param)
                    
            except Exception:
                return vec 
            return vec

        # 3. Apply
        # Note: We use xarray broadcasting. 
        # transform_params has (Y, X). target_data has (T_new, Y, X).
        # xarray automatically aligns them by Y and X.
        return xr.apply_ufunc(
            _inv_core,
            target_data,
            self.transform_methods,
            self.transform_params,
            input_core_dims=[['T'], [], []], # vectors along time
            output_core_dims=[['T']],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[float]
        )

    # ----------------------------------------------------------------------
    # 5. Unified Distribution Fitting (Grid + Cluster modes)
    # ----------------------------------------------------------------------
    def fit_best_distribution(self, use_transformed=False, mode="cluster"):
        """
        Fit best distributions using either clustering or per-grid methods.
        
        Parameters
        ----------
        mode : {'cluster', 'grid'}
        """
        # Data Selection
        data = self.transformed_data if (use_transformed and self.transformed_data is not None) else self.data
        Y, X = data.sizes["Y"], data.sizes["X"]
        coords = {"Y": data.Y, "X": data.X}

        # Map setup
        name_to_dist = {
            "norm": norm, "lognorm": lognorm, "expon": expon, 
            "gamma": gamma_dist, "weibull_min": weibull_min, 
            "t": t_dist, "t_dist": t_dist,
            "poisson": poisson, "nbinom": nbinom
        }
        dist_candidates = {
            name: (name_to_dist[name], code)
            for name, code in self.distribution_map.items() if name in name_to_dist
        }

        # --- Core Fitter Function ---
        def _fit_1d_core(sample, min_n=15):
            vals = sample[np.isfinite(sample)]
            if len(vals) < min_n:
                return np.nan, np.nan, np.nan, np.nan

            vals_pos = vals[vals > 0]
            
            best_aic = np.inf
            best_res = (np.nan, np.nan, np.nan, np.nan) # code, shape, loc, scale

            for name, (dist_obj, code) in dist_candidates.items():
                is_discrete = name in ("poisson", "nbinom")
                
                # 1. Select correct subset based on support
                if is_discrete:
                    # Enforce integer inputs for discrete distributions
                    current_sample = np.round(vals_pos)
                    # Filter out any lingering non-integers or zeros if strict
                    current_sample = current_sample[current_sample >= 0]
                elif name in ("lognorm", "gamma", "weibull_min", "expon"):
                    current_sample = vals_pos
                else:
                    current_sample = vals

                if len(current_sample) < min_n: continue

                try:
                    # 2. Fit Parameters
                    # Note: floc=0 enforces 0 bound for precip-like data
                    if name == "poisson":
                        mu = current_sample.mean()
                        if mu <= 0: continue
                        params = (mu, 0)
                        k = 1
                    elif name == "nbinom":
                        params = dist_obj.fit(current_sample, floc=0)
                        k = 2 # n, p (loc fixed)
                    elif name in ("lognorm", "gamma", "weibull_min", "expon"):
                        params = dist_obj.fit(current_sample, floc=0)
                        k = len(params) - 1
                    else:
                        params = dist_obj.fit(current_sample)
                        k = len(params)

                    # 3. Calculate AIC
                    if is_discrete:
                        if name == "poisson":
                            logL = np.sum(dist_obj.logpmf(current_sample, params[0], loc=params[1]))
                        else:
                            logL = np.sum(dist_obj.logpmf(current_sample, *params))
                    else:
                        logL = np.sum(dist_obj.logpdf(current_sample, *params))

                    aic = 2*k - 2*logL

                    if aic < best_aic:
                        best_aic = aic
                        # Standardize output to (code, shape, loc, scale)
                        # Helper to extract based on param length
                        p_len = len(params)
                        if name == "poisson":
                            # (mu, loc) -> shape=mu, loc=loc
                            best_res = (code, params[0], params[1], np.nan)
                        elif name == "nbinom":
                            # (n, p, loc) -> shape=n, loc=loc, scale=p
                            best_res = (code, params[0], params[2], params[1])
                        elif p_len == 2: # norm, expon (loc, scale)
                            best_res = (code, np.nan, params[0], params[1])
                        elif p_len == 3: # gamma, etc (shape, loc, scale)
                            best_res = (code, params[0], params[1], params[2])

                except Exception:
                    continue
            
            return best_res

        # --- MODE: Grid (Pixel by Pixel) ---
        if mode == "grid":
            res = xr.apply_ufunc(
                _fit_1d_core, data,
                input_core_dims=[['T']],
                output_core_dims=[[], [], [], []],
                vectorize=True, dask='parallelized',
                output_dtypes=[float, float, float, float]
            )
            return res[0], res[1], res[2], res[3], xr.full_like(res[0], np.nan)

        # --- MODE: Cluster (Regional Frequency Analysis) ---
        elif mode == "cluster":
            # 1. KMeans on Mean/Std
            df = data.to_dataframe().reset_index().dropna().drop(columns=['T'])
            if df.empty or len(df) < self.n_clusters:
                print("Not enough data for clustering.")
                nan_da = xr.full_like(data.isel(T=0), np.nan)
                return nan_da, nan_da, nan_da, nan_da, nan_da

            var_col = df.columns[-1] # Assuming data column is last
            kmeans = KMeans(n_clusters=self.n_clusters, n_init=10, random_state=42)
            df['cluster'] = kmeans.fit_predict(df[[var_col]])

            # Map clusters back to grid
            cluster_da = df.set_index(['Y', 'X'])['cluster'].to_xarray().reindex_like(data.isel(T=0))

            # 2. Loop over clusters (fast, N=1000 not N=1Million)
            best_code = xr.full_like(cluster_da, np.nan, dtype=float)
            best_shape = xr.full_like(cluster_da, np.nan, dtype=float)
            best_loc = xr.full_like(cluster_da, np.nan, dtype=float)
            best_scale = xr.full_like(cluster_da, np.nan, dtype=float)

            unique_clusters = np.unique(df['cluster'].values)
            
            for cl in unique_clusters:
                # Mask: All time steps for all pixels in this cluster
                mask = (cluster_da == cl)
                # Extract pooled data
                pooled = data.where(mask).values.flatten()
                
                # Fit
                c, s, l, sc = _fit_1d_core(pooled, min_n=30)
                
                # Broadcast results back to map
                if not np.isnan(c):
                    best_code = best_code.where(~mask, c)
                    best_shape = best_shape.where(~mask, s)
                    best_loc = best_loc.where(~mask, l)
                    best_scale = best_scale.where(~mask, sc)
            
            return best_code, best_shape, best_loc, best_scale, cluster_da

        else:
            raise ValueError("mode must be 'grid' or 'cluster'")

    # ----------------------------------------------------------------------
    # 6. Plotting
    # ----------------------------------------------------------------------
    # def plot_best_fit_map_(self, data_array, map_dict, output_file='map.png', 
    #                       title='Map', colors=None, show_plot=True):
    #     """Plot categorical map using Cartopy."""
    #     if colors is None:
    #         colors = ['#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', 
    #                   '#fb9a99', '#e31a1c', '#fdbf6f', '#ff7f00']
            
    #     # Filter map_dict to only what's in the data
    #     unique_vals = np.unique(data_array.values[~np.isnan(data_array.values)])
    #     if len(unique_vals) == 0:
    #         print("No data to plot.")
    #         return

    #     # Ensure we have enough colors
    #     if len(colors) < len(unique_vals):
    #         colors = colors * (len(unique_vals) // len(colors) + 1)
            
    #     cmap = ListedColormap(colors[:len(unique_vals)])
    #     bounds = np.arange(len(unique_vals) + 1)
    #     norm = BoundaryNorm(bounds, cmap.N)
        
    #     # Create lookup for legend
    #     val_to_idx = {v: i for i, v in enumerate(unique_vals)}
    #     # Code to Name
    #     code_to_name = {v: k for k, v in map_dict.items()}
        
    #     # Transform data to indices for plotting
    #     plot_data = xr.apply_ufunc(lambda x: val_to_idx.get(x, np.nan), data_array, vectorize=True)

    #     fig = plt.figure(figsize=(10, 6))
    #     ax = plt.axes(projection=ccrs.PlateCarree())
        
    #     extent = [float(data_array.X.min()), float(data_array.X.max()), 
    #               float(data_array.Y.min()), float(data_array.Y.max())]
    #     ax.set_extent(extent, crs=ccrs.PlateCarree())
        
    #     ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
    #     ax.add_feature(cfeature.BORDERS, linewidth=0.5)
    #     ax.gridlines(draw_labels=True, linestyle='--', alpha=0.5)

    #     mesh = plot_data.plot.pcolormesh(
    #         ax=ax, transform=ccrs.PlateCarree(),
    #         cmap=cmap, norm=norm, add_colorbar=False
    #     )
        
    #     # Custom colorbar
    #     cbar = plt.colorbar(mesh, ax=ax, ticks=bounds[:-1] + 0.5, pad=0.05)
    #     cbar.set_ticklabels([code_to_name.get(v, str(v)) for v in unique_vals])
        
    #     ax.set_title(title)
    #     plt.savefig(output_file, dpi=300, bbox_inches='tight')
    #     if show_plot:
    #         plt.show()
    #     plt.close()

    def plot_best_fit_map(self, data_array, map_dict, output_file='map.png', 
                              title='Map', colors=None, show_plot=True):
            if colors is None:
                colors = ['#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', 
                          '#fb9a99', '#e31a1c', '#fdbf6f', '#ff7f00']
                
            unique_vals = np.unique(data_array.values[~np.isnan(data_array.values)])
            if len(unique_vals) == 0:
                print("No data to plot.")
                return
    
            if len(colors) < len(unique_vals):
                colors = colors * (len(unique_vals) // len(colors) + 1)
                
            cmap = ListedColormap(colors[:len(unique_vals)])
            bounds = np.arange(len(unique_vals) + 1)
            norm = BoundaryNorm(bounds, cmap.N)
            
            val_to_idx = {v: i for i, v in enumerate(unique_vals)}
            code_to_name = {v: k for k, v in map_dict.items()}
            
            plot_data = xr.full_like(data_array, fill_value=np.nan, dtype=float)
            
            for val, idx in val_to_idx.items():
                plot_data = xr.where(data_array == val, idx, plot_data)
    
            fig = plt.figure(figsize=(10, 6))
            ax = plt.axes(projection=ccrs.PlateCarree())
            
            extent = [float(data_array.X.min()), float(data_array.X.max()), 
                      float(data_array.Y.min()), float(data_array.Y.max())]
            ax.set_extent(extent, crs=ccrs.PlateCarree())
            
            ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
            ax.add_feature(cfeature.BORDERS, linewidth=0.5)
            ax.gridlines(draw_labels=True, linestyle='--', alpha=0.5)
    
            mesh = plot_data.plot.pcolormesh(
                ax=ax, transform=ccrs.PlateCarree(),
                cmap=cmap, norm=norm, add_colorbar=False
            )
            
            cbar = plt.colorbar(mesh, ax=ax, ticks=bounds[:-1] + 0.5, pad=0.05)
            cbar.set_ticklabels([code_to_name.get(v, str(v)) for v in unique_vals])
            
            ax.set_title(title)
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            if show_plot:
                plt.show()
            plt.close()