from typing import Any, Callable, Tuple

import numpy as np
import pandas as pd


class EvalMetric:
    """
    Evaluation metrics for time-series forecasting with time-based weighting.
    
    This class provides both cohort-level (aggregated by days_since_ftd) and 
    sample-level metrics with time decay weighting based on business value periods.
    
    Attributes
    ----------
    base_X : pd.DataFrame
        Base dataframe containing days_since_ftd, y_true, and y_pred
    daily_agg : pd.DataFrame
        Daily aggregated metrics with time decay weights
    weights : np.ndarray or None
        Computed weights for samples
    agg_preds : float or None
        Aggregated predictions
    agg_true : float or None
        Aggregated true values
    """

    def __init__(self) -> None:
        """Initialize the EvalMetric class with empty dataframes."""
        self.base_X = pd.DataFrame(
            columns=["days_since_ftd", "y_true", "y_pred"]
        )
        self.daily_agg = pd.DataFrame(
            columns=["days_since_ftd", "time_decay_weight", "daily_agg_target", "daily_agg_preds"]
        )
        self.weights = None
        self.agg_preds = None
        self.agg_true = None
        pass

    def get_daily_agg(self, agg_func: str = "mean") -> pd.DataFrame:
        """
        Calculate the daily aggregated target and predictions grouped by days_since_ftd.

        Parameters
        ----------
        agg_func : str, default='mean'
            Aggregation function to use. Options: 'sum' or 'mean'

        Returns
        -------
        pd.DataFrame
            DataFrame with columns: days_since_ftd, daily_agg_target, daily_agg_preds
        """
        self.daily_agg = (
            self.base_X.groupby("days_since_ftd")
            .agg(
                daily_agg_target=("y_true", agg_func),
                daily_agg_preds=("y_pred", agg_func),
            )
            .reset_index()
        )

        return self.daily_agg

    def get_time_decay_weight(self) -> None:
        """
        Calculates segmented time-based weights based on business value periods.
        
        Weight scheme:
        - Days 0-7: 0.05 (very low weight for initial period)
        - Days 8-14: 0.85 (high weight for early engagement)
        - Days 15-45: 1.00 (maximum weight for peak value period)
        - Days 46-90: 0.90 (high weight)
        - Days 91-180: 0.60 (medium weight)
        - Days 181-270: 0.40 (lower weight)
        - Days 270+: 0.30 (lowest weight)
        
        Returns
        -------
        pd.Series
            Series of time decay weights for each day
        """
        days = self.daily_agg["days_since_ftd"]
        self.daily_agg["time_decay_weight"] = np.where(days <= 7, 0.05,
                                            np.where(days <= 14, 0.85,
                                            np.where(days <= 45, 1.00,
                                            np.where(days <= 90, 0.90,
                                            np.where(days <= 180, 0.60,
                                            np.where(days <= 270, 0.40, 
                                                    0.30))))))
        
        return self.daily_agg["time_decay_weight"]

    def _set_y_true_and_pred(self, y_true: Any, y_pred: Any) -> None:
        """
        Set the true and predicted values in base_X dataframe.
        
        Handles conversion from pandas Series, numpy arrays, or other array-like objects
        to flattened arrays.

        Parameters
        ----------
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values
        """
        if hasattr(y_true, "values"):
            y_true_flat = y_true.values.flatten()
        else:
            y_true_flat = np.array(y_true).flatten()
            
        if hasattr(y_pred, "values"):
            y_pred_flat = y_pred.values.flatten() if hasattr(y_pred, "values") else y_pred
        else:
            y_pred_flat = np.array(y_pred).flatten()
        
        self.base_X["y_true"] = y_true_flat
        self.base_X["y_pred"] = y_pred_flat

    # ========================================================================
    # COHORT-LEVEL METRICS (agregowane po days_since_ftd)
    # ========================================================================
    
    def cohort_weighted_mape(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the cohort-level weighted mean absolute percentage error.
        
        Aggregates predictions and targets by days_since_ftd, then calculates
        MAPE with time decay weighting. Excludes days where target is zero.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean absolute percentage error at cohort level
        """
        self.base_X = X[["days_since_ftd"]].copy().reset_index(drop=True)
        self._set_y_true_and_pred(y_true, y_pred)

        self.get_daily_agg()
        self.get_time_decay_weight()

        mask = self.daily_agg['daily_agg_target'] != 0

        percentage_errors = np.abs((self.daily_agg['daily_agg_target'][mask] - self.daily_agg['daily_agg_preds'][mask]) / self.daily_agg['daily_agg_target'][mask])
        weights = self.daily_agg['time_decay_weight'][mask]
        wmape = np.average(percentage_errors, weights=weights)

        return float(wmape)

    def cohort_weighted_mae(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the cohort-level weighted mean absolute error.
        
        Aggregates predictions and targets by days_since_ftd, then calculates
        MAE with time decay weighting.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean absolute error at cohort level
        """
        self.base_X = X[["days_since_ftd"]].copy().reset_index(drop=True)
        self._set_y_true_and_pred(y_true, y_pred)

        self.get_daily_agg()
        self.get_time_decay_weight()

        abs_errors = np.abs(self.daily_agg['daily_agg_target'] - self.daily_agg['daily_agg_preds'])
        wmae = np.average(abs_errors, weights=self.daily_agg['time_decay_weight'])

        return float(wmae)

    def cohort_weighted_mse(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the cohort-level weighted mean squared error.
        
        Aggregates predictions and targets by days_since_ftd, then calculates
        MSE with time decay weighting.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean squared error at cohort level
        """
        self.base_X = X[["days_since_ftd"]].copy().reset_index(drop=True)
        self._set_y_true_and_pred(y_true, y_pred)

        self.get_daily_agg()
        self.get_time_decay_weight()

        squared_errors = (self.daily_agg['daily_agg_target'] - self.daily_agg['daily_agg_preds'])**2
        wmse = np.average(squared_errors, weights=self.daily_agg['time_decay_weight'])

        return float(wmse)
    
    # ========================================================================
    # SAMPLE-LEVEL METRICS (na pojedynczych próbkach)
    # ========================================================================
    
    def _get_sample_weights(self, X: Any) -> np.ndarray:
        """
        Calculate time-based weights for individual samples.
        
        Uses the same weight scheme as cohort-level metrics but applies
        to individual samples without aggregation.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column

        Returns
        -------
        np.ndarray
            Array of weights for each sample
        """
        days = X["days_since_ftd"].values
        weights = np.where(
            days <= 7,0.05,
            np.where(days <= 14,0.85,
            np.where(days <= 45,1.00,
            np.where(days <= 90, 0.90,
            np.where(days <= 180, 0.60, 
            np.where(days <= 270, 0.40, 
            0.30))),
            ),
            ),
        )
        return weights
    
    def sample_weighted_mae(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the sample-level weighted mean absolute error.
        
        Calculates MAE on individual samples with time decay weighting,
        without aggregation by days_since_ftd.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean absolute error at sample level
        """
        X_reset = X[["days_since_ftd"]].copy().reset_index(drop=True)
        
        self._set_y_true_and_pred(y_true, y_pred)
        
        weights = self._get_sample_weights(X_reset)
        abs_errors = np.abs(self.base_X["y_true"] - self.base_X["y_pred"])
        wmae = np.average(abs_errors, weights=weights)
        
        return float(wmae)
    
    def sample_weighted_mse(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the sample-level weighted mean squared error.
        
        Calculates MSE on individual samples with time decay weighting,
        without aggregation by days_since_ftd.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean squared error at sample level
        """
        X_reset = X[["days_since_ftd"]].copy().reset_index(drop=True)
        
        self._set_y_true_and_pred(y_true, y_pred)
        
        weights = self._get_sample_weights(X_reset)
        squared_errors = (self.base_X["y_true"] - self.base_X["y_pred"])**2
        wmse = np.average(squared_errors, weights=weights)
        
        return float(wmse)
    
    def sample_weighted_mape(self, X: Any, y_true: Any, y_pred: Any) -> float:
        """
        Calculate the sample-level weighted mean absolute percentage error.
        
        Calculates MAPE on individual samples with time decay weighting,
        without aggregation by days_since_ftd. Excludes samples where y_true is zero.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing 'days_since_ftd' column
        y_true : array-like
            True target values
        y_pred : array-like
            Predicted values

        Returns
        -------
        float
            Weighted mean absolute percentage error at sample level
        """
        X_reset = X[["days_since_ftd"]].copy().reset_index(drop=True)
        
        self._set_y_true_and_pred(y_true, y_pred)
        
        weights = self._get_sample_weights(X_reset)
        
        mask = self.base_X["y_true"] != 0
        percentage_errors = np.abs((self.base_X["y_true"][mask] - self.base_X["y_pred"][mask]) / self.base_X["y_true"][mask])
        weights_filtered = weights[mask]
        
        wmape = np.average(percentage_errors, weights=weights_filtered)
        
        return float(wmape)



    def create_lgb_metric(
        self, X_eval: pd.DataFrame, metric_type: str = "mae", level: str = "cohort"
    ) -> Callable[[np.ndarray, np.ndarray], Tuple[str, float, bool]]:
        """
        Creates a LightGBM-compatible evaluation metric with access to X via closure.
        
        This method creates a custom metric function that can be passed to LightGBM's
        eval_metric parameter. The returned function has access to the evaluation
        features (X_eval) through closure, allowing time-weighted metrics.

        Parameters
        ----------
        X_eval : pd.DataFrame
            DataFrame containing 'days_since_ftd' column for evaluation set
        metric_type : str, default='mae'
            Type of metric: 'mae', 'mse', or 'mape'
        level : str, default='cohort'
            Level of aggregation: 
            - 'cohort': aggregated by days_since_ftd
            - 'sample': individual samples

        Returns
        -------
        Callable
            Custom eval metric for LightGBM with signature:
            (y_true, y_pred) -> (metric_name, metric_value, is_higher_better)
            
        Raises
        ------
        ValueError
            If invalid combination of level and metric_type is provided

        Examples
        --------
        >>> eval_metric = EvalMetric()
        >>> lgb_metric = eval_metric.create_lgb_metric(X_val, metric_type='mae', level='cohort')
        >>> model = lgb.train(params, train_data, valid_sets=[val_data], feval=lgb_metric)
        """
        X_eval_reset = X_eval[["days_since_ftd"]].copy().reset_index(drop=True)
    
        # Mapowanie: (level, metric_type) -> (metoda obliczająca, nazwa dla LightGBM)
        metrics_map = {
            ("cohort", "mae"): (self.cohort_weighted_mae, "cohort_wmae"),
            ("cohort", "mse"): (self.cohort_weighted_mse, "cohort_wmse"),
            ("cohort", "mape"): (self.cohort_weighted_mape, "cohort_wmape"),
            ("sample", "mae"): (self.sample_weighted_mae, "sample_wmae"),
            ("sample", "mse"): (self.sample_weighted_mse, "sample_wmse"),
            ("sample", "mape"): (self.sample_weighted_mape, "sample_wmape"),
        }
        
        # Walidacja i wybór metryki
        key = (level, metric_type)
        if key not in metrics_map:
            valid_combinations = ", ".join([f"{lvl}_{typ}" for lvl, typ in metrics_map.keys()])
            raise ValueError(
                f"Invalid combination: level='{level}', metric_type='{metric_type}'. "
                f"Valid combinations: {valid_combinations}"
            )
        
        metric_func, metric_name = metrics_map[key]
        
        # Wrapper dla LightGBM
        def weighted_metric(y_true: Any, y_pred: Any) -> Tuple[str, float, bool]:
            try:
                # Konwersja do pandas Series jeśli potrzeba
                if isinstance(y_true, np.ndarray):
                    y_true = pd.Series(y_true).reset_index(drop=True)
                if isinstance(y_pred, np.ndarray):
                    y_pred = pd.Series(y_pred).reset_index(drop=True)
                
                # Oblicz metrykę używając wybranej funkcji z klasy
                metric_value = metric_func(X_eval_reset, y_true, y_pred)
                
                # Zwróć (nazwa, wartość, czy_większe_lepsze)
                return metric_name, metric_value, False
                
            except Exception as e:
                print(f"❌ Error in {metric_name}: {e}")
                import traceback
                traceback.print_exc()
                return metric_name, 999999, False
        
        return weighted_metric



class DaysWeightedObjective:
    """
    Custom time-decayed objective function for LightGBM.
    
    This class implements a custom loss function with time-based weighting
    for gradient boosting. The objective applies time decay weights based
    on days_since_ftd to prioritize predictions in specific time periods.
    
    Attributes
    ----------
    epsilon : float
        Small constant for numerical stability in calculations
    X : pd.DataFrame or None
        Feature matrix containing days_since_ftd column
    days_col : str
        Name of the column containing days information
    agg_level : str
        Aggregation level ('sample' only currently supported)
    mode : str
        Loss function mode: 'mae', 'mse', or 'mape'
    time_decay : np.ndarray
        Computed time decay weights
    """
    
    def __init__(
        self, days_col="days_since_ftd", agg_level="sample", mode="mse", epsilon=1e-3
    ):
        """
        Initialize the custom objective function.

        Parameters
        ----------
        days_col : str, default='days_since_ftd'
            Name of the column containing days since first transaction
        agg_level : str, default='sample'
            Aggregation level (currently only 'sample' is supported)
        mode : str, default='mse'
            Loss function mode: 'mae', 'mse', or 'mape'
        epsilon : float, default=1e-3
            Small constant for numerical stability
        """
        self.epsilon = epsilon
        self.X = None
        self.days_col = days_col
        self.agg_level = agg_level.lower()
        self.mode = mode.lower()

    def __call__(self, y_true, y_pred):
        """
        Compute gradients and hessians for LightGBM custom objective.

        Parameters
        ----------
        y_true : np.ndarray
            True target values
        y_pred : np.ndarray
            Predicted values

        Returns
        -------
        tuple of np.ndarray
            (gradients, hessians) for LightGBM optimization

        Raises
        ------
        ValueError
            If agg_level is not 'sample'
        """
        if self.agg_level == "sample":
            return self.sample_level(y_true, y_pred)
        else:
            raise ValueError("agg_level must be 'sample'")

    def set_data(self, X):
        """
        Attach feature matrix for objective function computation.
        
        Must be called before using the objective in LightGBM training
        to provide access to the days_since_ftd column.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing days_since_ftd column
        """
        self.X = X[[self.days_col]].copy()

    def get_time_decay_weight(self, X):
        """
        Calculates uniform time-based weights (currently all set to 1).
        
        This method computes time decay weights based on days_since_ftd.
        The current implementation uses uniform weights (all 1.0) across
        all time periods. Can be modified to apply different weights for
        different business value periods.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix containing days_since_ftd column
        """
        days = X[self.days_col].values
        self.time_decay = np.where(
            days <= 7, 1,
            np.where(days <= 14,1,
            np.where(days <= 45,1,
            np.where(days <= 90, 1, 
            np.where(days <= 180, 1, 
            np.where(days <= 270, 1, 
            1))),
            ),
            ),
        )

    # def get_time_decay_weight(self, X):
        # """
        # Calculates the time decay weight
        # """
        # self.time_decay = np.exp(-X[self.days_col].values.astype(np.float64) / 365.0)


    def sample_level(self, y_true, y_pred):
        """
        Compute weighted gradients and hessians for sample-level optimization.
        
        Calculates first and second derivatives of the loss function with
        time decay weighting applied. Supports MAE, MSE, and MAPE modes.

        Parameters
        ----------
        y_true : np.ndarray
            True target values
        y_pred : np.ndarray
            Predicted values

        Returns
        -------
        tuple of np.ndarray
            (grad, hess) - Gradients and Hessians for LightGBM

        Raises
        ------
        ValueError
            If X has not been set via set_data(), or if mode is invalid
        """
        if self.X is None:
            raise ValueError(
                "Feature matrix X has not been set. Use set_data(X) before fitting."
            )

        self.get_time_decay_weight(self.X)
        error = y_pred - y_true
        sign = np.sign(error)

        if self.mode == "mape":
            denom = np.abs(y_true) + self.epsilon
            smooth = np.sqrt(error**2 + self.epsilon**2)

            grad = self.time_decay * error / (smooth * denom)
            hess = self.time_decay * (self.epsilon**2) / (smooth**3 * denom)

        elif self.mode == "mae":
            grad = self.time_decay * sign
            hess = self.time_decay

        elif self.mode == "mse":
            grad = 2 * self.time_decay * error
            hess = 2 * self.time_decay * np.ones_like(error)

        else:
            raise ValueError("mode must be 'mae' or 'mape' or 'mse'")

        return grad, hess
