"""
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.

This approach allows for:
1. Arbitrary distribution shapes (multimodal, skewed, etc.)
2. "Dragging" behavior via the sigma parameter (smoothing neighbors)
3. Single-model efficiency
"""

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

class DistributionRegressorSoftTarget(BaseEstimator, RegressorMixin):
    """
    Predicts probability distributions using a single gradient boosted model
    trained on soft targets (Gaussian kernel smoothing).

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

    sigma : float or str, default=1.0
        The standard deviation of the Gaussian kernel used to generate soft targets.
        - If float: Constant spread around the true y.
        - If 'auto': Data-driven estimation based on residual standard deviation from
          a preliminary LGBMRegressor fit. This reflects the inherent noise in the data.
          Set to max(std(residuals), grid_step) to ensure at least one grid interval.
        Controls the "dragging" effect:
        - Small sigma: Target looks like a sharp spike (One-Hot).
        - Large sigma: Target looks like a wide bell curve.

    use_base_model : bool, default=False
        If True, trains a base LGBMRegressor for point predictions and builds
        the soft-target model on CV residuals. This anchors point predictions
        to RMSE-optimal estimates and learns the residual distribution.
        If False, the soft-target model is trained directly on raw y.

    monte_carlo_training : bool, default=True
        If True, use Monte Carlo sampling for training expansion (K points per
        sample instead of the full grid). Much more memory-efficient.
        If False, use full grid expansion (each sample crossed with all n_bins
        grid points).

    mc_samples : int, default=5
        Number of Monte Carlo sample points per observation when
        monte_carlo_training=True. Ignored when monte_carlo_training=False.

    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',
        use_base_model=False,
        n_estimators=100,
        learning_rate=0.1,
        monte_carlo_training=True,
        mc_samples=5,
        random_state=None,
        **kwargs
    ):
        self.n_bins = n_bins
        self.sigma = sigma
        self.use_base_model = use_base_model
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.monte_carlo_training = monte_carlo_training
        self.mc_samples = mc_samples
        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 _to_dataframe(self, X):
        """Convert X to DataFrame if it isn't one already, using stored feature names."""
        if isinstance(X, pd.DataFrame):
            return X
        return pd.DataFrame(X, columns=self.feature_names_in_)

    def _expand_with_grid_points(self, X_df, grid_points, K):
        """
        Repeat each row of X_df K times and append grid_point column.
        Preserves original dtypes.
        """
        idx = np.repeat(np.arange(len(X_df)), K)
        X_expanded = X_df.iloc[idx].reset_index(drop=True)
        X_expanded['grid_point'] = grid_points
        return X_expanded

    def _per_sample_grids(self, base_preds):
        """Compute per-sample residual grids bounded to [y_min, y_max] in absolute space."""
        r_min = np.maximum(self.grid_[0], self.y_min_ - base_preds)
        r_max = np.minimum(self.grid_[-1], self.y_max_ - base_preds)
        t = np.linspace(0, 1, self.n_bins)
        return r_min[:, None] + t[None, :] * (r_max - r_min)[:, None]

    def fit(self, X, y, sample_weight=None):
        """
        Fit the model by expanding the dataset and training on soft targets.

        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 — preserve original dtypes throughout
        self._is_dataframe = isinstance(X, pd.DataFrame)
        if self._is_dataframe:
            self.feature_names_in_ = X.columns.tolist()
        else:
            self.feature_names_in_ = [f"feature_{i}" for i in range(np.asarray(X).shape[1])]

        X_df = self._to_dataframe(X)
        y_array = np.asarray(y, dtype=np.float64)
        self.y_min_ = float(np.min(y_array))
        self.y_max_ = float(np.max(y_array))
        self.n_features_in_ = X_df.shape[1]
        n_samples = X_df.shape[0]

        if sample_weight is not None:
            sample_weight = np.asarray(sample_weight, dtype=np.float64)

        # Full params for the base model (same as user-specified)
        full_params = {
            'n_estimators': self.n_estimators,
            'learning_rate': self.learning_rate,
            'random_state': self.random_state,
            'verbose': -1
        }
        full_params.update(self.lgbm_kwargs)

        baseline_params = {**full_params, 'n_estimators': min(self.n_estimators, 100)}
        cv_params = {'sample_weight': sample_weight} if sample_weight is not None else None

        # 2. Base model + OOF residuals (optional)
        if self.use_base_model:
            self.base_model_ = LGBMRegressor(**full_params)
            self.base_model_.fit(X_df, y_array, sample_weight=sample_weight)

            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_df):
                fold_model = LGBMRegressor(**full_params)
                fold_sw = sample_weight[train_idx] if sample_weight is not None else None
                fold_model.fit(X_df.iloc[train_idx], y_array[train_idx], sample_weight=fold_sw)
                oof_predictions[val_idx] = fold_model.predict(X_df.iloc[val_idx])

            target_for_grid = y_array - oof_predictions
        else:
            self.base_model_ = None
            target_for_grid = y_array

        # 3. Define the Grid
        t_min = float(np.min(target_for_grid))
        t_max = float(np.max(target_for_grid))
        self.grid_ = np.linspace(t_min, t_max, self.n_bins)

        # 4. Resolve Sigma
        if self.sigma == 'auto':
            if self.use_base_model:
                residuals = target_for_grid
            else:
                baseline_model = LGBMRegressor(**baseline_params)
                y_pred_cv = cross_val_predict(baseline_model, X_df, y_array, cv=5, params=cv_params)
                residuals = y_array - y_pred_cv

            grid_step = (self.grid_[-1] - self.grid_[0]) / (self.n_bins - 1)

            log_resid_sq = np.log((residuals ** 2) + 1e-9)
            sigma_model = LGBMRegressor(**baseline_params)
            log_sigma_sq_pred = cross_val_predict(sigma_model, X_df, log_resid_sq, cv=5, params=cv_params)

            sigma_pred = np.sqrt(np.exp(log_sigma_sq_pred))
            self.sigma_val_ = np.maximum(sigma_pred, grid_step * 0.5)
        else:
            self.sigma_val_ = float(self.sigma)

        # ------------------------------------------------------------------
        # 5. Training Expansion
        # ------------------------------------------------------------------
        if self.use_base_model:
            # Per-sample bounded residual ranges
            r_min = np.maximum(t_min, self.y_min_ - oof_predictions)
            r_max = np.minimum(t_max, self.y_max_ - oof_predictions)

        if self.monte_carlo_training:
            K = self.mc_samples
            rng = np.random.default_rng(self.random_state)
            g_points = np.empty((n_samples, K))

            # Point 0: exact peak
            g_points[:, 0] = target_for_grid

            # Gaussian slope points
            n_gauss = max((K - 1) // 2, 1)
            sig_b = self.sigma_val_ if np.ndim(self.sigma_val_) == 0 else self.sigma_val_[:, None]
            g_points[:, 1:1+n_gauss] = target_for_grid[:, None] + rng.standard_normal((n_samples, n_gauss)) * sig_b

            # Uniform tail points
            n_uniform = K - 1 - n_gauss
            if n_uniform > 0:
                if self.use_base_model:
                    g_points[:, 1+n_gauss:] = r_min[:, None] + rng.uniform(0, 1, size=(n_samples, n_uniform)) * (r_max - r_min)[:, None]
                else:
                    g_points[:, 1+n_gauss:] = rng.uniform(t_min, t_max, size=(n_samples, n_uniform))

            # Clip MC points to valid per-sample range
            if self.use_base_model:
                g_points = np.clip(g_points, r_min[:, None], r_max[:, None])

            grid_points_flat = g_points.ravel()

            # Importance weights
            domain_width = t_max - t_min
            iw = np.empty((n_samples, K))
            iw[:, 0] = 1.0 / K

            gauss_points = g_points[:, 1:1+n_gauss]
            sig_arr = self.sigma_val_ if np.ndim(self.sigma_val_) == 0 else self.sigma_val_[:, None]
            gauss_pdf = np.exp(-((gauss_points - target_for_grid[:, None]) ** 2) / (2 * sig_arr ** 2)) / (sig_arr * np.sqrt(2 * np.pi))
            iw[:, 1:1+n_gauss] = (1.0 / domain_width) / (gauss_pdf + 1e-12)

            if n_uniform > 0:
                iw[:, 1+n_gauss:] = 1.0

            iw /= iw.mean()
            mc_weights = iw.ravel()
        else:
            K = self.n_bins
            if self.use_base_model:
                # Per-sample grids bounded to [y_min, y_max]
                per_sample_grids = self._per_sample_grids(oof_predictions)
                grid_points_flat = per_sample_grids.ravel()
            else:
                grid_points_flat = np.tile(self.grid_, n_samples)

        # Build expanded DataFrame (preserves original dtypes)
        X_expanded = self._expand_with_grid_points(X_df, grid_points_flat, K)

        # 6. Generate Soft Targets
        diff_sq = (np.repeat(target_for_grid, K) - grid_points_flat) ** 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

        # 7. Configure and Train LightGBM
        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)

        if self.monte_carlo_training:
            sample_weight_expanded = mc_weights
            if sample_weight is not None:
                sample_weight_expanded = sample_weight_expanded * np.repeat(sample_weight, K)
        else:
            sample_weight_expanded = np.repeat(sample_weight, K) if sample_weight is not None else None

        self.model_.fit(X_expanded, y_targets_soft, sample_weight=sample_weight_expanded)

        return self

    def _predict_internal_distribution(self, X, output_smoothing=0):
        """
        Predict distribution on per-sample grids.

        Returns
        -------
        grids : array of shape (n_samples, n_bins)
            Per-sample grid points (residual space if base model is used).
        distributions : array of shape (n_samples, n_bins)
            Probability distribution for each sample.
        """
        check_is_fitted(self)

        X_df = self._to_dataframe(X)
        n_samples = X_df.shape[0]

        if self.base_model_ is not None:
            base_preds = self.base_model_.predict(X_df)
            grids = self._per_sample_grids(base_preds)
        else:
            grids = np.tile(self.grid_, (n_samples, 1))

        grid_points_flat = grids.ravel()
        X_expanded = self._expand_with_grid_points(X_df, grid_points_flat, self.n_bins)

        pred_scores = self.model_.predict(X_expanded)
        distributions = pred_scores.reshape(n_samples, self.n_bins)

        row_sums = distributions.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1.0
        distributions /= row_sums

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

        return grids, distributions

    def predict_distribution(self, X, output_smoothing=0):
        """
        Returns grid points and probability distribution for each sample.

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

        output_smoothing : float, default=0
            Standard deviation for Gaussian smoothing of the output distribution.
            Helps reduce jaggedness in the predicted PDF. Set to 0 to disable.

        Returns
        -------
        grids : array of shape (n_samples, n_bins)
            Per-sample grid points in absolute space.

        distributions : array of shape (n_samples, n_bins)
            Probability distribution for each sample at each grid point.

        base_offsets : array of shape (n_samples,)
            Per-sample shift from the base model prediction.
            Zeros when use_base_model=False.
        """
        grids, dists = self._predict_internal_distribution(X, output_smoothing)
        base_offsets = self.base_model_.predict(X) if self.base_model_ is not None else np.zeros(len(X))
        # Convert residual grids to absolute space
        absolute_grids = grids + base_offsets[:, None]
        return absolute_grids, dists, base_offsets

    def predict(self, X):
        """Default: Predict Mean"""
        return self.predict_mean(X)

    def predict_mean(self, X):
        """
        Predict the mean of the distribution for each sample.
        When use_base_model=True, returns the base model's prediction directly.

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

        Returns
        -------
        means : array of shape (n_samples,)
            Predicted means.
        """
        if self.base_model_ is not None:
            return self.base_model_.predict(X)
        grids, dists = self._predict_internal_distribution(X)
        return np.sum(grids * dists, axis=1)

    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.
        """
        grids, dists = self._predict_internal_distribution(X)
        max_indices = np.argmax(dists, axis=1)
        modes = grids[np.arange(len(grids)), max_indices]
        if self.base_model_ is not None:
            modes += self.base_model_.predict(X)
        return 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.
        """
        grids, dists = self._predict_internal_distribution(X)
        cdfs = np.cumsum(dists, axis=1)
        cdfs[:, -1] = 1.0

        q_arr = np.atleast_1d(q)
        is_scalar = np.ndim(q) == 0

        mask = cdfs[:, None, :] >= q_arr[None, :, None]
        indices = mask.argmax(axis=2)
        rows = np.arange(len(grids))[:, None]
        quantiles = grids[rows, indices]

        if self.base_model_ is not None:
            base_pred = self.base_model_.predict(X)
            quantiles = 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 whether 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.
        """
        grids, dists = self._predict_internal_distribution(X)
        means = np.sum(grids * dists, axis=1)
        expected_sq = np.sum(grids ** 2 * dists, axis=1)
        variance = expected_sq - means ** 2
        return np.sqrt(np.maximum(variance, 0.0))

    def pdf(self, X, y, output_smoothing=0, eps=1e-10):
        """
        Evaluate the predicted probability density at given y values.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.
        y : array-like of shape (n_samples,)
            Target values at which to evaluate the density.
        output_smoothing : float, default=0
            Standard deviation for Gaussian smoothing of the distribution.
        eps : float, default=1e-10
            Floor value to avoid zero densities.

        Returns
        -------
        densities : array of shape (n_samples,)
            Predicted density f(y|X) for each sample.
        """
        grids, dists = self._predict_internal_distribution(X, output_smoothing)
        y_array = np.asarray(y, dtype=np.float64)

        # Per-sample bin widths
        bin_widths = grids[:, 1] - grids[:, 0]
        density_grid = dists / bin_widths[:, None]

        # Shift y into residual space if base model is used
        if self.base_model_ is not None:
            y_residual = y_array - self.base_model_.predict(X)
        else:
            y_residual = y_array

        # Per-sample linear interpolation
        y_clipped = np.clip(y_residual, grids[:, 0], grids[:, -1])
        frac_idx = (y_clipped - grids[:, 0]) / bin_widths
        idx_lo = np.floor(frac_idx).astype(int)
        idx_hi = np.minimum(idx_lo + 1, self.n_bins - 1)
        alpha = frac_idx - idx_lo

        rows = np.arange(len(y_array))
        densities = (1 - alpha) * density_grid[rows, idx_lo] + alpha * density_grid[rows, idx_hi]
        return np.maximum(densities, eps)

    def logpdf(self, X, y, output_smoothing=0, eps=1e-10):
        """
        Evaluate the log probability density at given y values.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.
        y : array-like of shape (n_samples,)
            Target values at which to evaluate the log-density.
        output_smoothing : float, default=0
            Standard deviation for Gaussian smoothing of the distribution.
        eps : float, default=1e-10
            Floor value to avoid log(0).

        Returns
        -------
        log_densities : array of shape (n_samples,)
            Log predicted density log f(y|X) for each sample.
        """
        return np.log(self.pdf(X, y, output_smoothing=output_smoothing, eps=eps))

    def cdf(self, X, y, output_smoothing=0):
        """
        Evaluate the predicted cumulative distribution function at given y values.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.
        y : array-like of shape (n_samples,)
            Target values at which to evaluate the CDF.
        output_smoothing : float, default=0
            Standard deviation for Gaussian smoothing of the distribution.

        Returns
        -------
        probabilities : array of shape (n_samples,)
            Predicted P(Y <= y | X) for each sample.
        """
        grids, dists = self._predict_internal_distribution(X, output_smoothing)
        y_array = np.asarray(y, dtype=np.float64)

        # Shift y into residual space if base model is used
        if self.base_model_ is not None:
            y_residual = y_array - self.base_model_.predict(X)
        else:
            y_residual = y_array

        # Cumulative sums of probability masses
        cumsum = np.cumsum(dists, axis=1)
        bin_widths = grids[:, 1] - grids[:, 0]

        # Per-sample linear interpolation on the CDF
        y_clipped = np.clip(y_residual, grids[:, 0], grids[:, -1])
        frac_idx = (y_clipped - grids[:, 0]) / bin_widths
        idx_lo = np.floor(frac_idx).astype(int)
        idx_hi = np.minimum(idx_lo + 1, self.n_bins - 1)
        alpha = frac_idx - idx_lo

        rows = np.arange(len(y_array))
        cdf_vals = (1 - alpha) * cumsum[rows, idx_lo] + alpha * cumsum[rows, idx_hi]
        return np.clip(cdf_vals, 0.0, 1.0)

    def ppf(self, X, q=0.5):
        """
        Percent point function (inverse CDF / quantile function).

        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 quantile values.
        """
        return self.predict_quantile(X, q=q)

    def nll(self, X, y, output_smoothing=0, eps=1e-10):
        """
        Compute mean negative log-likelihood.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.
        y : array-like of shape (n_samples,)
            True target values.
        output_smoothing : float, default=0
            Standard deviation for Gaussian smoothing of the distribution.
        eps : float, default=1e-10
            Floor value to avoid log(0).

        Returns
        -------
        nll : float
            Mean negative log-likelihood: -mean(log f(y|X)).
        """
        return -np.mean(self.logpdf(X, y, output_smoothing=output_smoothing, eps=eps))
