"""
NGR Analysis and Visualization Module

This module provides functions for analyzing and visualizing NGR (Net Gaming Revenue)
predictions and actual values, including segmentation and error metrics calculation.
"""

from typing import Optional, Dict, Tuple
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


def print_ngrs_improved(
    test_df: pd.DataFrame,
    additional_text: str = '',
    late_stage_day: Optional[int] = None,
    late_stage_correction: Optional[float] = None
) -> Dict[str, float]:
    """
    Calculate and visualize NGR metrics grouped by days_since_ftd with business-optimal weighting.
    
    This function aggregates predictions and actual values, applies optional late-stage corrections,
    calculates various error metrics (MAE, MAPE, ME, MPE) with business weights, and generates
    a comprehensive visualization.
    
    Parameters
    ----------
    test_df : pd.DataFrame
        DataFrame containing prediction results with columns:
        - 'days_since_ftd': int, days since first time deposit
        - 'remaining_ngr': float, actual remaining NGR
        - 'remaining_ngr_pred': float, predicted remaining NGR
    additional_text : str, default=''
        Additional text to append to plot title
    late_stage_day : int, optional
        Day threshold after which to apply correction
    late_stage_correction : float, optional
        Correction factor to apply (0.0-1.0) linearly after late_stage_day
    
    Returns
    -------
    Dict[str, float]
        Dictionary containing calculated metrics:
        - 'Standard MAE': Mean Absolute Error
        - 'Standard MAPE (%)': Mean Absolute Percentage Error
        - 'Business Optimal MAE': Weighted MAE
        - 'Business Optimal MAPE (%)': Weighted MAPE
    
    Notes
    -----
    Business-optimal weights by day ranges:
    - Days 1-7: 0.05 (learning period - insufficient data)
    - Days 8-14: 0.85 (early signals emerging)
    - Days 15-45: 1.00 (🥇 SUPER CRITICAL - optimal intervention window)
    - Days 46-90: 0.90 (confirmation period - high value)
    - Days 91-180: 0.60 (established patterns - moderate value)
    - Days 181-270: 0.40 (mature behavior - operational value)
    - Days 271+: 0.30 (end-game precision - tactical value)
    
    Examples
    --------
    >>> metrics = print_ngrs_improved(
    ...     test_df,
    ...     additional_text='(Q4 2024)',
    ...     late_stage_day=320,
    ...     late_stage_correction=0.95
    ... )
    >>> print(f"Business MAE: {metrics['Business Optimal MAE']:.2f}")
    """
    # Calculate mean values grouped by days_since_ftd
    mean_values = test_df.groupby('days_since_ftd').agg({
        'remaining_ngr': 'mean',
        'remaining_ngr_pred': 'mean',
    }).reset_index()
    mean_values = mean_values[mean_values['days_since_ftd'] < 365]

    # Apply late stage correction if specified
    if late_stage_day is not None and late_stage_correction is not None:
        days_array = mean_values['days_since_ftd'].values
        days_from_start = days_array - late_stage_day
        max_days = 364 - late_stage_day
        
        correction_factor = np.where(
            days_array <= late_stage_day,
            1.0,  # Before late_stage_day: no changes
            1.0 - (1.0 - late_stage_correction) * (days_from_start / max_days)  # Linear transition
        )
        mean_values['remaining_ngr_pred'] = mean_values['remaining_ngr_pred'] * correction_factor
    
    days = mean_values['days_since_ftd'].values.astype(np.float64)
    
    # Business-optimal weighting
    weights = np.where(days <= 7, 0.05,      # Learning period - insufficient data
             np.where(days <= 14, 0.85,     # Early signals emerging
             np.where(days <= 45, 1.00,     # 🥇 SUPER CRITICAL - optimal intervention window
             np.where(days <= 90, 0.90,     # Confirmation period - high value
             np.where(days <= 180, 0.60,    # Established patterns - moderate value
             np.where(days <= 270, 0.40,    # Mature behavior - operational value
                      0.30))))))            # End-game precision - tactical value

    abs_errors = np.abs(mean_values['remaining_ngr'] - mean_values['remaining_ngr_pred'])
    errors = mean_values['remaining_ngr'] - mean_values['remaining_ngr_pred']

    # Define bins for error analysis by period
    bins = [0, 7, 14, 45, 90, 180, 270, 364]
    labels = ['1-7', '8-14', '15-45', '46-90', '91-180', '181-270', '271-364']

    # Calculate standard metrics
    business_optimal_tmae = (np.sum(abs_errors * weights) / np.sum(weights))
    business_optimal_me = (np.sum(errors * weights) / np.sum(weights))
    standard_mae = round(np.mean(abs_errors), 4)
    standard_me = round(np.mean(errors), 4)
    
    # Calculate percentage-based metrics
    # Avoid division by zero - use small epsilon for very small values
    epsilon = 1e-6
    safe_remaining_ngr = np.where(
        np.abs(mean_values['remaining_ngr']) < epsilon,
        epsilon,
        mean_values['remaining_ngr']
    )

    percentage_abs_errors = np.abs(
        (mean_values['remaining_ngr'] - mean_values['remaining_ngr_pred']) / safe_remaining_ngr
    ) * 100
    errors_percentage = (
        (mean_values['remaining_ngr'] - mean_values['remaining_ngr_pred']) / safe_remaining_ngr
    ) * 100

    # Assign each day to a period bin
    day_bins = pd.cut(
        mean_values['days_since_ftd'],
        bins=bins,
        labels=labels,
        right=True,
        include_lowest=True
    )
    
    # Calculate mean errors by period
    mean_abs_error_by_bin = pd.DataFrame({
        'bin': labels,
        'mean_abs_error': [abs_errors[day_bins == label].mean() for label in labels],
        'mean_error': [errors[day_bins == label].mean() for label in labels],
        'mape': [percentage_abs_errors[day_bins == label].mean() for label in labels]
    })

    # Calculate business-optimal percentage metrics
    business_optimal_mape = (np.sum(percentage_abs_errors * weights) / np.sum(weights)) 
    standard_mape = np.mean(percentage_abs_errors)

    business_optimal_mpe = (np.sum(errors_percentage * weights) / np.sum(weights)) 
    standard_mpe = np.mean(errors_percentage)

    # Create visualization
    plt.figure(figsize=(14, 8))

    # Plot both lines on same axes
    plt.plot(
        mean_values['days_since_ftd'],
        mean_values['remaining_ngr'],
        'b-',
        label='Actual NGR',
        linewidth=2
    )
    plt.plot(
        mean_values['days_since_ftd'],
        mean_values['remaining_ngr_pred'],
        'r-',
        label='Predicted NGR',
        linewidth=2
    )

    # Add business-optimal period reference lines
    plt.axvline(x=7, color='red', linestyle='--', alpha=0.8, linewidth=1)
    plt.axvline(x=14, color='orange', linestyle='--', alpha=0.8, linewidth=1)
    plt.axvline(x=45, color='gold', linestyle='--', alpha=0.9, linewidth=1)
    plt.axvline(x=90, color='green', linestyle='--', alpha=0.8, linewidth=1)
    plt.axvline(x=180, color='blue', linestyle='--', alpha=0.7, linewidth=1)
    plt.axvline(x=270, color='purple', linestyle='--', alpha=0.6, linewidth=1)

    # Add weight labels at the top of each period
    y_max = max(mean_values['remaining_ngr'].max(), mean_values['remaining_ngr_pred'].max())
    y_text = y_max * 0.85  # Position text at 85% of max height
    
    plt.text(3.5, y_text, '0.05', ha='center', va='center', fontsize=10, fontweight='bold')
    plt.text(10.5, y_text, '0.85', ha='center', va='center', fontsize=10, fontweight='bold')
    plt.text(29.5, y_text, '1.00', ha='center', va='center', fontsize=10, fontweight='bold')
    plt.text(67.5, y_text, '0.90', ha='center', va='center', fontsize=10, fontweight='bold')
    plt.text(135, y_text, '0.60', ha='center', va='center', fontsize=10, fontweight='bold')
    plt.text(225, y_text, '0.40', ha='center', va='center', fontsize=10, fontweight='bold')
    
    # For the last period, calculate center dynamically
    max_days = mean_values['days_since_ftd'].max()
    if max_days > 270:
        last_center = (270 + max_days) / 2
        plt.text(last_center, y_text, '0.30', ha='center', va='center', fontsize=10, fontweight='bold')

    plt.xlabel('Days Since First Time Deposit', fontsize=12)
    plt.ylabel('Average Remaining NGR', fontsize=12)
    title = f'Actual vs Predicted Average Remaining NGR by Days Since FTD {additional_text}'
    plt.title(title, fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    
    # Create metrics text with MAE and MAPE
    metrics_text = "Model Performance Metrics:\n"
    metrics_text += f"Business Optimal MAE: {business_optimal_tmae:.2f}\n"
    metrics_text += f"Business Optimal ME: {business_optimal_me:.2f}\n"
    metrics_text += f"Business Optimal MAPE: {business_optimal_mape:.1f}%\n"
    metrics_text += f"Business Optimal MPE: {business_optimal_mpe:.1f}%\n\n"
    
    # Add mean_abs_error_by_bin to the metrics text
    metrics_text += "Mean Abs Error by Period:\n"
    for _, row in mean_abs_error_by_bin.iterrows():
        metrics_text += f"{row['bin']}: {row['mean_abs_error']:.2f}\n"

    metrics_text += "\nMean ABS PERCENT Error by Period:\n"
    for _, row in mean_abs_error_by_bin.iterrows():
        metrics_text += f"{row['bin']}: {row['mape']:.2f}%\n"
    
    # Add metrics in top right corner
    plt.text(
        0.98, 0.98, metrics_text,
        transform=plt.gca().transAxes,
        verticalalignment='top',
        horizontalalignment='right',
        bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.9),
        fontsize=9, fontweight='bold', fontfamily='monospace'
    )

    plt.tight_layout()
    plt.show()
    plt.close()
    
    # Return dict with MAE and MAPE metrics
    return {
        'Standard MAE': standard_mae,
        'Standard MAPE (%)': round(standard_mape, 1),
        'Business Optimal MAE': round(business_optimal_tmae, 2),
        'Business Optimal MAPE (%)': round(business_optimal_mape, 1)
    }


def get_df_split_by_ngr(
    X_test: pd.DataFrame,
    y_test: pd.Series,
    y_pred: pd.Series,
    id_test: pd.DataFrame,
    high_threshold: float = 5000,
    medium_threshold: float = 100,
    low_threshold: float = -100
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """
    Split data into high, medium, low, and negative NGR groups.
    
    This function segments data based on ngr_after_1_year values for granular
    analysis of model performance across different customer value segments.
    Works with both test and validation sets.
    
    Parameters
    ----------
    X_test : pd.DataFrame
        Features dataframe
    y_test : pd.Series
        Actual target values (remaining_ngr)
    y_pred : pd.Series
        Predicted target values
    id_test : pd.DataFrame
        DataFrame containing 'ngr_after_1_year' column for segmentation
    high_threshold : float, default=5000
        Minimum NGR for 'high' segment
    medium_threshold : float, default=100
        Minimum NGR for 'medium' segment
    low_threshold : float, default=-100
        Minimum NGR for 'low' segment
    
    Returns
    -------
    test_df : pd.DataFrame
        Complete dataframe with all segments
    test_df_high : pd.DataFrame
        High NGR segment (>= high_threshold)
    test_df_medium : pd.DataFrame
        Medium NGR segment (>= medium_threshold and < high_threshold)
    test_df_low : pd.DataFrame
        Low NGR segment (>= low_threshold and < medium_threshold)
    test_df_negative : pd.DataFrame
        Negative NGR segment (< low_threshold)
    
    Examples
    --------
    >>> test_df, high, medium, low, negative = get_df_split_by_ngr(
    ...     X_test, y_test, y_pred, id_test
    ... )
    High NGR proportion: 0.0523
    Medium NGR proportion: 0.2145
    ...
    """
    # Combine X_test and y_test into one dataframe for easier grouping
    test_df = X_test.copy()
    test_df['ngr_after_1_year'] = id_test['ngr_after_1_year']
    test_df['remaining_ngr'] = y_test
    test_df['remaining_ngr_pred'] = y_pred

    # Split data into high, medium, low, and negative NGR groups
    test_df_high = test_df[test_df['ngr_after_1_year'] >= high_threshold]
    test_df_medium = test_df[
        (test_df['ngr_after_1_year'] >= medium_threshold) & 
        (test_df['ngr_after_1_year'] < high_threshold)
    ]
    test_df_low = test_df[
        (test_df['ngr_after_1_year'] >= low_threshold) & 
        (test_df['ngr_after_1_year'] < medium_threshold)
    ]
    test_df_negative = test_df[test_df['ngr_after_1_year'] < low_threshold]

    # Print segment statistics
    total = len(test_df)
    print("\n" + "="*60)
    print("NGR Segmentation Statistics")
    print("="*60)
    
    print(f"\nSegment Proportions:")
    print(f"  High NGR:     {len(test_df_high)/total:6.2%} ({len(test_df_high):>8,} records)")
    print(f"  Medium NGR:   {len(test_df_medium)/total:6.2%} ({len(test_df_medium):>8,} records)")
    print(f"  Low NGR:      {len(test_df_low)/total:6.2%} ({len(test_df_low):>8,} records)")
    print(f"  Negative NGR: {len(test_df_negative)/total:6.2%} ({len(test_df_negative):>8,} records)")
    
    print(f"\nMean NGR after 1 Year by Segment:")
    print(f"  High NGR:     €{test_df_high['ngr_after_1_year'].mean():>10,.2f}")
    print(f"  Medium NGR:   €{test_df_medium['ngr_after_1_year'].mean():>10,.2f}")
    print(f"  Low NGR:      €{test_df_low['ngr_after_1_year'].mean():>10,.2f}")
    print(f"  Negative NGR: €{test_df_negative['ngr_after_1_year'].mean():>10,.2f}")
    print("="*60 + "\n")

    return test_df, test_df_high, test_df_medium, test_df_low, test_df_negative