"""
Two-Stage Distribution Regressor: LGBM Baseline + Soft-Target Residual Distribution

Stage 1: L2 regression baseline (5-fold CV for honest residuals, refit on full data)
Stage 2: Cross-entropy soft-target model on residuals (X × residual_grid → P(residual))

Benefits over single-stage approach:
- Better expected value prediction (L2 baseline)
- Better grid resolution (residuals have narrower range than raw y)
- No need for heteroscedastic sigma estimation
- Similar runtime to the CV auto-sigma approach
"""

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 lightgbm import LGBMRegressor
from scipy.ndimage import gaussian_filter1d
from sklearn.model_selection import cross_val_predict


class DistributionRegressorSoftTargetLGBMBaseline(BaseEstimator, RegressorMixin):
    """
    Two-stage distribution regressor:
    1. LGBM L2 regression baseline for point prediction
    2. Soft-target cross-entropy model for residual distribution

    At prediction time, the residual distribution is shifted by the baseline
    prediction to produce a full distribution in the original y-space.

    Parameters
    ----------
    n_bins : int, default=100
        Number of grid points for residual distribution.

    sigma : float or str, default='auto'
        Gaussian kernel width for soft targets on residuals.
        - 'auto': 1.5 × grid_step (covers ~3 bins)
        - float: explicit width

    output_smoothing : float, default=1.0
        Gaussian smoothing of output distribution. 0 = disabled.

    n_estimators : int, default=100
        Number of boosting trees for both baseline and distribution models.

    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 (both models).
    """

    def __init__(
        self,
        n_bins=100,
        sigma='auto',
        output_smoothing=1.0,
        n_estimators=100,
        learning_rate=0.1,
        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.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 fit(self, X, y, sample_weight=None):
        """
        Fit the two-stage model:
        1. Train L2 baseline via 5-fold CV to get honest residuals, then refit on full data.
        2. Build a residual grid and train a soft-target cross-entropy model on 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_array = X.values
            y_array = np.asarray(y)
        else:
            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}")

        # 2. Stage 1: Baseline L2 regression with 5-fold CV
        lgbm_params = {
            'n_estimators': self.n_estimators,
            'learning_rate': self.learning_rate,
            'random_state': self.random_state,
            'verbose': -1
        }
        lgbm_params.update(self.lgbm_kwargs)

        self.baseline_model_ = LGBMRegressor(**lgbm_params)

        # Use DataFrame consistently for baseline to avoid feature-name warnings
        X_baseline = pd.DataFrame(X_array, columns=self.feature_names_in_) if self._is_dataframe else X_array

        # 5-fold CV to get honest out-of-sample residuals
        cv_params = {'sample_weight': sample_weight} if sample_weight is not None else None
        y_pred_cv = cross_val_predict(self.baseline_model_, X_baseline, y_array,
                                      cv=5, params=cv_params)

        # Refit on full data for prediction
        self.baseline_model_.fit(X_baseline, y_array, sample_weight=sample_weight)

        # 3. Compute residuals from CV predictions
        residuals = y_array - y_pred_cv

        # 4. Build residual grid (centered around 0 with margin)
        r_min = float(np.min(residuals))
        r_max = float(np.max(residuals))
        margin = (r_max - r_min) * 0.1
        self.residual_grid_ = np.linspace(r_min - margin, r_max + margin, self.n_bins)

        # 5. Sigma for soft targets — CV-based estimation on residuals
        grid_step = (self.residual_grid_[-1] - self.residual_grid_[0]) / (self.n_bins - 1)
        if self.sigma == 'auto':
            # 3-fold CV on residuals to get honest "residuals of residuals"
            sigma_base_params = {**lgbm_params, 'n_estimators': min(self.n_estimators, 100)}
            sigma_model = LGBMRegressor(**sigma_base_params)
            resid_pred_cv = cross_val_predict(sigma_model, X_baseline, residuals,
                                              cv=3, params=cv_params)
            resid_of_resid = residuals - resid_pred_cv

            # Heteroscedastic sigma: predict per-sample noise level
            log_resid_sq = np.log(resid_of_resid ** 2 + 1e-9)
            sigma_het_model = LGBMRegressor(**sigma_base_params)
            sigma_het_model.fit(X_baseline, log_resid_sq, sample_weight=sample_weight)
            log_sigma_sq_pred = sigma_het_model.predict(X_baseline)
            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)

        # 6. Monte Carlo Training Expansion (K=5 instead of n_bins copies)
        K = 5
        n_total_rows = n_samples * K
        X_final = np.empty((n_total_rows, self.n_features_in_ + 1), dtype=X_array.dtype)
        X_final[:, :-1].reshape(n_samples, K, -1)[:] = X_array[:, None, :]

        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:5] = rng.uniform(r_min - margin, r_max + margin, size=(n_samples, 2))  # Uniform tails (full grid range)

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

        # 7. Soft targets for the K sampled points
        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

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

        # 9. Train distribution model on residuals
        dist_params = {
            'objective': 'cross_entropy',
            'metric': 'cross_entropy',
            'n_estimators': self.n_estimators,
            'learning_rate': self.learning_rate,
            'random_state': self.random_state,
            'verbose': -1
        }
        dist_params.update(self.lgbm_kwargs)

        self.model_ = LGBMRegressor(**dist_params)

        sample_weight_expanded = np.repeat(sample_weight, K) if sample_weight is not None else None
        self.model_.fit(X_final, y_targets_soft, sample_weight=sample_weight_expanded)

        return self

    def predict_distribution(self, X):
        """
        Returns per-sample grid (in original y-space) and probability distributions.

        The residual distribution is shifted by the baseline prediction so the
        returned grid is in the original target space.

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

        Returns
        -------
        grid : array of shape (n_samples, n_bins)
            Grid points in original y-space (residual_grid + baseline_prediction).

        distributions : array of shape (n_samples, n_bins)
            Probability distribution for each sample at each grid point.
        """
        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]

        # Stage 1: Baseline prediction (use DataFrame if fitted with one to avoid warnings)
        X_baseline = pd.DataFrame(X_array, columns=self.feature_names_in_) if self._is_dataframe else X_array
        y_base = self.baseline_model_.predict(X_baseline)

        # Stage 2: Residual distribution (NumPy-only, no DataFrame)
        n_total_rows = n_samples * self.n_bins
        X_expanded = np.empty((n_total_rows, self.n_features_in_ + 1), dtype=X_array.dtype)
        X_expanded[:, :-1].reshape(n_samples, self.n_bins, -1)[:] = X_array[:, None, :]
        X_expanded[:, -1].reshape(n_samples, self.n_bins)[:] = self.residual_grid_

        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 self.output_smoothing > 0:
            gaussian_filter1d(distributions, sigma=self.output_smoothing, axis=1, output=distributions)
            row_sums = distributions.sum(axis=1, keepdims=True)
            row_sums[row_sums == 0] = 1.0
            distributions /= row_sums

        # Shift residual grid to absolute y-space per sample
        grid = self.residual_grid_[None, :] + y_base[:, None]  # (n_samples, n_bins)

        return grid, distributions

    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.

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

        Returns
        -------
        means : array of shape (n_samples,)
            Predicted means.
        """
        grid, dists = self.predict_distribution(X)
        return np.einsum('ij,ij->i', dists, grid)

    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.
        """
        grid, dists = self.predict_distribution(X)
        max_indices = np.argmax(dists, axis=1)
        return grid[np.arange(len(grid)), max_indices]

    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.
        """
        grid, dists = self.predict_distribution(X)
        cdfs = np.cumsum(dists, axis=1)
        cdfs[:, -1] = 1.0

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

        # Vectorized: broadcast (n_samples, 1, n_bins) >= (1, n_quantiles, 1)
        mask = cdfs[:, None, :] >= q_arr[None, :, None]
        indices = mask.argmax(axis=2)
        quantiles = grid[np.arange(grid.shape[0])[:, None], indices]

        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
        return self.predict_quantile(X, q=[alpha / 2.0, 1.0 - alpha / 2.0])

    def predict_std(self, X):
        """
        Predict standard deviation of the distribution for each sample.

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

        Returns
        -------
        stds : array of shape (n_samples,)
            Predicted standard deviations.
        """
        grid, dists = self.predict_distribution(X)
        means = np.einsum('ij,ij->i', dists, grid)
        expected_sq = np.einsum('ij,ij->i', dists, grid ** 2)
        variance = expected_sq - means ** 2
        return np.sqrt(np.maximum(variance, 0.0))
