"""
Visualization and Validation of NN-Only Optimization Results

This script:
1. Plots the Pareto front from NN optimization
2. Compares NN predictions vs actual FEA data
3. Shows prediction confidence and error analysis
4. Validates selected NN-optimal designs with FEA data
"""
import sys
from pathlib import Path
import json
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend for headless operation
import matplotlib.pyplot as plt
import optuna

# Add project paths
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from optimization_engine.processors.surrogates.simple_mlp_surrogate import SimpleSurrogate

def load_fea_data_from_database(db_path: str, study_name: str):
    """Load actual FEA results from database for comparison."""
    storage = optuna.storages.RDBStorage(f"sqlite:///{db_path}")
    study = optuna.load_study(study_name=study_name, storage=storage)

    completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]

    data = []
    for trial in completed_trials:
        if len(trial.values) < 2:
            continue
        mass = trial.values[0]
        # Handle both formats: some stores -freq (for minimization), some store +freq
        raw_freq = trial.values[1]
        # If frequency is stored as negative (minimization convention), flip it
        frequency = -raw_freq if raw_freq < 0 else raw_freq

        # Skip invalid
        if np.isinf(mass) or np.isinf(frequency) or frequency <= 0:
            continue

        data.append({
            'params': trial.params,
            'mass': mass,
            'frequency': frequency,
            'max_displacement': trial.user_attrs.get('max_displacement', 0),
            'max_stress': trial.user_attrs.get('max_stress', 0),
        })

    return data

def plot_pareto_comparison(nn_results, fea_data, surrogate):
    """Plot Pareto fronts: NN optimization vs FEA data."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))

    # Extract NN Pareto front
    nn_mass = [d['mass'] for d in nn_results['pareto_designs']]
    nn_freq = [d['frequency'] for d in nn_results['pareto_designs']]

    # Extract FEA data
    fea_mass = [d['mass'] for d in fea_data]
    fea_freq = [d['frequency'] for d in fea_data]

    # 1. Pareto Front Comparison
    ax = axes[0, 0]
    ax.scatter(fea_mass, fea_freq, alpha=0.5, label='FEA Trials', c='blue', s=30)
    ax.scatter(nn_mass, nn_freq, alpha=0.7, label='NN Pareto Front', c='red', s=20, marker='x')
    ax.set_xlabel('Mass (g)')
    ax.set_ylabel('Frequency (Hz)')
    ax.set_title('Pareto Front: NN Optimization vs FEA Data')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # 2. NN Prediction Error on FEA Data
    ax = axes[0, 1]
    nn_pred_mass = []
    nn_pred_freq = []
    actual_mass = []
    actual_freq = []

    for d in fea_data:
        pred = surrogate.predict(d['params'])
        nn_pred_mass.append(pred['mass'])
        nn_pred_freq.append(pred['frequency'])
        actual_mass.append(d['mass'])
        actual_freq.append(d['frequency'])

    # Mass prediction error
    mass_errors = np.array(nn_pred_mass) - np.array(actual_mass)
    freq_errors = np.array(nn_pred_freq) - np.array(actual_freq)

    scatter = ax.scatter(actual_mass, actual_freq, c=np.abs(mass_errors),
                         cmap='RdYlGn_r', s=50, alpha=0.7)
    plt.colorbar(scatter, ax=ax, label='Mass Prediction Error (g)')
    ax.set_xlabel('Actual Mass (g)')
    ax.set_ylabel('Actual Frequency (Hz)')
    ax.set_title('FEA Points Colored by NN Mass Error')
    ax.grid(True, alpha=0.3)

    # 3. Prediction vs Actual (Mass)
    ax = axes[1, 0]
    ax.scatter(actual_mass, nn_pred_mass, alpha=0.6, s=30)
    min_val, max_val = min(actual_mass), max(actual_mass)
    ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Prediction')
    ax.set_xlabel('Actual Mass (g)')
    ax.set_ylabel('NN Predicted Mass (g)')
    ax.set_title(f'Mass: NN vs FEA\nMAPE: {np.mean(np.abs(mass_errors)/np.array(actual_mass))*100:.1f}%')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # 4. Prediction vs Actual (Frequency)
    ax = axes[1, 1]
    ax.scatter(actual_freq, nn_pred_freq, alpha=0.6, s=30)
    min_val, max_val = min(actual_freq), max(actual_freq)
    ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Prediction')
    ax.set_xlabel('Actual Frequency (Hz)')
    ax.set_ylabel('NN Predicted Frequency (Hz)')
    ax.set_title(f'Frequency: NN vs FEA\nMAPE: {np.mean(np.abs(freq_errors)/np.array(actual_freq))*100:.1f}%')
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(project_root / 'nn_optimization_analysis.png', dpi=150)
    print(f"Saved: nn_optimization_analysis.png")
    plt.close()

    return mass_errors, freq_errors

def plot_design_space_coverage(nn_results, fea_data):
    """Show how well NN explored the design space."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    param_names = ['beam_half_core_thickness', 'beam_face_thickness',
                   'holes_diameter', 'hole_count']

    for idx, (param1, param2) in enumerate([
        ('beam_half_core_thickness', 'beam_face_thickness'),
        ('holes_diameter', 'hole_count'),
        ('beam_half_core_thickness', 'holes_diameter'),
        ('beam_face_thickness', 'hole_count')
    ]):
        ax = axes[idx // 2, idx % 2]

        # FEA data points
        fea_p1 = [d['params'].get(param1, 0) for d in fea_data]
        fea_p2 = [d['params'].get(param2, 0) for d in fea_data]

        # NN Pareto designs
        nn_p1 = [d['params'].get(param1, 0) for d in nn_results['pareto_designs']]
        nn_p2 = [d['params'].get(param2, 0) for d in nn_results['pareto_designs']]

        ax.scatter(fea_p1, fea_p2, alpha=0.4, label='FEA Trials', c='blue', s=30)
        ax.scatter(nn_p1, nn_p2, alpha=0.7, label='NN Pareto', c='red', s=20, marker='x')

        ax.set_xlabel(param1.replace('_', ' ').title())
        ax.set_ylabel(param2.replace('_', ' ').title())
        ax.legend()
        ax.grid(True, alpha=0.3)

    plt.suptitle('Design Space Coverage: FEA vs NN Pareto Designs', fontsize=14)
    plt.tight_layout()
    plt.savefig(project_root / 'nn_design_space_coverage.png', dpi=150)
    print(f"Saved: nn_design_space_coverage.png")
    plt.close()

def plot_error_distribution(mass_errors, freq_errors):
    """Plot error distributions to understand prediction confidence."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    # Mass error histogram
    ax = axes[0]
    ax.hist(mass_errors, bins=30, edgecolor='black', alpha=0.7)
    ax.axvline(0, color='r', linestyle='--', label='Zero Error')
    ax.axvline(np.mean(mass_errors), color='g', linestyle='-',
               label=f'Mean: {np.mean(mass_errors):.1f}g')
    ax.set_xlabel('Mass Prediction Error (g)')
    ax.set_ylabel('Count')
    ax.set_title(f'Mass Error Distribution\nStd: {np.std(mass_errors):.1f}g')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # Frequency error histogram
    ax = axes[1]
    ax.hist(freq_errors, bins=30, edgecolor='black', alpha=0.7)
    ax.axvline(0, color='r', linestyle='--', label='Zero Error')
    ax.axvline(np.mean(freq_errors), color='g', linestyle='-',
               label=f'Mean: {np.mean(freq_errors):.1f}Hz')
    ax.set_xlabel('Frequency Prediction Error (Hz)')
    ax.set_ylabel('Count')
    ax.set_title(f'Frequency Error Distribution\nStd: {np.std(freq_errors):.1f}Hz')
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(project_root / 'nn_error_distribution.png', dpi=150)
    print(f"Saved: nn_error_distribution.png")
    plt.close()

def find_closest_fea_validation(nn_pareto, fea_data):
    """Find FEA data points closest to NN Pareto designs for validation."""
    print("\n" + "="*70)
    print("VALIDATION: Comparing NN Pareto Designs to Nearest FEA Points")
    print("="*70)

    # Get unique NN Pareto designs (remove duplicates)
    seen = set()
    unique_pareto = []
    for d in nn_pareto:
        key = (round(d['mass'], 1), round(d['frequency'], 1))
        if key not in seen:
            seen.add(key)
            unique_pareto.append(d)

    # Sort by mass
    unique_pareto.sort(key=lambda x: x['mass'])

    # Sample 10 designs across the Pareto front
    indices = np.linspace(0, len(unique_pareto)-1, min(10, len(unique_pareto)), dtype=int)
    sampled_designs = [unique_pareto[i] for i in indices]

    print(f"\nSampled {len(sampled_designs)} designs from NN Pareto front:")
    print("-"*70)

    for i, nn_design in enumerate(sampled_designs):
        # Find closest FEA point (by parameter distance)
        min_dist = float('inf')
        closest_fea = None

        for fea in fea_data:
            dist = sum((nn_design['params'].get(k, 0) - fea['params'].get(k, 0))**2
                      for k in nn_design['params'])
            if dist < min_dist:
                min_dist = dist
                closest_fea = fea

        if closest_fea:
            mass_err = nn_design['mass'] - closest_fea['mass']
            freq_err = nn_design['frequency'] - closest_fea['frequency']

            print(f"\n{i+1}. NN Design: mass={nn_design['mass']:.1f}g, freq={nn_design['frequency']:.1f}Hz")
            print(f"   Closest FEA: mass={closest_fea['mass']:.1f}g, freq={closest_fea['frequency']:.1f}Hz")
            print(f"   Error: mass={mass_err:+.1f}g ({mass_err/closest_fea['mass']*100:+.1f}%), "
                  f"freq={freq_err:+.1f}Hz ({freq_err/closest_fea['frequency']*100:+.1f}%)")
            print(f"   Parameter Distance: {np.sqrt(min_dist):.2f}")

def print_optimization_summary(nn_results, fea_data, mass_errors, freq_errors):
    """Print summary statistics."""
    print("\n" + "="*70)
    print("OPTIMIZATION SUMMARY")
    print("="*70)

    print(f"\n1. NN Optimization Performance:")
    print(f"   - Trials: {nn_results['n_trials']}")
    print(f"   - Time: {nn_results['total_time_s']:.1f}s ({nn_results['trials_per_second']:.1f} trials/sec)")
    print(f"   - Pareto Front Size: {nn_results['pareto_front_size']}")

    print(f"\n2. NN Prediction Accuracy (on FEA data):")
    print(f"   - Mass MAPE: {np.mean(np.abs(mass_errors)/np.array([d['mass'] for d in fea_data]))*100:.1f}%")
    print(f"   - Mass Mean Error: {np.mean(mass_errors):.1f}g (Std: {np.std(mass_errors):.1f}g)")
    print(f"   - Freq MAPE: {np.mean(np.abs(freq_errors)/np.array([d['frequency'] for d in fea_data]))*100:.1f}%")
    print(f"   - Freq Mean Error: {np.mean(freq_errors):.1f}Hz (Std: {np.std(freq_errors):.1f}Hz)")

    print(f"\n3. Design Space:")
    nn_mass = [d['mass'] for d in nn_results['pareto_designs']]
    nn_freq = [d['frequency'] for d in nn_results['pareto_designs']]
    fea_mass = [d['mass'] for d in fea_data]
    fea_freq = [d['frequency'] for d in fea_data]

    print(f"   - NN Pareto Mass Range: {min(nn_mass):.1f}g - {max(nn_mass):.1f}g")
    print(f"   - NN Pareto Freq Range: {min(nn_freq):.1f}Hz - {max(nn_freq):.1f}Hz")
    print(f"   - FEA Mass Range: {min(fea_mass):.1f}g - {max(fea_mass):.1f}g")
    print(f"   - FEA Freq Range: {min(fea_freq):.1f}Hz - {max(fea_freq):.1f}Hz")

    print(f"\n4. Confidence Assessment:")
    if np.std(mass_errors) < 100 and np.std(freq_errors) < 5:
        print("   [OK] LOW prediction variance - NN is fairly confident")
    else:
        print("   [!] HIGH prediction variance - consider more training data")

    if abs(np.mean(mass_errors)) < 50:
        print("   [OK] LOW mass bias - predictions are well-centered")
    else:
        print(f"   [!] Mass bias detected ({np.mean(mass_errors):+.1f}g) - systematic error")

    if abs(np.mean(freq_errors)) < 2:
        print("   [OK] LOW frequency bias - predictions are well-centered")
    else:
        print(f"   [!] Frequency bias detected ({np.mean(freq_errors):+.1f}Hz) - systematic error")

def main():
    print("="*70)
    print("NN Optimization Visualization & Validation")
    print("="*70)

    # Load NN optimization results
    results_file = project_root / "nn_optimization_results.json"
    if not results_file.exists():
        print(f"ERROR: {results_file} not found. Run NN optimization first.")
        return

    with open(results_file) as f:
        nn_results = json.load(f)

    print(f"\nLoaded NN results: {nn_results['n_trials']} trials, "
          f"{nn_results['pareto_front_size']} Pareto designs")

    # Load surrogate model
    model_path = project_root / "simple_mlp_surrogate.pt"
    if not model_path.exists():
        print(f"ERROR: {model_path} not found.")
        return

    surrogate = SimpleSurrogate.load(model_path)
    print(f"Loaded surrogate model with {len(surrogate.design_var_names)} design variables")

    # Load FEA data from original study
    # Try each database path with its matching study name
    db_options = [
        (project_root / "studies/uav_arm_optimization/2_results/study.db", "uav_arm_optimization"),
        (project_root / "studies/uav_arm_atomizerfield_test/2_results/study.db", "uav_arm_atomizerfield_test"),
    ]

    db_path = None
    study_name = None
    for path, name in db_options:
        if path.exists():
            db_path = path
            study_name = name
            break

    if db_path and study_name:
        fea_data = load_fea_data_from_database(str(db_path), study_name)
        print(f"Loaded {len(fea_data)} FEA data points from {study_name}")
    else:
        print("WARNING: No FEA database found. Using only NN results.")
        fea_data = []

    if fea_data:
        # Generate all plots
        mass_errors, freq_errors = plot_pareto_comparison(nn_results, fea_data, surrogate)
        plot_design_space_coverage(nn_results, fea_data)
        plot_error_distribution(mass_errors, freq_errors)

        # Print validation analysis
        find_closest_fea_validation(nn_results['pareto_designs'], fea_data)
        print_optimization_summary(nn_results, fea_data, mass_errors, freq_errors)
    else:
        # Just plot NN results
        print("\nPlotting NN Pareto front only (no FEA data for comparison)")
        nn_mass = [d['mass'] for d in nn_results['pareto_designs']]
        nn_freq = [d['frequency'] for d in nn_results['pareto_designs']]

        plt.figure(figsize=(10, 6))
        plt.scatter(nn_mass, nn_freq, alpha=0.7, c='red', s=30)
        plt.xlabel('Mass (g)')
        plt.ylabel('Frequency (Hz)')
        plt.title('NN Optimization Pareto Front')
        plt.grid(True, alpha=0.3)
        plt.savefig(project_root / 'nn_pareto_front.png', dpi=150)
        plt.close()
        print(f"Saved: nn_pareto_front.png")

if __name__ == "__main__":
    main()
