"""
Soft-Target Distribution Regressor (Single LGBM)

A scikit-learn-compatible regressor that predicts full probability distributions.
It uses a single LightGBM model trained on an expanded dataset (X crossed with Grid)
using Soft Targets generated by a Gaussian kernel around the true y.

Two-stage approach:
- Stage 1: A base LightGBM regressor provides an RMSE-optimal point estimate.
- Stage 2: The soft-target model learns the residual distribution around that estimate,
  using K-fold cross-validated residuals to avoid in-sample bias.

This approach allows for:
1. Arbitrary distribution shapes (multimodal, skewed, etc.)
2. RMSE-competitive point predictions (anchored by the base model)
3. Calibrated uncertainty via honest CV residuals
4. Single-model efficiency for the distributional component
"""

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from sklearn.model_selection import KFold
from lightgbm import LGBMRegressor
from scipy.ndimage import gaussian_filter1d

class DistributionRegressorSoftTargetV3(BaseEstimator, RegressorMixin):
    """
    Predicts probability distributions using a two-stage approach:
    1. A base LightGBM regressor for RMSE-optimal point predictions.
    2. A soft-target model that learns the residual distribution using
       K-fold cross-validated residuals.

    Parameters
    ----------
    n_bins : int, default=50
        Number of grid points to discretize the residual range.
        Higher = higher resolution but more memory usage (RAM usage is ~ n_samples * n_bins).

    sigma : float or str, default='auto'
        The standard deviation of the Gaussian kernel used to generate soft targets.
        Operates in residual space (not raw target space).
        - If float: Constant spread around the true residual.
        - If 'auto': Data-driven estimation using 3-fold CV on residuals.
        Controls the "dragging" effect:
        - Small sigma: Target looks like a sharp spike (One-Hot).
        - Large sigma: Target looks like a wide bell curve.

    output_smoothing : float, default=1.0
        Standard deviation for Gaussian smoothing of the output distribution.
        Applies a Gaussian filter to smooth the predicted probability distribution,
        which helps reduce jaggedness in the output PDF. Set to 0.0 to disable.
        Higher values produce smoother distributions.

    n_estimators : int, default=100
        Number of boosting trees.

    learning_rate : float, default=0.1
        Boosting learning rate.

    random_state : int or None, default=None
        Random seed.

    **kwargs : dict
        Additional parameters passed to LGBMRegressor.
    """

    def __init__(
        self,
        n_bins=50,
        sigma='auto',
        output_smoothing=1.0,
        n_estimators=100,
        learning_rate=0.1,
        k_mc=5,
        random_state=None,
        **kwargs
    ):
        self.n_bins = n_bins
        self.sigma = sigma
        self.output_smoothing = output_smoothing
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.k_mc = k_mc
        self.random_state = random_state
        self.lgbm_kwargs = kwargs

    def _validate_params(self):
        if self.n_bins < 2:
            raise ValueError("n_bins must be >= 2")

    def _get_base_lgbm_params(self):
        """Get parameters for the base LGBM model."""
        params = {
            'n_estimators': self.n_estimators,
            'learning_rate': self.learning_rate,
            'random_state': self.random_state,
            'verbose': -1
        }
        params.update(self.lgbm_kwargs)
        return params

    def fit(self, X, y, sample_weight=None):
        """
        Fit the model in two stages:
        1. Train a base LightGBM for point estimates + compute 5-fold CV residuals.
        2. Train the soft-target distribution model on those residuals.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values.
        sample_weight : array-like of shape (n_samples,), optional
            Individual weights for each sample.
        """
        self._validate_params()

        # 1. Prepare Data
        self._is_dataframe = isinstance(X, pd.DataFrame)
        if self._is_dataframe:
            self.feature_names_in_ = X.columns.tolist()
            X_orig = X  # Keep original DataFrame for model fitting
            X_array = X.values
            y_array = np.asarray(y)
        else:
            X_orig = X
            X_array = X
            y_array = y

        X_array, y_array = check_X_y(X_array, y_array, accept_sparse=False, dtype=np.float64)
        self.n_features_in_ = X_array.shape[1]
        n_samples = X_array.shape[0]

        # Validate sample_weight
        if sample_weight is not None:
            sample_weight = np.asarray(sample_weight, dtype=np.float64)
            if sample_weight.shape[0] != n_samples:
                raise ValueError(f"sample_weight has {sample_weight.shape[0]} samples, expected {n_samples}")

        # =====================================================
        # Stage 1: Base LGBM + 5-fold CV residuals
        # =====================================================
        base_params = self._get_base_lgbm_params()

        # Train base model on full data (used at test time)
        self.base_model_ = LGBMRegressor(**base_params)
        self.base_model_.fit(X_orig, y_array, sample_weight=sample_weight)

        # 5-fold CV for honest out-of-fold residuals
        kf = KFold(n_splits=5, shuffle=True, random_state=self.random_state)
        oof_predictions = np.zeros(n_samples)
        for train_idx, val_idx in kf.split(X_array):
            fold_model = LGBMRegressor(**base_params)
            fold_sw = sample_weight[train_idx] if sample_weight is not None else None
            if self._is_dataframe:
                fold_model.fit(X_orig.iloc[train_idx], y_array[train_idx], sample_weight=fold_sw)
                oof_predictions[val_idx] = fold_model.predict(X_orig.iloc[val_idx])
            else:
                fold_model.fit(X_array[train_idx], y_array[train_idx], sample_weight=fold_sw)
                oof_predictions[val_idx] = fold_model.predict(X_array[val_idx])

        residuals = y_array - oof_predictions

        # =====================================================
        # Stage 2: Build residual grid
        # =====================================================
        r_min = float(np.min(residuals))
        r_max = float(np.max(residuals))
        margin = (r_max - r_min) * 0.1  # 10% margin for test-time safety
        self.grid_ = np.linspace(r_min - margin, r_max + margin, self.n_bins)

        # =====================================================
        # Stage 3: Resolve Sigma (using 3-fold CV on residuals)
        # =====================================================
        if self.sigma == 'auto':
            # 3-fold CV to get honest "residuals of residuals"
            sigma_base_params = {**base_params, 'n_estimators': min(self.n_estimators, 100)}
            kf_sigma = KFold(n_splits=3, shuffle=True, random_state=self.random_state)
            oof_resid_pred = np.zeros(n_samples)

            for train_idx, val_idx in kf_sigma.split(X_array):
                fold_model = LGBMRegressor(**sigma_base_params)
                fold_sw = sample_weight[train_idx] if sample_weight is not None else None
                if self._is_dataframe:
                    fold_model.fit(X_orig.iloc[train_idx], residuals[train_idx], sample_weight=fold_sw)
                    oof_resid_pred[val_idx] = fold_model.predict(X_orig.iloc[val_idx])
                else:
                    fold_model.fit(X_array[train_idx], residuals[train_idx], sample_weight=fold_sw)
                    oof_resid_pred[val_idx] = fold_model.predict(X_array[val_idx])

            # Second-order residuals (honest noise estimate)
            resid_of_resid = residuals - oof_resid_pred

            # Heteroscedastic sigma model
            log_resid_sq = np.log(resid_of_resid ** 2 + 1e-9)
            sigma_model = LGBMRegressor(**sigma_base_params)
            sigma_model.fit(X_orig, log_resid_sq, sample_weight=sample_weight)
            log_sigma_sq_pred = sigma_model.predict(X_orig)
            sigma_pred = np.sqrt(np.exp(log_sigma_sq_pred))

            grid_step = (self.grid_[-1] - self.grid_[0]) / (self.n_bins - 1)
            min_sigma = grid_step * 0.5
            self.sigma_val_ = np.maximum(sigma_pred, min_sigma)
        else:
            self.sigma_val_ = float(self.sigma)

        # =====================================================
        # Stage 4: K Monte Carlo Expansion (on residuals)
        # Layout: [1 exact peak] [2 Gaussian slope] [K-3 uniform tails]
        # =====================================================
        K = self.k_mc
        n_total_rows = n_samples * K
        X_final = np.empty((n_total_rows, self.n_features_in_ + 1), dtype=X_array.dtype)

        # Fill feature columns using broadcasting
        X_final[:, :-1].reshape(n_samples, K, -1)[:] = X_array[:, None, :]

        # Generate targeted evaluation points
        rng = np.random.default_rng(self.random_state)
        g_points = np.empty((n_samples, K), dtype=X_array.dtype)

        sig_b = self.sigma_val_ if np.ndim(self.sigma_val_) == 0 else self.sigma_val_[:, None]
        g_points[:, 0] = residuals                                                              # Exact peak
        g_points[:, 1:3] = residuals[:, None] + rng.standard_normal((n_samples, 2)) * sig_b     # Gaussian slope
        g_points[:, 3:self.k_mc] = rng.uniform(r_min - margin, r_max + margin, size=(n_samples, K - 3))  # Uniform tails

        X_final[:, -1] = g_points.ravel()

        # Soft targets: Gaussian around true RESIDUAL
        diff_sq = (np.repeat(residuals, K) - X_final[:, -1]) ** 2

        if np.ndim(self.sigma_val_) == 0:
            sigma_sq = self.sigma_val_ ** 2
        else:
            sigma_sq = np.repeat(self.sigma_val_, K) ** 2

        y_targets_soft = np.exp(-diff_sq / (2 * sigma_sq))
        y_targets_soft[y_targets_soft < 1e-5] = 0.0

        # Feature Names
        if self._is_dataframe:
            feature_names = self.feature_names_in_ + ['grid_point']
        else:
            feature_names = [f"feature_{i}" for i in range(self.n_features_in_)] + ['grid_point']

        # =====================================================
        # Stage 5: Train Soft-Target LightGBM (on residuals)
        # =====================================================
        params = {
            'objective': 'cross_entropy',
            'metric': 'cross_entropy',
            'n_estimators': self.n_estimators,
            'learning_rate': self.learning_rate,
            'random_state': self.random_state,
            'verbose': -1
        }
        params.update(self.lgbm_kwargs)

        self.model_ = LGBMRegressor(**params)

        # Expand sample_weight to match expanded dataset
        if sample_weight is not None:
            sample_weight_expanded = np.repeat(sample_weight, K)
        else:
            sample_weight_expanded = None

        self.model_.fit(X_final, y_targets_soft, sample_weight=sample_weight_expanded, feature_name=feature_names)

        return self

    def _predict_residual_distribution(self, X):
        """
        Internal: predict the residual distribution (before shifting by base model).

        Returns
        -------
        grid : array of shape (n_bins,)
            Residual grid points.
        distributions : array of shape (n_samples, n_bins)
            Probability distribution over residuals.
        """
        check_is_fitted(self)

        if isinstance(X, pd.DataFrame):
            X_array = X.values
        else:
            X_array = X
        X_array = check_array(X_array, accept_sparse=False)
        n_samples = X_array.shape[0]

        # Expand X for prediction
        X_expanded = np.repeat(X_array, self.n_bins, axis=0)
        grid_tile = np.tile(self.grid_, n_samples)

        if self._is_dataframe:
            feature_names = self.feature_names_in_
        else:
            feature_names = [f"feature_{i}" for i in range(self.n_features_in_)]

        X_df_expanded = pd.DataFrame(X_expanded, columns=feature_names)
        X_df_expanded['grid_point'] = grid_tile

        pred_scores = self.model_.predict(X_df_expanded)
        scores_matrix = pred_scores.reshape(n_samples, self.n_bins)

        # Normalize to valid probability distribution
        row_sums = scores_matrix.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1.0
        distributions = scores_matrix / row_sums

        # Apply output smoothing
        if self.output_smoothing > 0:
            distributions = gaussian_filter1d(distributions, sigma=self.output_smoothing, axis=1)
            row_sums = distributions.sum(axis=1, keepdims=True)
            row_sums[row_sums == 0] = 1.0
            distributions = distributions / row_sums

        return self.grid_, distributions

    def _get_base_predictions(self, X):
        """Get base LGBM point predictions for X."""
        return self.base_model_.predict(X)

    def predict_distribution(self, X):
        """
        Returns grid points and probability distribution for each sample.
        The grid is shifted per sample by the base model prediction (absolute scale).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples for which to predict distributions.

        Returns
        -------
        grids : array of shape (n_samples, n_bins)
            Per-sample grid points (residual grid shifted by base prediction).

        distributions : array of shape (n_samples, n_bins)
            Probability distribution for each sample at each grid point.
        """
        resid_grid, dists = self._predict_residual_distribution(X)
        base_pred = self._get_base_predictions(X)
        shifted_grids = resid_grid[None, :] + base_pred[:, None]
        return shifted_grids, dists

    def predict(self, X):
        """Default: Predict Mean (base LGBM + expected residual)."""
        return self.predict_mean(X)

    def predict_mean(self, X):
        """
        Predict the mean of the distribution for each sample.
        Returns base_prediction + E[residual_distribution].

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        Returns
        -------
        means : array of shape (n_samples,)
            Predicted means.
        """
        resid_grid, dists = self._predict_residual_distribution(X)
        residual_mean = np.sum(dists * resid_grid, axis=1)
        base_pred = self._get_base_predictions(X)
        return base_pred + residual_mean

    def predict_mode(self, X):
        """
        Predict the mode (peak) of the distribution for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        Returns
        -------
        modes : array of shape (n_samples,)
            Predicted modes.
        """
        resid_grid, dists = self._predict_residual_distribution(X)
        max_indices = np.argmax(dists, axis=1)
        residual_modes = resid_grid[max_indices]
        base_pred = self._get_base_predictions(X)
        return base_pred + residual_modes

    def predict_quantile(self, X, q=0.5):
        """
        Predict quantile(s) of the distribution for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        q : float or array-like, default=0.5
            Quantile(s) to compute, in [0, 1].

        Returns
        -------
        quantiles : array of shape (n_samples,) or (n_samples, n_quantiles)
            Predicted quantiles.
        """
        resid_grid, dists = self._predict_residual_distribution(X)
        cdfs = np.cumsum(dists, axis=1)

        q_arr = np.asarray(q)
        is_scalar = q_arr.ndim == 0
        if is_scalar:
            q_arr = q_arr.reshape(1)

        n_samples = dists.shape[0]
        n_quantiles = len(q_arr)
        quantiles = np.zeros((n_samples, n_quantiles))

        for i in range(n_samples):
            indices = np.searchsorted(cdfs[i], q_arr)
            indices = np.clip(indices, 0, self.n_bins - 1)
            quantiles[i] = resid_grid[indices]

        # Shift from residual space to absolute space
        base_pred = self._get_base_predictions(X)
        quantiles += base_pred[:, None]

        if is_scalar:
            return quantiles.flatten()
        return quantiles

    def predict_interval(self, X, confidence=0.95):
        """
        Predict prediction interval for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        confidence : float, default=0.95
            Confidence level (e.g., 0.95 for 95% interval).

        Returns
        -------
        intervals : array of shape (n_samples, 2)
            Lower and upper bounds of the interval.
        """
        alpha = 1.0 - confidence
        q_low = alpha / 2.0
        q_high = 1.0 - alpha / 2.0

        return self.predict_quantile(X, q=[q_low, q_high])

    def predict_std(self, X):
        """
        Predict standard deviation of the distribution for each sample.
        (Shift-invariant: same in residual or absolute space.)

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        Returns
        -------
        stds : array of shape (n_samples,)
            Predicted standard deviations.
        """
        resid_grid, dists = self._predict_residual_distribution(X)
        means = np.sum(dists * resid_grid, axis=1, keepdims=True)
        variance = np.sum(dists * (resid_grid - means)**2, axis=1)
        return np.sqrt(variance)
