"""
Example: Fan Design Optimization using the Surrogate Model

This script demonstrates how to use the surrogate model for
rapid design space exploration and optimization.
"""

import numpy as np
from model import FanDesignSurrogate


def design_optimization_example():
    """Find the fan design that maximizes efficiency at a target pressure rise."""
    
    # Load surrogate
    surrogate = FanDesignSurrogate.from_pretrained("harshaperla/fan-design-surrogate")
    
    # Baseline design
    baseline = {
        'blade_inlet_angle_deg': 55.0,
        'blade_turning_angle_deg': 15.0,
        'chord_length_mm': 100.0,
        'blade_thickness_ratio': 0.06,
        'stagger_angle_deg': 45.0,
        'hub_tip_ratio': 0.5,
        'tip_clearance_ratio': 0.015,
        'num_blades': 12,
        'aspect_ratio': 2.5,
        'solidity': 1.0,
        'sweep_angle_deg': 0.0,
        'flow_coefficient': 0.5,
        'rotational_speed_rpm': 3000,
        'tip_radius_mm': 300.0,
    }
    
    print("Baseline Design Performance:")
    baseline_perf = surrogate.predict(baseline)
    for k, v in baseline_perf.items():
        print(f"  {k}: {v:.4f}")
    
    # =========================================================================
    # 1. Sensitivity Analysis
    # =========================================================================
    print("\n" + "=" * 50)
    print("Sensitivity Analysis (efficiency vs each parameter)")
    print("=" * 50)
    
    for param in ['blade_turning_angle_deg', 'solidity', 'tip_clearance_ratio', 
                  'flow_coefficient', 'hub_tip_ratio']:
        values, predictions = surrogate.sensitivity_analysis(baseline, param)
        etas = [p['isentropic_efficiency'] for p in predictions]
        best_idx = np.argmax(etas)
        print(f"  {param}: best η={etas[best_idx]:.4f} at {param}={values[best_idx]:.3f}")
    
    # =========================================================================
    # 2. Random Search Optimization
    # =========================================================================
    print("\n" + "=" * 50)
    print("Random Search Optimization (maximize η, target ΔPt > 1500 Pa)")
    print("=" * 50)
    
    bounds = surrogate.bounds
    n_samples = 50000
    
    # Generate random designs
    designs = []
    for _ in range(n_samples):
        d = {}
        for col in surrogate.input_cols:
            lo, hi = bounds[col]
            d[col] = np.random.uniform(lo, hi)
        designs.append(d)
    
    # Batch prediction (fast!)
    predictions = surrogate.predict_batch(designs)
    
    # Filter: pressure rise > 1500 Pa, noise < 100 dBA
    best_eta = 0
    best_design = None
    best_perf = None
    
    for d, p in zip(designs, predictions):
        if (p['total_pressure_rise_Pa'] > 1500 and 
            p['noise_estimate_dBA'] < 100 and
            p['isentropic_efficiency'] > best_eta):
            best_eta = p['isentropic_efficiency']
            best_design = d
            best_perf = p
    
    if best_design:
        print(f"\n  Best design (η = {best_eta:.4f}):")
        for k, v in best_design.items():
            print(f"    {k}: {v:.3f}")
        print(f"\n  Performance:")
        for k, v in best_perf.items():
            print(f"    {k}: {v:.4f}")
    
    # =========================================================================
    # 3. Pareto Front (Efficiency vs Noise)
    # =========================================================================
    print("\n" + "=" * 50)
    print("Pareto Front: Efficiency vs Noise")
    print("=" * 50)
    
    # Filter feasible designs
    feasible = [(p['isentropic_efficiency'], p['noise_estimate_dBA']) 
                for p in predictions if p['total_pressure_rise_Pa'] > 1000]
    
    if feasible:
        etas = [f[0] for f in feasible]
        noises = [f[1] for f in feasible]
        
        # Simple Pareto extraction
        pareto = []
        sorted_by_eta = sorted(zip(etas, noises), key=lambda x: -x[0])
        min_noise = float('inf')
        for eta, noise in sorted_by_eta:
            if noise < min_noise:
                pareto.append((eta, noise))
                min_noise = noise
        
        print(f"  Feasible designs: {len(feasible)}/{n_samples}")
        print(f"  Pareto front points: {len(pareto)}")
        print(f"\n  {'η':>8} {'Noise (dBA)':>12}")
        print(f"  {'-'*8} {'-'*12}")
        for eta, noise in pareto[:10]:
            print(f"  {eta:>8.4f} {noise:>11.1f}")


if __name__ == '__main__':
    design_optimization_example()