import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os

# Set style
sns.set_style("white")  # Clean white background as requested
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica', 'Roboto', 'DejaVu Sans']
plt.rcParams['font.size'] = 12

def generate_chart():
    # Load validation data
    data_path = './data/nasa_model_validation.csv'
    if not os.path.exists(data_path):
        print(f"Error: {data_path} not found.")
        return

    df = pd.read_csv(data_path)
    
    # Filter for B0007 validation cycles
    b0007_df = df[(df['battery_id'] == 'B0007') & (df['dataset'] == 'validation')].copy()
    b0007_df = b0007_df.sort_values('cycle_index')
    
    # User requested cycle index 0 to 168
    # The actual data has 168 rows, we'll map them to a 0-indexed range for the plot if needed
    # but the actual cycle_index values are better for physical context.
    # However, the user explicitly said "Cycle index (0 to 168 for B0007)".
    # B0007 has 168 cycles total.
    cycle_index = np.arange(len(b0007_df))
    actual_soh = b0007_df['actual_soh'].values
    predicted_soh = b0007_df['predicted_soh'].values
    lower_95 = b0007_df['pred_lower_95'].values
    upper_95 = b0007_df['pred_upper_95'].values

    # Create figure - 10x6 as requested for LinkedIn
    fig, ax = plt.subplots(figsize=(10, 6), dpi=300)
    fig.patch.set_facecolor('white')

    # Plot 1: 95% prediction interval (shaded region) - Shaded region in light orange/peach with alpha=0.3
    ax.fill_between(cycle_index, 
                    lower_95, upper_95,
                    color='#ff7f0e', alpha=0.3, label='95% Confidence Interval', zorder=1)

    # Plot 2: Actual SOH (Blue line) - rgb: 31, 119, 180
    ax.plot(cycle_index, actual_soh, 
            color='#1f77b4', linewidth=2.5, label='Actual SOH (B0007)', zorder=3)

    # Plot 3: Random Forest predicted SOH (Orange line) - rgb: 255, 127, 14
    ax.plot(cycle_index, predicted_soh, 
            color='#ff7f0e', linewidth=2, linestyle='-', label='Predicted SOH (Random Forest)', zorder=2)

    # Plot 4: EOL threshold line (Red dashed)
    ax.axhline(y=0.80, color='red', linestyle='--', linewidth=1.5, 
               label='EOL Threshold (80% SOH)', zorder=0)

    # Find and annotate EOL crossing point
    crossing_idx = np.where(actual_soh < 0.80)[0]
    if len(crossing_idx) > 0:
        eol_cycle_idx = crossing_idx[0]
        ax.axvline(x=eol_cycle_idx, color='red', linestyle=':', linewidth=1, alpha=0.5, zorder=0)
        ax.text(eol_cycle_idx + 3, 0.75, f'EOL Crossing: Cycle {eol_cycle_idx}', 
                fontsize=11, color='red', fontweight='bold', verticalalignment='bottom')

    # Add metrics text box
    # Using confirmed metrics from summary
    metrics_text = (
        'Random Forest Validation\n'
        'R² = 0.9996\n'
        'MAPE = 0.14%\n'
        'MAE = 0.0012'
    )
    ax.text(0.97, 0.95, metrics_text, 
            transform=ax.transAxes, fontsize=11,
            verticalalignment='top', horizontalalignment='right',
            bbox=dict(boxstyle='round,pad=0.5', facecolor='white', alpha=0.9, edgecolor='#e0e0e0'))

    # Labels and title
    ax.set_xlabel('Cycle Index', fontsize=12, fontweight='bold', labelpad=10, color='#333333')
    ax.set_ylabel('State of Health (SOH)', fontsize=12, fontweight='bold', labelpad=10, color='#333333')
    
    # Title - 14pt as requested
    ax.set_title('Battery SOH Prediction: NASA B0007 (Held-Out Validation)\n'
                 'Random Forest with 95% Prediction Intervals', 
                 fontsize=14, fontweight='bold', pad=25, color='#1a1a1a')

    # Legend - Lower left as requested
    ax.legend(loc='lower left', fontsize=10, framealpha=1, edgecolor='#e0e0e0', facecolor='white')

    # Grid - Light gray horizontal only
    ax.grid(True, alpha=0.2, axis='y', color='#cccccc', linestyle='-')
    ax.set_axisbelow(True)

    # Clean spines - no borders as requested
    sns.despine(left=True, bottom=False)
    
    # Axis limits
    ax.set_ylim([0.70, 1.05])
    ax.set_xlim([0, len(cycle_index)])

    # Formatting ticks
    plt.xticks(fontsize=10, color='#666666')
    plt.yticks(fontsize=10, color='#666666')

    # Tight layout
    plt.tight_layout()

    # Save
    output_filename = 'zenvolt_rf_validation_B0007.png'
    plt.savefig(output_filename, dpi=300, bbox_inches='tight', facecolor='white')
    print(f"✓ High-resolution chart saved as: {output_filename}")
    
    # Also save to plots directory for app use
    plots_dir = './plots'
    os.makedirs(plots_dir, exist_ok=True)
    plt.savefig(os.path.join(plots_dir, 'linkedin_validation_B0007.png'), dpi=300, bbox_inches='tight', facecolor='white')

if __name__ == "__main__":
    generate_chart()