import pandas as pd
import numpy as np
import logging
from typing import Dict, List, Tuple, Optional, Union, Any
from scipy.stats import nbinom
from scipy.stats import gamma

from ._gbm_trainer import GBMTrainer 

__all__ = ['GBMTrainerCount']

# Set up a logger for this module
logger = logging.getLogger(__name__)

class GBMTrainerCount:
    """
    A class to train a Gradient Boosting Machine (GBM) for count data.

    This class uses a multi-class classification GBM to model the probability
    distribution of counts (treated as discrete "delays" or classes).
    It then transforms these probabilities to estimate counts and provides
    functionality for conformal prediction intervals.

    Attributes:
        objective (str): The objective function for the GBM ('multiclass').
        eval_metric (str): The evaluation metric for the GBM ('multi_logloss').
        n_estimators (int): Maximum number of boosting rounds.
        stopping_rounds (int): Rounds to wait for metric improvement before stopping.
        train_size (float): Proportion of data to use for training (vs. validation).
        log_period (int | None): Period to print logging messages during training.
        random_state (int): Seed for reproducibility.
        model_p (Optional[GBMTrainer]): The underlying GBMTrainer instance.
        delays (Optional[np.ndarray]): Sorted array of unique class labels (counts).
        conformal_quantiles (Optional[List[float]]): Calibrated quantiles for each delay.
        alpha (Optional[float]): The significance level for conformal prediction.

    """

    def __init__(
        self,
        n_estimators: int = 10000,
        stopping_rounds: int = 10,
        train_size: float = 0.8,
        log_period: Optional[int] = None,
        random_state: int = 0
    ):
        """
        Initializes the GBMTrainerCount parameters.
        """
        self.objective = 'multiclass'
        self.eval_metric = 'multi_logloss'
        self.n_estimators = n_estimators
        self.stopping_rounds = stopping_rounds
        self.train_size = train_size
        self.log_period = log_period
        self.random_state = random_state

        # Attributes to be populated during fitting
        self.model_p: Optional[GBMTrainer] = None
        self.delays: Optional[np.ndarray] = None
        self.conformal_quantiles: Dict[Tuple, Dict[str, Any]] = {}

    def fit(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        verbose: bool = False,
        refit: bool = True
    ) -> None:
        """
        Fits the underlying multi-class GBM model.

        Args:
            X: Feature DataFrame.
            y: Target variable (counts/classes).
            verbose: Whether to print training logs.
            refit: Whether to refit the model if it already exists.
        """
        logger.info("Fitting probability model (GBM multiclass)...")
        self.model_p = GBMTrainer(
            objective=self.objective,
            eval_metric=self.eval_metric,
            n_estimators=self.n_estimators,
            stopping_rounds=self.stopping_rounds,
            train_size=self.train_size,
            log_period=self.log_period,
            random_state=self.random_state
        )
        self.model_p.fit(X, y, verbose=verbose, refit=refit)
        self.delays = np.sort(self.model_p.model.classes_)
        logger.info("Probability model fitting complete.")

    def predict_parameters(
        self,
        X: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Predicts probabilities and derived parameters.

        This calculates:
        1.  Class probabilities: P(Y=k)
        2.  Cumulative probabilities: P(Y<=k)
        3.  Variance estimate: Based on the assumption that the count follows
            a geometric-like distribution where the "success" probability
            is p = P(Y<=k). The variance of such a count (1/p) is
            approximated by (1-p) / p^2.

        Args:
            X: Feature DataFrame.

        Returns:
            A tuple containing:
            - probabilities (np.ndarray): P(Y=k) for each class.
            - cumulative_probabilities (np.ndarray): P(Y<=k) for each class.
            - variance_estimate (np.ndarray): Estimated variance of the count.
        """
        if self.model_p is None or self.model_p.model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        features = self.model_p.model.feature_name_
        model_classes = self.model_p.model.classes_
        sorted_indices = np.argsort(model_classes)

        # Predict probabilities and sort them by class label
        probabilities = self.model_p.model.predict_proba(X[features])[:, sorted_indices]
        
        # Calculate cumulative probability P(Y <= k)
        cumulative_probabilities = np.cumsum(probabilities, axis=1)[:, :-1]
        
        # Estimate variance: Var(1/p) approx (1-p)/p^2
        variance_estimate = (1 - cumulative_probabilities) / (cumulative_probabilities ** 2)

        return probabilities, cumulative_probabilities, variance_estimate

    def predict_count(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        index_agg: pd.DataFrame
    ) -> List[Tuple[np.ndarray, np.ndarray, np.ndarray, pd.MultiIndex]]:
        """
        Predicts aggregated counts and variances for each delay threshold.

        Args:
            X: Feature DataFrame.
            y: Target variable (used for filtering by delay).
            index_agg: DataFrame with columns to group by (must include
                       'DT_SIN_PRI_SEM').

        Returns:
            A list of tuples. Each tuple corresponds to a delay threshold
            and contains:
            - aggregated_count (np.ndarray): Aggregated simple row counts (sum of ones).
            - aggregated_pred (np.ndarray): Aggregated predicted counts (sum of 1/p).
            - aggregated_variance (np.ndarray): Aggregated variance.
            - aggregated_index (pd.MultiIndex): The unique indices for the aggregates.
        """
        if self.model_p is None or self.model_p.model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        if 'DT_SIN_PRI_SEM' not in index_agg.columns:
            raise ValueError("'DT_SIN_PRI_SEM' must be a column in index_agg.")
            
        features = self.model_p.model.feature_name_
        _, cumulative_probabilities, variance_estimate = self.predict_parameters(X[features])

        aggregated_counts_by_delay = []
        for i, delay in enumerate(self.delays[:-1]):
            # Filter data up to the current delay
            mask_delay = (y <= int(delay))

            cum_prob_delay = cumulative_probabilities[mask_delay, i]
            var_delay = variance_estimate[mask_delay, i]
            
            # Point estimate for count: E[X] = 1/p
            point_count_pred = 1.0 / cum_prob_delay
            index_delay_df = index_agg.loc[mask_delay]

            # Aggregate counts and variances
            data_to_agg = {
                'count': np.ones(len(point_count_pred)), 
                'pred': point_count_pred,
                'var': var_delay
            }

            agg_df = self._aggregate_by_columns(data_to_agg, index_delay_df)

            aggregated_count = agg_df['count'].to_numpy()
            aggregated_pred = agg_df['pred'].to_numpy()
            aggregated_variance = agg_df['var'].to_numpy()
            aggregated_index = agg_df.index # This is now a MultiIndex

            aggregated_counts_by_delay.append(
                (aggregated_count, aggregated_pred, aggregated_variance, aggregated_index)
            )

        return aggregated_counts_by_delay

    def _aggregate_by_columns(
        self,
        data: Dict[str, np.ndarray],
        index_df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Aggregates data by the columns of a given DataFrame.

        Args:
            data: A dictionary where keys are column names (e.g., 'count', 'var')
                  and values are numpy arrays of the data.
            index_df: A DataFrame whose columns will be used as
                      aggregation keys.

        Returns:
            A DataFrame with data aggregated (summed) by the columns
            of index_df.
        """
        if not all(len(arr) == len(index_df) for arr in data.values()):
            raise ValueError("Data arrays and index DataFrame must have the same length.")
            
        full_df = index_df.copy()
        for col_name, col_data in data.items():
            full_df[col_name] = col_data
            
        agg_cols = list(index_df.columns)
        agg_df = full_df.groupby(agg_cols).sum()

        # if ('var' in agg_df.columns) and ('pred' in agg_df.columns):
        #     pred_mean = full_df[agg_cols+['pred']].groupby(agg_cols).mean().reset_index().rename(columns={'pred':'pred_mean'})
        #     agg_df_aux = (full_df[agg_cols+['pred']]
        #      .merge(pred_mean, on=agg_cols)
        #      .assign(squared_residuals=lambda x: (x['pred'] - x['pred_mean'])**2) 
        #         .groupby(agg_cols)
        #         .sum()
        #     )
        #     agg_df = agg_df.merge(agg_df_aux['squared_residuals'], left_index=True, right_index=True, how='left')
        #     agg_df['var'] = np.where(agg_df['count'] > 1, agg_df['var'] + (agg_df['count']/(agg_df['count']-1) )* agg_df['squared_residuals'], agg_df['var'])
        #     agg_df = agg_df.drop(columns=['squared_residuals'])

        return agg_df

    def count_aggregation(
        self,
        index_agg: pd.DataFrame
    ) -> Tuple[np.ndarray, pd.MultiIndex]:
        """
        Aggregates a simple count (1 for each record) by the given index.

        Args:
            index_agg: DataFrame with columns to group by (must include
                       'DT_SIN_PRI_SEM').

        Returns:
            A tuple containing:
            - count (np.ndarray): The sum of counts for each unique index.
            - index_count (pd.MultiIndex): The unique indices.
        """
        data_to_agg = {'count': np.ones(len(index_agg))}
        agg_df = self._aggregate_by_columns(data_to_agg, index_agg)
        
        count = agg_df['count'].to_numpy()
        index_count = agg_df.index
        
        return count, index_count

    def fit_conformal(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        index_agg: pd.DataFrame,
        index_agg_result: pd.DataFrame,
        alpha: Tuple[float, float] = (0.05, 0.05),
        method: str = "asym_residuals"
    ) -> None:
        """
        Fits the conformal prediction quantiles using a calibration dataset.

        Args:
            X: Calibration feature DataFrame.
            y: Calibration target variable.
            index_agg: Calibration index for pre-aggregated data (DataFrame).
            index_agg_result: Calibration index for true aggregated results (DataFrame).
            alpha: Significance level (e.g., 0.05 for 95% intervals).
            method: Method type for conformal scores ("asym_residuals", "abs_residuals", "nbinom", "gamma").
        """
        if self.delays is None:
            raise RuntimeError("Model not fitted. Call fit() first.")

        # Validate aggregation columns
        agg_cols = tuple(sorted(index_agg.columns))
        agg_cols_result = tuple(sorted(index_agg_result.columns))

        if agg_cols != agg_cols_result:
            raise ValueError(
                "Columns of index_agg and index_agg_result must match."
            )
        
        if 'DT_SIN_PRI_SEM' not in agg_cols:
            raise ValueError(
                "'DT_SIN_PRI_SEM' must be in aggregation columns."
            )
            
        logger.info(
            f"Fitting conformal quantiles with alpha={alpha} "
            f"for agg_cols={agg_cols}"
        )
            
        # Get date range from the "truth" dataframe
        dt_min = index_agg_result['DT_SIN_PRI_SEM'].min()
        dt_max = index_agg_result['DT_SIN_PRI_SEM'].max()
        scores_by_delay = self._calculate_nonconformity_scores(
            X, y, index_agg, index_agg_result, method=method, alpha=alpha
        )

        # Calculate quantiles and store them in the new structure
        calculated_quantiles = []
        for scores in scores_by_delay:
            if len(scores) == 0:
                logger.warning(
                    "No matching data found for a delay threshold during "
                    "conformal fitting. Appending quantile=0. This may "
                    "indicate a data mismatch."
                )
                quantiles_to_append = [0.0, 0.0] if method == "asym_residuals" else [0.0]
            else:

                if method in ["nbinom", "gamma"]:
                    q_value = np.quantile(scores, 1 - alpha[0])
                    quantiles_to_append = [q_value]
                elif method == "asym_residuals":
                    q_lower = np.quantile(np.maximum( -scores, 0), 1 - alpha[0]/2) 
                    q_upper = np.quantile(np.maximum(  scores, 0), 1 - alpha[0]/2) 
                    quantiles_to_append = [q_lower, q_upper]
                elif method == "abs_residuals":
                    q_value = np.quantile(scores, 1 - alpha[0])
                    quantiles_to_append = [q_value]

            calculated_quantiles.append(quantiles_to_append)

        storage_key = (agg_cols, alpha, method)
        payload = {
            'quantiles': calculated_quantiles,
            'dt_min': dt_min,
            'dt_max': dt_max,
            'scores_by_delay': scores_by_delay
        }

        if storage_key in self.conformal_quantiles:
            logger.warning(
                f"Overwriting existing conformal quantiles for key: {storage_key}"
            )
            
        self.conformal_quantiles[storage_key] = payload
        logger.info(
            f"Conformal quantiles fitted and stored for key: {storage_key}. "
            f"Calibration date range: {dt_min} to {dt_max}."
        )


    def _calculate_nonconformity_scores(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        index_agg: pd.DataFrame,
        index_agg_result: pd.DataFrame,
        method: str,
        alpha: Tuple[float, float]
    ) -> List[np.ndarray]:
        """
        Calculates scaled nonconformity scores for conformal fitting.

        Args:
            X: Calibration feature DataFrame.
            y: Calibration target variable.
            index_agg: Calibration index for pre-aggregated data (DataFrame).
            index_agg_result: Calibration index for true aggregated results (DataFrame).
            method: Method type for conformal scores ("asym_residuals", "abs_residuals", "nbinom", "gamma").
            alpha: Significance level (e.g., 0.05 for 95% intervals).
        Returns:
            A list of numpy arrays, where each array contains the
            nonconformity scores for a specific delay threshold.
        """

        # Get predicted aggregated counts
        predicted_counts_by_delay = self.predict_count(X, y, index_agg)
        
        # Get true aggregated counts
        true_counts, true_counts_index = self.count_aggregation(index_agg_result)
        true_df = pd.DataFrame(data={'count': true_counts}, index=true_counts_index)

        scores_by_delay = []
        for i in range(len(self.delays[:-1])):
            count_agg, pred_agg_count, pred_agg_var, pred_agg_index = predicted_counts_by_delay[i]

            pred_df = pd.DataFrame(
                data={'pred_count': pred_agg_count, 'pred_var': pred_agg_var,'count_aux': count_agg},
                index=pred_agg_index
            )

            # Align predictions with true values
            df_merged = pred_df.merge(true_df, left_index=True, right_index=True, how='inner')

            # Calculate scaled residuals (nonconformity scores)
            residuals = df_merged['count'] - df_merged['pred_count']
            pred_std = np.sqrt(df_merged['pred_var'])
            if method == "nbinom":
                mu = df_merged['pred_count'] - df_merged['count_aux']
                sigma2 = df_merged['pred_var']
                n = mu**2 / (sigma2 - mu)
                p = mu/sigma2
                quantile_lower = nbinom.ppf(alpha[1]/2, n, p ) + df_merged['count_aux']
                quantile_upper = nbinom.ppf(1 - alpha[1]/2, n , p ) + df_merged['count_aux']
                nonconformity_scores = np.maximum( quantile_lower - df_merged['count']
                                                    , df_merged['count'] - quantile_upper)      
            elif method == "gamma":
                mu = df_merged['pred_count'] - df_merged['count_aux']
                sigma2 = df_merged['pred_var']
                a = mu**2 / sigma2
                scale = sigma2 / mu
                quantile_lower = gamma.ppf(alpha[1]/2, a=a,loc=0, scale=scale) + df_merged['count_aux']
                quantile_upper = gamma.ppf(1 - alpha[1]/2, a=a, loc=0, scale=scale) + df_merged['count_aux']
                nonconformity_scores = np.maximum( quantile_lower - df_merged['count']
                                                    , df_merged['count'] - quantile_upper)    

            elif method == "asym_residuals":
                nonconformity_scores = residuals / pred_std

            elif method == "abs_residuals":
                nonconformity_scores = np.abs(residuals) / pred_std

            scores_by_delay.append(nonconformity_scores)

        return scores_by_delay

    def predict_count_conformal(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        index_agg: pd.DataFrame,
        alpha: Tuple[float, float] = (0.05, 0.05),
        method: str = "asym_residuals",
        constrained: bool = False
    ) -> List[Tuple[np.ndarray, np.ndarray, np.ndarray, pd.MultiIndex, np.ndarray, np.ndarray]]:
        """
        Predicts aggregated counts with conformal prediction intervals.

        Args:
            X: Feature DataFrame for prediction.
            y: Target variable for prediction (used for filtering).
            index_agg: DataFrame with columns to group by (must include
                       'DT_SIN_PRI_SEM').
            alpha: Significance level (e.g., 0.05) to *retrieve* quantiles for.
            constrained: If True, the lower and upper bounds
                are constrained to be at least the simple aggregated count.            

        Returns:
            A list of tuples. Each tuple corresponds to a delay threshold
            and contains:
            - aggregated_count (np.ndarray): Aggregated simple row counts.
            - aggregated_pred (np.ndarray): Aggregated predicted counts.
            - aggregated_variance (np.ndarray): Aggregated variance.
            - aggregated_index (pd.MultiIndex): The unique indices.
            - lower_bound (np.ndarray): Lower bound of the conformal interval.
            - upper_bound (np.ndarray): Upper bound of the conformal interval.
        """

        if not self.conformal_quantiles:
            raise RuntimeError("No conformal models have been fitted. Call fit_conformal() first.")
        if self.delays is None:
            raise RuntimeError("Model not fitted. Call fit() first.")

        # Find the correct quantiles from storage
        if 'DT_SIN_PRI_SEM' not in index_agg.columns:
            raise ValueError("'DT_SIN_PRI_SEM' must be a column in index_agg.")
            
        agg_cols = tuple(sorted(index_agg.columns))
        storage_key = (agg_cols, alpha, method)
        
        if storage_key not in self.conformal_quantiles:
            raise ValueError(
                f"No conformal quantiles found for agg_cols={agg_cols}, "
                f"alpha={alpha}, and method={method}. "
                "Run fit_conformal() first with these exact parameters."
            )
            
        conformal_data = self.conformal_quantiles[storage_key]
        quantiles_list = conformal_data['quantiles']
        dt_min_cal = conformal_data['dt_min']
        dt_max_cal = conformal_data['dt_max']
        
        logger.info(
            f"Applying quantiles for key {storage_key}, "
            f"calibrated on {dt_min_cal} to {dt_max_cal}."
        )
        predicted_counts_by_delay = self.predict_count(X, y, index_agg)

        conformal_predictions = []
        for i in range(len(self.delays[:-1])):
            agg_count, pred_agg_count, pred_agg_var, pred_agg_index = predicted_counts_by_delay[i]

            quantile_values = quantiles_list[i]
            pred_std = np.sqrt(pred_agg_var)

            if method == "nbinom":
                q_value = quantile_values[0]

                mu = pred_agg_count - agg_count
                sigma2 = pred_agg_var
                n = mu**2 / (sigma2 - mu)
                p = mu/sigma2

                lower_bound_calc = nbinom.ppf(alpha[1]/2, n, p) + agg_count - q_value
                upper_bound_calc = nbinom.ppf(1 - alpha[1]/2, n, p) + agg_count + q_value
                pred_agg_count = nbinom.ppf(.5, n, p) + agg_count 

            elif method == "gamma":
                q_value = quantile_values[0]

                mu = pred_agg_count - agg_count
                sigma2 = pred_agg_var

                a = mu**2 / sigma2
                scale = sigma2 / mu
                lower_bound_calc = gamma.ppf(alpha[1]/2,  a=a,loc=0, scale=scale) + agg_count - q_value
                upper_bound_calc = gamma.ppf(1 - alpha[1]/2,  a=a,loc=0, scale=scale) + agg_count + q_value
                pred_agg_count = gamma.ppf(.5,  a=a,loc=0, scale=scale) + agg_count
            
            elif method == 'asym_residuals':
                q_lower = quantile_values[0]
                q_upper = quantile_values[1]
                margin_lower = q_lower * pred_std  
                margin_upper = q_upper * pred_std
                lower_bound_calc = pred_agg_count - margin_lower 
                upper_bound_calc = pred_agg_count + margin_upper

            elif method == 'abs_residuals':
                q_value = quantile_values[0]
                margin = q_value * pred_std
                lower_bound_calc = pred_agg_count - margin
                upper_bound_calc = pred_agg_count + margin

            if constrained:
                lower_bound = np.maximum( np.ceil(lower_bound_calc), agg_count)
                upper_bound = np.maximum(np.floor(upper_bound_calc), agg_count)
                pred_agg_count = np.maximum(np.round(pred_agg_count), agg_count)
            else:
                lower_bound = lower_bound_calc
                upper_bound = upper_bound_calc
            
            conformal_predictions.append(
                (agg_count, pred_agg_count, pred_agg_var, pred_agg_index, lower_bound, upper_bound)
            )

        return conformal_predictions