#!/usr/bin/env python3
"""
Create visualizations and graphs for StoxChai NSE Stock Prediction Models
"""

import json
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

# Set style for better looking plots
plt.style.use('seaborn-v0_8')
sns.set_palette("husl")

def load_training_summary():
    """Load the training summary data"""
    try:
        with open('comprehensive_training_summary.json', 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        print("❌ Training summary not found. Please run training first.")
        return None

def create_performance_comparison():
    """Create performance comparison charts"""
    print("📊 Creating performance comparison charts...")
    
    # Sample performance data (you can replace with actual metrics from training)
    models = ['RandomForest', 'GradientBoosting', 'LinearRegression', 'Ridge', 'Lasso', 'SVR', 'XGBoost', 'LightGBM']
    
    # Sample metrics (replace with actual values from your training)
    mse_scores = [0.85, 0.92, 1.15, 1.12, 1.18, 0.95, 0.88, 0.90]
    mae_scores = [0.72, 0.78, 0.89, 0.87, 0.91, 0.79, 0.74, 0.76]
    r2_scores = [0.92, 0.91, 0.88, 0.89, 0.87, 0.90, 0.91, 0.90]
    
    # Create subplots
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('StoxChai NSE Stock Prediction Models - Performance Comparison', fontsize=16, fontweight='bold')
    
    # 1. MSE Comparison
    axes[0, 0].bar(models, mse_scores, color='skyblue', alpha=0.8)
    axes[0, 0].set_title('Mean Squared Error (Lower is Better)', fontweight='bold')
    axes[0, 0].set_ylabel('MSE')
    axes[0, 0].tick_params(axis='x', rotation=45)
    axes[0, 0].grid(True, alpha=0.3)
    
    # 2. MAE Comparison
    axes[0, 1].bar(models, mae_scores, color='lightcoral', alpha=0.8)
    axes[0, 1].set_title('Mean Absolute Error (Lower is Better)', fontweight='bold')
    axes[0, 1].set_ylabel('MAE')
    axes[0, 1].tick_params(axis='x', rotation=45)
    axes[0, 1].grid(True, alpha=0.3)
    
    # 3. R² Comparison
    axes[1, 0].bar(models, r2_scores, color='lightgreen', alpha=0.8)
    axes[1, 0].set_title('R-squared Score (Higher is Better)', fontweight='bold')
    axes[1, 0].set_ylabel('R²')
    axes[1, 0].tick_params(axis='x', rotation=45)
    axes[1, 0].grid(True, alpha=0.3)
    
    # 4. Combined Performance Heatmap
    performance_data = pd.DataFrame({
        'MSE': mse_scores,
        'MAE': mae_scores,
        'R²': r2_scores
    }, index=models)
    
    # Normalize for better visualization (0-1 scale)
    normalized_data = performance_data.copy()
    normalized_data['MSE'] = 1 - (normalized_data['MSE'] - normalized_data['MSE'].min()) / (normalized_data['MSE'].max() - normalized_data['MSE'].min())
    normalized_data['MAE'] = 1 - (normalized_data['MAE'] - normalized_data['MAE'].min()) / (normalized_data['MAE'].max() - normalized_data['MAE'].min())
    
    sns.heatmap(normalized_data, annot=True, cmap='RdYlGn', ax=axes[1, 1], cbar_kws={'label': 'Normalized Score'})
    axes[1, 1].set_title('Performance Heatmap (Normalized)', fontweight='bold')
    
    plt.tight_layout()
    plt.savefig('model_performance_comparison.png', dpi=300, bbox_inches='tight')
    print("✅ Performance comparison chart saved as 'model_performance_comparison.png'")
    
    return fig

def create_feature_importance_chart():
    """Create feature importance visualization"""
    print("🔍 Creating feature importance chart...")
    
    # Feature names in order
    features = [
        'OpnPric', 'HghPric', 'LwPric', 'LastPric', 'PrvsClsgPric',
        'Price_Range', 'Price_Change', 'Price_Change_Pct', 'Volume_Price_Ratio',
        'SMA_5', 'SMA_20', 'Price_Momentum', 'Volume_MA', 'Volume_Ratio',
        'TtlTradgVol', 'TtlTrfVal'
    ]
    
    # Sample feature importance scores (RandomForest style - you can replace with actual values)
    importance_scores = [0.12, 0.11, 0.10, 0.09, 0.08, 0.07, 0.06, 0.05, 0.04, 0.03, 0.02, 0.01, 0.01, 0.01, 0.01, 0.01]
    
    # Create feature importance plot
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Sort features by importance
    feature_importance = pd.DataFrame({
        'Feature': features,
        'Importance': importance_scores
    }).sort_values('Importance', ascending=True)
    
    # Create horizontal bar chart
    bars = ax.barh(range(len(feature_importance)), feature_importance['Importance'], 
                   color='steelblue', alpha=0.8)
    
    # Customize the plot
    ax.set_yticks(range(len(feature_importance)))
    ax.set_yticklabels(feature_importance['Feature'])
    ax.set_xlabel('Feature Importance Score')
    ax.set_title('Feature Importance for Stock Price Prediction', fontsize=16, fontweight='bold')
    ax.grid(True, alpha=0.3)
    
    # Add value labels on bars
    for i, bar in enumerate(bars):
        width = bar.get_width()
        ax.text(width + 0.001, bar.get_y() + bar.get_height()/2, 
                f'{width:.3f}', ha='left', va='center', fontweight='bold')
    
    plt.tight_layout()
    plt.savefig('feature_importance.png', dpi=300, bbox_inches='tight')
    print("✅ Feature importance chart saved as 'feature_importance.png'")
    
    return fig

def create_data_overview():
    """Create data overview visualization"""
    print("📈 Creating data overview charts...")
    
    # Create subplots
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle('NSE Bhavcopy Data Overview', fontsize=16, fontweight='bold')
    
    # 1. Data Timeline
    dates = pd.date_range('2025-01-01', '2025-08-20', freq='D')
    trading_days = [d for d in dates if d.weekday() < 5]  # Weekdays only
    
    axes[0, 0].plot(trading_days, range(len(trading_days)), marker='o', linewidth=2, markersize=4)
    axes[0, 0].set_title('Trading Days Timeline', fontweight='bold')
    axes[0, 0].set_xlabel('Date')
    axes[0, 0].set_ylabel('Trading Day Number')
    axes[0, 0].tick_params(axis='x', rotation=45)
    axes[0, 0].grid(True, alpha=0.3)
    
    # 2. Stock Coverage
    stock_counts = [1000, 1500, 2000, 2500, 3000, 3257]  # Sample progression
    months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun-Aug']
    
    axes[0, 1].bar(months, stock_counts, color='lightblue', alpha=0.8)
    axes[0, 1].set_title('Cumulative Stock Coverage', fontweight='bold')
    axes[0, 1].set_ylabel('Number of Stocks')
    axes[0, 1].grid(True, alpha=0.3)
    
    # 3. Data Volume Distribution
    data_volumes = [50000, 75000, 100000, 125000, 150000, 175000, 200000, 225000, 250000, 275000, 300000, 325000, 350000, 375000, 400000, 425000, 450000, 464548]
    axes[1, 0].hist(data_volumes, bins=20, color='lightgreen', alpha=0.8, edgecolor='black')
    axes[1, 0].set_title('Data Volume Distribution', fontweight='bold')
    axes[1, 0].set_xlabel('Number of Records')
    axes[1, 0].set_ylabel('Frequency')
    axes[1, 0].grid(True, alpha=0.3)
    
    # 4. Feature Categories
    feature_categories = ['Price Data', 'Technical Indicators', 'Moving Averages', 'Volume Data']
    feature_counts = [5, 4, 2, 5]
    
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4']
    wedges, texts, autotexts = axes[1, 1].pie(feature_counts, labels=feature_categories, 
                                               autopct='%1.1f%%', colors=colors, startangle=90)
    axes[1, 1].set_title('Feature Categories Distribution', fontweight='bold')
    
    plt.tight_layout()
    plt.savefig('data_overview.png', dpi=300, bbox_inches='tight')
    print("✅ Data overview charts saved as 'data_overview.png'")
    
    return fig

def create_model_architecture_diagram():
    """Create model architecture visualization"""
    print("🏗️ Creating model architecture diagram...")
    
    fig, ax = plt.subplots(figsize=(14, 10))
    
    # Define model architectures
    models = {
        'RandomForest': {'type': 'Ensemble', 'estimators': 100, 'depth': 10},
        'GradientBoosting': {'type': 'Boosting', 'estimators': 100, 'learning_rate': 0.1},
        'LinearRegression': {'type': 'Linear', 'regularization': 'None'},
        'Ridge': {'type': 'Linear', 'regularization': 'L2'},
        'Lasso': {'type': 'Linear', 'regularization': 'L1'},
        'SVR': {'type': 'Kernel', 'kernel': 'RBF'},
        'XGBoost': {'type': 'Boosting', 'estimators': 100, 'learning_rate': 0.1},
        'LightGBM': {'type': 'Boosting', 'estimators': 100, 'learning_rate': 0.1}
    }
    
    # Create architecture visualization
    y_pos = np.arange(len(models))
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7', '#DDA0DD', '#98D8C8', '#F7DC6F']
    
    bars = ax.barh(y_pos, [1]*len(models), color=colors, alpha=0.8)
    
    # Add model details
    for i, (model, details) in enumerate(models.items()):
        ax.text(0.1, i, f"{model}\n{details['type']}", va='center', fontweight='bold')
        if details['type'] == 'Ensemble':
            ax.text(0.6, i, f"Estimators: {details['estimators']}", va='center')
        elif details['type'] == 'Boosting':
            ax.text(0.6, i, f"Estimators: {details['estimators']}\nLR: {details['learning_rate']}", va='center')
        elif details['type'] == 'Linear':
            ax.text(0.6, i, f"Reg: {details['regularization']}", va='center')
        elif details['type'] == 'Kernel':
            ax.text(0.6, i, f"Kernel: {details['kernel']}", va='center')
    
    ax.set_yticks(y_pos)
    ax.set_yticklabels([''] * len(models))
    ax.set_xlim(0, 1)
    ax.set_title('Model Architecture Overview', fontsize=16, fontweight='bold')
    ax.set_xlabel('Model Complexity')
    ax.grid(True, alpha=0.3)
    
    # Remove x-axis ticks
    ax.set_xticks([])
    
    plt.tight_layout()
    plt.savefig('model_architecture.png', dpi=300, bbox_inches='tight')
    print("✅ Model architecture diagram saved as 'model_architecture.png'")
    
    return fig

def main():
    """Main function to create all visualizations"""
    print("🎨 Creating StoxChai Model Visualizations")
    print("=" * 50)
    
    try:
        # Create all visualizations
        create_performance_comparison()
        create_feature_importance_chart()
        create_data_overview()
        create_model_architecture_diagram()
        
        print("\n🎉 All visualizations created successfully!")
        print("📁 Generated files:")
        print("  - model_performance_comparison.png")
        print("  - feature_importance.png")
        print("  - data_overview.png")
        print("  - model_architecture.png")
        
        return True
        
    except Exception as e:
        print(f"❌ Error creating visualizations: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    if success:
        print("\n✨ Visualizations are ready for Hugging Face!")
    else:
        print("\n💥 Some visualizations failed to create.")