pdf_chart_generator.py

"""
PDF Chart Generator
Generates matplotlib charts as base64 images for embedding in PDF reports.

MIT/Stanford Quality:
- Publication-quality charts
- Optimized for print (300 DPI)
- Professional styling
- Fast rendering
"""

import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import BytesIO
import base64
from typing import List, Dict, Optional
import seaborn as sns

# Set style
plt.style.use('seaborn-v0_8-darkgrid')
sns.set_palette("husl")


def generate_health_timeline_chart(
    history_data: List[Dict],
    asset_name: str,
    width: int = 8,
    height: int = 4,
    dpi: int = 150
) -> str:
    """
    Generate health score timeline chart.
    
    Parameters:
    -----------
    history_data : list
        List of prediction dictionaries with health_score and prediction_date
    asset_name : str
        Asset name for title
    width, height : int
        Figure dimensions in inches
    dpi : int
        Resolution (150 for screen, 300 for print)
    
    Returns:
    --------
    str: Base64-encoded PNG image
    """
    fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
    
    # Extract data
    dates = [pd.to_datetime(p['prediction_date']) for p in history_data]
    scores = [p['health_score'] for p in history_data]
    lower_bounds = [p.get('soh_lower_95', 0) * 100 for p in history_data]
    upper_bounds = [p.get('soh_upper_95', 1) * 100 for p in history_data]
    
    # Plot
    ax.plot(dates, scores, 'o-', linewidth=2.5, markersize=8, label='Health Score', color='#22c55e')
    ax.fill_between(dates, lower_bounds, upper_bounds, alpha=0.2, color='#22c55e', label='95% Confidence')
    
    # Styling
    ax.set_xlabel('Date', fontsize=11, fontweight='bold')
    ax.set_ylabel('Health Score', fontsize=11, fontweight='bold')
    ax.set_title(f'Health Timeline - {asset_name}', fontsize=13, fontweight='bold', pad=15)
    ax.grid(True, alpha=0.3, linestyle='--')
    ax.legend(loc='best', frameon=True, shadow=True)
    ax.set_ylim([0, 105])
    
    # Add threshold line
    ax.axhline(y=80, color='#ef4444', linestyle='--', linewidth=1.5, alpha=0.7, label='EOL Threshold (80%)')
    
    plt.tight_layout()
    
    # Convert to base64
    buffer = BytesIO()
    plt.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buffer.seek(0)
    
    img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    return f"data:image/png;base64,{img_base64}"


def generate_degradation_curve_chart(
    history_data: List[Dict],
    asset_name: str,
    width: int = 8,
    height: int = 4,
    dpi: int = 150
) -> str:
    """Generate capacity degradation curve chart."""
    fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
    
    # Extract data
    cycles = [p['cycle_index'] for p in history_data]
    soh_values = [p['current_soh'] for p in history_data]
    
    # Plot actual
    ax.plot(cycles, soh_values, 'o-', linewidth=2.5, markersize=8, color='#3b82f6', label='Actual SOH')
    
    # Fit trend line
    if len(cycles) > 2:
        z = np.polyfit(cycles, soh_values, 2)  # Quadratic fit
        p = np.poly1d(z)
        cycle_range = np.linspace(min(cycles), max(cycles) * 1.5, 100)
        ax.plot(cycle_range, p(cycle_range), '--', linewidth=2, alpha=0.7, color='#ef4444', label='Projected Trend')
    
    # Styling
    ax.set_xlabel('Cycle Index', fontsize=11, fontweight='bold')
    ax.set_ylabel('State of Health (SOH)', fontsize=11, fontweight='bold')
    ax.set_title(f'Degradation Curve - {asset_name}', fontsize=13, fontweight='bold', pad=15)
    ax.grid(True, alpha=0.3, linestyle='--')
    ax.legend(loc='best', frameon=True, shadow=True)
    ax.set_ylim([0, 1.05])
    
    # Add EOL threshold
    ax.axhline(y=0.8, color='#f59e0b', linestyle='--', linewidth=1.5, alpha=0.7, label='EOL (80%)')
    
    plt.tight_layout()
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buffer.seek(0)
    
    return f"data:image/png;base64,{base64.b64encode(buffer.read()).decode('utf-8')}"


def generate_voltage_temperature_chart(
    voltage_data: List[float],
    temperature_data: List[float],
    asset_name: str,
    width: int = 8,
    height: int = 4,
    dpi: int = 150
) -> str:
    """Generate voltage and temperature dual-axis chart."""
    fig, ax1 = plt.subplots(figsize=(width, height), dpi=dpi)
    
    time_points = list(range(len(voltage_data)))
    
    # Voltage on left axis
    color = '#3b82f6'
    ax1.set_xlabel('Time Point', fontsize=11, fontweight='bold')
    ax1.set_ylabel('Voltage (V)', fontsize=11, fontweight='bold', color=color)
    ax1.plot(time_points, voltage_data, 'o-', color=color, linewidth=2, markersize=6, label='Voltage')
    ax1.tick_params(axis='y', labelcolor=color)
    ax1.grid(True, alpha=0.3, linestyle='--')
    
    # Temperature on right axis
    ax2 = ax1.twinx()
    color = '#ef4444'
    ax2.set_ylabel('Temperature (°C)', fontsize=11, fontweight='bold', color=color)
    ax2.plot(time_points, temperature_data, 's-', color=color, linewidth=2, markersize=6, label='Temperature')
    ax2.tick_params(axis='y', labelcolor=color)
    
    # Title
    ax1.set_title(f'Voltage & Temperature Trends - {asset_name}', fontsize=13, fontweight='bold', pad=15)
    
    # Combined legend
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='best', frameon=True, shadow=True)
    
    plt.tight_layout()
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buffer.seek(0)
    
    return f"data:image/png;base64,{base64.b64encode(buffer.read()).decode('utf-8')}"


def generate_feature_importance_chart(
    features: List[Dict],
    asset_name: str,
    width: int = 8,
    height: int = 5,
    dpi: int = 150
) -> str:
    """Generate feature importance bar chart."""
    fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
    
    # Extract top features
    names = [f['name'][:30] for f in features[:10]]  # Truncate long names
    importances = [f['importance'] for f in features[:10]]
    
    # Create horizontal bar chart
    y_pos = np.arange(len(names))
    colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(names)))
    
    bars = ax.barh(y_pos, importances, color=colors, edgecolor='black', linewidth=0.5)
    
    # Styling
    ax.set_yticks(y_pos)
    ax.set_yticklabels(names, fontsize=9)
    ax.invert_yaxis()  # Top features at top
    ax.set_xlabel('Feature Importance', fontsize=11, fontweight='bold')
    ax.set_title(f'Top Predictive Features - {asset_name}', fontsize=13, fontweight='bold', pad=15)
    ax.grid(True, alpha=0.3, axis='x', linestyle='--')
    
    # Add value labels on bars
    for i, (bar, val) in enumerate(zip(bars, importances)):
        ax.text(val, bar.get_y() + bar.get_height()/2, f' {val:.3f}', 
                va='center', fontsize=8, fontweight='bold')
    
    plt.tight_layout()
    
    buffer = BytesIO()
    plt.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buffer.seek(0)
    
    return f"data:image/png;base64,{base64.b64encode(buffer.read()).decode('utf-8')}"


def generate_all_charts_for_asset(
    asset: Dict,
    history: List[Dict]
) -> Dict[str, str]:
    """
    Generate all charts for an asset.
    
    Parameters:
    -----------
    asset : dict
        Asset data
    history : list
        Historical predictions
    
    Returns:
    --------
    dict: {chart_name: base64_image}
    """
    charts = {}
    
    # 1. Health timeline
    if history and len(history) > 0:
        charts['health_timeline'] = generate_health_timeline_chart(
            history, 
            asset.get('projectName', asset.get('project_name', 'Unknown'))
        )
    
    # 2. Degradation curve
    if history and len(history) > 0:
        charts['degradation_curve'] = generate_degradation_curve_chart(
            history,
            asset.get('projectName', asset.get('project_name', 'Unknown'))
        )
    
    # 3. Voltage/Temperature (if available)
    if 'voltageData' in asset or 'voltage' in asset:
        voltage_data = asset.get('voltageData', [asset.get('voltage', 700)] * 8)
        temp_data = asset.get('temperatureData', [asset.get('temperature', 25)] * 8)
        
        if len(voltage_data) > 0 and len(temp_data) > 0:
            charts['voltage_temperature'] = generate_voltage_temperature_chart(
                voltage_data,
                temp_data,
                asset.get('projectName', asset.get('project_name', 'Unknown'))
            )
    
    # 4. Feature importance (if available)
    if history and len(history) > 0 and 'shap_top_features' in history[-1]:
        features_str = history[-1]['shap_top_features']
        if features_str:
            import json
            features = json.loads(features_str) if isinstance(features_str, str) else features_str
            if features and len(features) > 0:
                charts['feature_importance'] = generate_feature_importance_chart(
                    features,
                    asset.get('projectName', asset.get('project_name', 'Unknown'))
                )
    
    return charts


if __name__ == "__main__":
    # Test chart generation
    print("Testing PDF chart generation...")
    
    # Mock data
    history = [
        {'prediction_date': '2025-10-01', 'health_score': 95, 'current_soh': 0.95, 'soh_lower_95': 0.93, 'soh_upper_95': 0.97, 'cycle_index': 40},
        {'prediction_date': '2025-11-01', 'health_score': 90, 'current_soh': 0.90, 'soh_lower_95': 0.88, 'soh_upper_95': 0.92, 'cycle_index': 80},
        {'prediction_date': '2025-12-01', 'health_score': 85, 'current_soh': 0.85, 'soh_lower_95': 0.83, 'soh_upper_95': 0.87, 'cycle_index': 120},
        {'prediction_date': '2026-01-01', 'health_score': 80, 'current_soh': 0.80, 'soh_lower_95': 0.78, 'soh_upper_95': 0.82, 'cycle_index': 160},
    ]
    
    asset = {'projectName': 'Test Battery', 'voltage': 700, 'temperature': 25}
    
    # Generate charts
    chart1 = generate_health_timeline_chart(history, "Test Battery")
    print(f"✓ Health timeline: {len(chart1)} chars")
    
    chart2 = generate_degradation_curve_chart(history, "Test Battery")
    print(f"✓ Degradation curve: {len(chart2)} chars")
    
    chart3 = generate_voltage_temperature_chart([700, 705, 710, 708, 710, 712, 710, 710], 
                                                  [25, 26, 27, 27, 26, 28, 27, 26], 
                                                  "Test Battery")
    print(f"✓ Voltage/Temperature: {len(chart3)} chars")
    
    print("\n✅ Chart generation test complete!")
    print("   All charts ready for PDF embedding")