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 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,
        index_agg_unique = None # XXXX
    ) -> 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)

            # XXXX
            if index_agg_unique is not None: 
                list_agg_columns = list(index_agg.columns)
                agg_df2 = (index_agg_unique[list_agg_columns]
                          .merge(agg_df.reset_index(), how='left', on= list_agg_columns )
                          .fillna(0)
                          .set_index(list_agg_columns) ) 
                agg_df = agg_df2
            # XXXX          

            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()

        return agg_df

    def count_aggregation(
        self,
        index_agg: pd.DataFrame,
        index_agg_unique = None # XXXX
    ) -> 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)
        
        # XXXX
        if index_agg_unique is not None: 
            list_agg_columns = list(index_agg.columns)
            agg_df2 = (index_agg_unique[list_agg_columns]
                        .merge(agg_df.reset_index(), how='left', on= list_agg_columns )
                        .fillna(0)
                        .set_index(list_agg_columns) ) 
            agg_df = agg_df2
        # XXXX     

        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,
        index_agg_unique = None, # XXXX
        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, mask_valid_by_delay = self._calculate_nonconformity_scores( # XXXX
            X, y, index_agg, index_agg_result, method=method, alpha=alpha,
        index_agg_unique = index_agg_unique # XXXX
        )

        # Calculate quantiles and store them in the new structure
        calculated_quantiles = []
        calculated_quantiles_zero = [] # XXXX

        for scores, mask_valid in zip(scores_by_delay, mask_valid_by_delay): # XXXX
            quantiles_to_append = [0.0, 0.0] if method == "asym_residuals" else [0.0]
            quantiles_zero_to_append = [0.0]
            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."
                )
            else:
                # XXXX
                mask_zero = ~mask_valid
                if mask_zero.any():
                    q_zero = np.quantile(scores.loc[mask_zero], 1 - alpha[0])
                    quantiles_zero_to_append = [q_zero]

                if mask_valid.any():
                    v_scores = scores.loc[mask_valid]

                    if method in ["nbinom", "gamma", "abs_residuals"]:
                        q_value = np.quantile(v_scores, 1 - alpha[0])
                        quantiles_to_append = [q_value]

                    elif method == "asym_residuals":
                        q_lower = np.quantile(np.maximum(-v_scores, 0), 1 - alpha[0]/2) 
                        q_upper = np.quantile(np.maximum( v_scores, 0), 1 - alpha[0]/2) 
                        quantiles_to_append = [q_lower, q_upper]
                # XXXX

            calculated_quantiles.append(quantiles_to_append)
            calculated_quantiles_zero.append(quantiles_zero_to_append) # XXXX

        storage_key = (agg_cols, alpha, method)
        payload = {
            'quantiles': calculated_quantiles,
            'quantiles_zero': calculated_quantiles_zero,  # XXXX
            'dt_min': dt_min,
            'dt_max': dt_max,
            'scores_by_delay': scores_by_delay, 
            'mask_valid_by_delay': mask_valid_by_delay  # XXXX
        }

        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],
        index_agg_unique = None # XXXX
    ) -> Tuple[List[pd.Series], List[pd.Series]]:
        """
        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, index_agg_unique) # XXXX
        
        # Get true aggregated counts
        true_counts, true_counts_index = self.count_aggregation(index_agg_result, index_agg_unique) # XXXX
        true_df = pd.DataFrame(data={'count': true_counts}, index=true_counts_index)

        scores_by_delay = []
        mask_valid_by_delay = [] # XXXX
        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')

            # XXXX
            nonconformity_scores = pd.Series(index=df_merged.index, dtype=float)

            mask_zero = df_merged['pred_count'] == 0
            mask_valid = ~mask_zero

            nonconformity_scores.loc[mask_zero] = df_merged.loc[mask_zero, 'count']

            if mask_valid.any():
                v_count = df_merged.loc[mask_valid, 'count']
                v_pred_count = df_merged.loc[mask_valid, 'pred_count']
                v_pred_var = df_merged.loc[mask_valid, 'pred_var']

                if method in ["nbinom","gamma"]:
                    v_count_aux = df_merged.loc[mask_valid, 'count_aux']
                    mu = v_pred_count - v_count_aux
                    sigma2 = v_pred_var

                    if method == "nbinom":
                        n = mu**2 / (sigma2 - mu) 
                        p = mu / sigma2
                        quantile_lower = nbinom.ppf(alpha[1]/2, n, p) + v_count_aux
                        quantile_upper = nbinom.ppf(1 - alpha[1]/2, n, p) + v_count_aux

                    elif method == "gamma":
                        a = mu**2 / sigma2
                        scale = sigma2 / mu
                        quantile_lower = gamma.ppf(alpha[1]/2, a=a, loc=0, scale=scale) + v_count_aux
                        quantile_upper = gamma.ppf(1 - alpha[1]/2, a=a, loc=0, scale=scale) + v_count_aux

                    nonconformity_scores.loc[mask_valid] = np.maximum(quantile_lower - v_count, v_count - quantile_upper)

                elif method in ["asym_residuals","abs_residuals"]:
                    v_residuals = v_count - v_pred_count
                    v_pred_std = np.sqrt(v_pred_var)

                    if method == "asym_residuals":
                        nonconformity_scores.loc[mask_valid] = v_residuals / v_pred_std
 
                    elif method == "abs_residuals":
                        nonconformity_scores.loc[mask_valid] = np.abs(v_residuals) / v_pred_std
            
            scores_by_delay.append(nonconformity_scores)
            mask_valid_by_delay.append(mask_valid)

            # XXXX

        return scores_by_delay, mask_valid_by_delay

    def predict_count_conformal(
        self,
        X: pd.DataFrame,
        y: Union[pd.Series, np.ndarray],
        index_agg: pd.DataFrame,
        index_agg_unique = None, # XXXX
        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, 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']
        quantiles_zero_list = conformal_data['quantiles_zero'] # XXXX
        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, index_agg_unique)

        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]
            quantiles_zero_values = quantiles_zero_list[i] # XXXX

            lower_bound = np.zeros_like(pred_agg_count, dtype=float)
            upper_bound = np.zeros_like(pred_agg_count, dtype=float)

            mask_zero = pred_agg_count == 0 # XXXX
            mask_valid = ~mask_zero # XXXX

            lower_bound[mask_zero] = 0
            upper_bound[mask_zero] = quantiles_zero_values[0]

            if mask_valid.any():
                v_pred_agg_count = pred_agg_count[mask_valid]  
                v_agg_count = agg_count[mask_valid]
                v_pred_agg_var = pred_agg_var[mask_valid]

                if method in ["nbinom", "gamma"]:
                    q_value = quantile_values[0]
                    mu = v_pred_agg_count - v_agg_count
                    sigma2 = v_pred_agg_var

                    if method == "nbinom":
                        n = mu**2 / (sigma2 - mu)
                        p = mu/sigma2
                        lower_bound_calc = nbinom.ppf(alpha[1]/2, n, p) + v_agg_count - q_value
                        upper_bound_calc = nbinom.ppf(1 - alpha[1]/2, n, p) + v_agg_count + q_value
                        v_pred_agg_count = nbinom.ppf(.5, n, p) + v_agg_count 

                    elif method == "gamma":
                        a = mu**2 / sigma2
                        scale = sigma2 / mu
                        lower_bound_calc = gamma.ppf(alpha[1]/2,  a=a,loc=0, scale=scale) + v_agg_count - q_value
                        upper_bound_calc = gamma.ppf(1 - alpha[1]/2,  a=a,loc=0, scale=scale) + v_agg_count + q_value
                        v_pred_agg_count = gamma.ppf(.5,  a=a,loc=0, scale=scale) + v_agg_count
                
                elif method in ['asym_residuals','abs_residuals']:
                    pred_std = np.sqrt(v_pred_agg_var)

                    if 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 = v_pred_agg_count - margin_lower 
                        upper_bound_calc = v_pred_agg_count + margin_upper

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

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

        return conformal_predictions