
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

def plot_binning(bin_stats, feature_name, show=True, filename=None):
    """
    Visualize binning results.
    Supports single dataset (list of dicts) or multiple datasets (dict of lists).
    """
    if not bin_stats:
        print(f"No binning stats available for {feature_name}")
        return

    fig, ax1 = plt.subplots(figsize=(12, 6))

    # Handle multiple datasets (Comparison Mode)
    if isinstance(bin_stats, dict):
        # Assume 'train' is the primary dataset for bars if present, else first one
        primary_key = 'train' if 'train' in bin_stats else list(bin_stats.keys())[0]
        
        # Prepare data for plotting
        datasets = list(bin_stats.keys())
        
        # Use primary dataset for X-axis labels and Bar chart (Count/Distribution)
        primary_stats = pd.DataFrame(bin_stats[primary_key])
        
        # Create label
        if 'min' in primary_stats.columns and 'max' in primary_stats.columns:
            primary_stats['label'] = primary_stats.apply(lambda x: f"[{x['min']}, {x['max']})", axis=1)
        elif 'bin_range' in primary_stats.columns:
            primary_stats['label'] = primary_stats['bin_range']
        else:
            primary_stats['label'] = primary_stats.index
            
        # Plot Bar (Count) for Primary Dataset
        # Use alpha to make it background
        ax1.bar(primary_stats.index, primary_stats['total'], color='skyblue', alpha=0.3, label=f'{primary_key.title()} Count')
        ax1.set_ylabel('Sample Count', color='blue')
        ax1.tick_params(axis='y', labelcolor='blue')
        ax1.set_xticks(primary_stats.index)
        ax1.set_xticklabels(primary_stats['label'], rotation=45, ha='right')
        ax1.set_xlabel('Bin')
        
        # Plot Bad Rates for ALL datasets on secondary axis
        ax2 = ax1.twinx()
        colors = ['red', 'green', 'orange', 'purple', 'brown']
        
        for i, ds_name in enumerate(datasets):
            df = pd.DataFrame(bin_stats[ds_name])
            color = colors[i % len(colors)]
            ax2.plot(df.index, df['bad_rate'], marker='o', linestyle='-', linewidth=2, color=color, label=f'{ds_name.title()} Bad Rate')
            
        ax2.set_ylabel('Bad Rate', color='black')
        ax2.tick_params(axis='y', labelcolor='black')
        
        # Combine legends
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax2.legend(lines1 + lines2, labels1 + labels2, loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=len(datasets)+1)
        
        plt.title(f"Binning Comparison: {feature_name}")

    # Handle single dataset (Standard Mode)
    else:
        df_stats = pd.DataFrame(bin_stats)
        
        # Create label
        if 'min' in df_stats.columns and 'max' in df_stats.columns:
            df_stats['label'] = df_stats.apply(lambda x: f"[{x['min']}, {x['max']})", axis=1)
        elif 'bin_range' in df_stats.columns:
            df_stats['label'] = df_stats['bin_range']
        else:
            df_stats['label'] = df_stats.index
        
        # Bar chart for Total Count
        ax1.bar(df_stats.index, df_stats['total'], color='skyblue', alpha=0.6, label='Count')
        ax1.set_ylabel('Sample Count', color='blue')
        ax1.tick_params(axis='y', labelcolor='blue')
        ax1.set_xticks(df_stats.index)
        ax1.set_xticklabels(df_stats['label'], rotation=45, ha='right')
        ax1.set_xlabel('Bin')
        
        # Line chart for Bad Rate
        ax2 = ax1.twinx()
        ax2.plot(df_stats.index, df_stats['bad_rate'], color='red', marker='o', linestyle='-', linewidth=2, label='Bad Rate')
        ax2.set_ylabel('Bad Rate', color='red')
        ax2.tick_params(axis='y', labelcolor='red')
        
        # Add IV info
        total_iv = df_stats['iv'].sum() if 'iv' in df_stats.columns else 0
        plt.title(f"Binning: {feature_name} (IV: {total_iv:.4f})")
        
        # Combine legends
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax2.legend(lines1 + lines2, labels1 + labels2, loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=2)
    
    # Layout
    plt.tight_layout()
    
    if filename:
        plt.savefig(filename)
        
    if show:
        plt.show()
    else:
        plt.close()
