"""
Production-Ready NN Optimization with Confidence Bounds

This script runs multi-objective optimization using the CV-validated neural network
with proper extrapolation warnings and confidence-bounded results.

Key Features:
1. Uses CV-validated model with known accuracy (1.8% mass, 1.1% freq MAPE)
2. Warns when extrapolating outside training data range
3. Reads optimization bounds from study's optimization_config.json
4. Constrains optimization to prescribed bounds for reliable predictions
5. Marks designs needing FEA validation
"""
import sys
from pathlib import Path
import time
import json
import argparse
import torch
import numpy as np
import optuna
from optuna.samplers import NSGAIISampler
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

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


def load_config_bounds(study_path: Path) -> dict:
    """Load design variable bounds from optimization_config.json

    Returns dict: {param_name: (min, max, is_int)}
    """
    config_path = study_path / "1_setup" / "optimization_config.json"

    if not config_path.exists():
        raise FileNotFoundError(f"Config not found: {config_path}")

    with open(config_path) as f:
        config = json.load(f)

    bounds = {}
    for var in config.get('design_variables', []):
        # Support both 'parameter' and 'name' keys
        name = var.get('parameter') or var.get('name')

        # Support both "bounds": [min, max] and "min_value"/"max_value" formats
        if 'bounds' in var:
            min_val, max_val = var['bounds']
        else:
            min_val = var.get('min_value', var.get('min', 0))
            max_val = var.get('max_value', var.get('max', 1))

        # Detect integer type based on name or explicit type
        is_int = (var.get('type') == 'integer' or
                  'count' in name.lower() or
                  (isinstance(min_val, int) and isinstance(max_val, int) and max_val - min_val < 20))

        bounds[name] = (min_val, max_val, is_int)

    return bounds

from optimization_engine.processors.surrogates.active_learning_surrogate import EnsembleMLP


class ValidatedSurrogate:
    """Surrogate with CV-validated accuracy and extrapolation detection."""

    def __init__(self, model_path: str):
        state = torch.load(model_path, map_location='cpu')

        self.model = EnsembleMLP(
            input_dim=len(state['design_var_names']),
            output_dim=4,  # mass, freq, disp, stress
            hidden_dims=state['hidden_dims']
        )
        self.model.load_state_dict(state['model'])
        self.model.eval()

        self.input_mean = np.array(state['input_mean'])
        self.input_std = np.array(state['input_std'])
        self.output_mean = np.array(state['output_mean'])
        self.output_std = np.array(state['output_std'])
        self.design_var_names = state['design_var_names']

        # CV metrics
        self.cv_mass_mape = state['cv_mass_mape']
        self.cv_freq_mape = state['cv_freq_mape']
        self.cv_mass_std = state['cv_mass_std']
        self.cv_freq_std = state['cv_freq_std']
        self.n_training_samples = state['n_samples']

        # Training bounds (for extrapolation detection)
        self.bounds_min = self.input_mean - 2 * self.input_std
        self.bounds_max = self.input_mean + 2 * self.input_std

    def predict(self, params: dict) -> dict:
        """Predict with extrapolation check."""
        x = np.array([[params.get(name, 0.0) for name in self.design_var_names]], dtype=np.float32)

        # Check for extrapolation
        extrapolation_score = 0.0
        for i, name in enumerate(self.design_var_names):
            val = x[0, i]
            if val < self.bounds_min[i]:
                extrapolation_score += (self.bounds_min[i] - val) / (self.input_std[i] + 1e-8)
            elif val > self.bounds_max[i]:
                extrapolation_score += (val - self.bounds_max[i]) / (self.input_std[i] + 1e-8)

        # Normalize input
        x_norm = (x - self.input_mean) / (self.input_std + 1e-8)
        x_t = torch.FloatTensor(x_norm)

        # Predict
        with torch.no_grad():
            pred_norm = self.model(x_t).numpy()

        pred = pred_norm * (self.output_std + 1e-8) + self.output_mean

        # Calculate confidence-adjusted uncertainty
        base_uncertainty = self.cv_mass_mape / 100  # Base uncertainty from CV
        extrapolation_penalty = min(extrapolation_score * 0.1, 0.5)  # Max 50% extra uncertainty
        total_uncertainty = base_uncertainty + extrapolation_penalty

        return {
            'mass': float(pred[0, 0]),
            'frequency': float(pred[0, 1]),
            'max_displacement': float(pred[0, 2]),
            'max_stress': float(pred[0, 3]),
            'uncertainty': total_uncertainty,
            'extrapolating': extrapolation_score > 0.1,
            'extrapolation_score': extrapolation_score,
            'needs_fea_validation': extrapolation_score > 0.5
        }


def run_optimization(surrogate: ValidatedSurrogate, bounds: dict, n_trials: int = 1000):
    """Run multi-objective optimization.

    Args:
        surrogate: ValidatedSurrogate model
        bounds: Dict from load_config_bounds {name: (min, max, is_int)}
        n_trials: Number of optimization trials
    """
    print(f"\nOptimization bounds (from config):")
    for name, (low, high, is_int) in bounds.items():
        type_str = "int" if is_int else "float"
        print(f"  {name}: [{low}, {high}] ({type_str})")

    # Create study
    study = optuna.create_study(
        directions=["minimize", "minimize"],  # mass, -frequency
        sampler=NSGAIISampler()
    )

    # Track stats
    start_time = time.time()
    trial_times = []
    extrapolation_count = 0

    def objective(trial: optuna.Trial):
        nonlocal extrapolation_count

        params = {}
        for name, (low, high, is_int) in bounds.items():
            if is_int:
                params[name] = trial.suggest_int(name, int(low), int(high))
            else:
                params[name] = trial.suggest_float(name, float(low), float(high))

        trial_start = time.time()
        result = surrogate.predict(params)
        trial_time = (time.time() - trial_start) * 1000
        trial_times.append(trial_time)

        # Track extrapolation
        if result['extrapolating']:
            extrapolation_count += 1

        # Store metadata
        trial.set_user_attr('uncertainty', result['uncertainty'])
        trial.set_user_attr('extrapolating', result['extrapolating'])
        trial.set_user_attr('needs_fea', result['needs_fea_validation'])
        trial.set_user_attr('params', params)

        if trial.number % 200 == 0:
            print(f"  Trial {trial.number}: mass={result['mass']:.1f}g, freq={result['frequency']:.2f}Hz, extrap={result['extrapolating']}")

        return result['mass'], -result['frequency']

    print(f"\nRunning {n_trials} trials...")
    study.optimize(objective, n_trials=n_trials, show_progress_bar=True)

    total_time = time.time() - start_time

    return study, {
        'total_time': total_time,
        'avg_trial_time_ms': np.mean(trial_times),
        'trials_per_second': n_trials / total_time,
        'extrapolation_count': extrapolation_count,
        'extrapolation_pct': extrapolation_count / n_trials * 100
    }


def analyze_pareto_front(study, surrogate):
    """Analyze Pareto front with confidence information."""
    pareto_trials = study.best_trials

    results = {
        'total_pareto_designs': len(pareto_trials),
        'confident_designs': 0,
        'needs_fea_designs': 0,
        'designs': []
    }

    for trial in pareto_trials:
        mass = trial.values[0]
        freq = -trial.values[1]
        uncertainty = trial.user_attrs.get('uncertainty', 0)
        needs_fea = trial.user_attrs.get('needs_fea', False)
        params = trial.user_attrs.get('params', trial.params)

        if needs_fea:
            results['needs_fea_designs'] += 1
        else:
            results['confident_designs'] += 1

        results['designs'].append({
            'mass': mass,
            'frequency': freq,
            'uncertainty': uncertainty,
            'needs_fea': needs_fea,
            'params': params
        })

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

    return results


def plot_results(study, pareto_analysis, surrogate, save_path):
    """Generate visualization of optimization results."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))

    # 1. Pareto Front with Confidence
    ax = axes[0, 0]
    pareto = pareto_analysis['designs']

    confident = [d for d in pareto if not d['needs_fea']]
    needs_fea = [d for d in pareto if d['needs_fea']]

    if confident:
        ax.scatter([d['mass'] for d in confident],
                   [d['frequency'] for d in confident],
                   c='green', s=50, alpha=0.7, label=f'Confident ({len(confident)})')

    if needs_fea:
        ax.scatter([d['mass'] for d in needs_fea],
                   [d['frequency'] for d in needs_fea],
                   c='orange', s=50, alpha=0.7, marker='^', label=f'Needs FEA ({len(needs_fea)})')

    ax.set_xlabel('Mass (g)')
    ax.set_ylabel('Frequency (Hz)')
    ax.set_title('Pareto Front with Confidence Assessment')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # 2. Uncertainty vs Performance
    ax = axes[0, 1]
    all_trials = study.trials
    masses = [t.values[0] for t in all_trials if t.values]
    freqs = [-t.values[1] for t in all_trials if t.values]
    uncertainties = [t.user_attrs.get('uncertainty', 0) for t in all_trials if t.values]

    sc = ax.scatter(masses, freqs, c=uncertainties, cmap='RdYlGn_r', s=10, alpha=0.5)
    plt.colorbar(sc, ax=ax, label='Uncertainty')
    ax.set_xlabel('Mass (g)')
    ax.set_ylabel('Frequency (Hz)')
    ax.set_title('All Trials Colored by Uncertainty')
    ax.grid(True, alpha=0.3)

    # 3. Top 10 Pareto Designs
    ax = axes[1, 0]
    top_10 = pareto_analysis['designs'][:10]

    y_pos = np.arange(len(top_10))
    bars = ax.barh(y_pos, [d['mass'] for d in top_10],
                   color=['green' if not d['needs_fea'] else 'orange' for d in top_10])

    ax.set_yticks(y_pos)
    ax.set_yticklabels([f"#{i+1}: {d['frequency']:.1f}Hz" for i, d in enumerate(top_10)])
    ax.set_xlabel('Mass (g)')
    ax.set_title('Top 10 Lowest Mass Designs')
    ax.grid(True, alpha=0.3, axis='x')

    # Add confidence text
    for i, (bar, d) in enumerate(zip(bars, top_10)):
        status = "FEA!" if d['needs_fea'] else "OK"
        ax.text(bar.get_width() + 10, bar.get_y() + bar.get_height()/2,
                status, va='center', fontsize=9)

    # 4. Summary Text
    ax = axes[1, 1]
    ax.axis('off')

    summary_text = f"""
OPTIMIZATION SUMMARY
====================

CV-Validated Model Accuracy:
  Mass MAPE: {surrogate.cv_mass_mape:.1f}% +/- {surrogate.cv_mass_std:.1f}%
  Freq MAPE: {surrogate.cv_freq_mape:.1f}% +/- {surrogate.cv_freq_std:.1f}%
  Training samples: {surrogate.n_training_samples}

Pareto Front:
  Total designs: {pareto_analysis['total_pareto_designs']}
  High confidence: {pareto_analysis['confident_designs']}
  Needs FEA validation: {pareto_analysis['needs_fea_designs']}

Best Confident Design:
"""
    # Find best confident design
    best_confident = [d for d in pareto_analysis['designs'] if not d['needs_fea']]
    if best_confident:
        best = best_confident[0]
        summary_text += f"""
  Mass: {best['mass']:.1f}g
  Frequency: {best['frequency']:.2f}Hz
  Parameters:
"""
        for k, v in best['params'].items():
            if isinstance(v, float):
                summary_text += f"    {k}: {v:.3f}\n"
            else:
                summary_text += f"    {k}: {v}\n"
    else:
        summary_text += "  No confident designs found - run more FEA!"

    ax.text(0.05, 0.95, summary_text, transform=ax.transAxes,
            fontfamily='monospace', fontsize=10, verticalalignment='top')

    plt.suptitle('NN-Based Multi-Objective Optimization Results', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    print(f"Saved plot: {save_path}")


def main():
    parser = argparse.ArgumentParser(description='NN Optimization with Confidence Bounds')
    parser.add_argument('--study', default='uav_arm_optimization',
                        help='Study name (e.g., uav_arm_optimization)')
    parser.add_argument('--model', default=None,
                        help='Path to CV-validated model (default: cv_validated_surrogate.pt)')
    parser.add_argument('--trials', type=int, default=2000,
                        help='Number of optimization trials')
    args = parser.parse_args()

    print("="*70)
    print("Production-Ready NN Optimization with Confidence Bounds")
    print("="*70)

    # Load bounds from study config
    study_path = project_root / "studies" / args.study
    if not study_path.exists():
        print(f"ERROR: Study not found: {study_path}")
        return

    print(f"\nLoading bounds from study: {args.study}")
    bounds = load_config_bounds(study_path)
    print(f"  Loaded {len(bounds)} design variables from config")

    # Load CV-validated model
    model_path = Path(args.model) if args.model else project_root / "cv_validated_surrogate.pt"
    if not model_path.exists():
        print(f"ERROR: Run validate_surrogate_real_data.py first to create {model_path}")
        return

    print(f"\nLoading CV-validated model: {model_path}")
    surrogate = ValidatedSurrogate(str(model_path))

    print(f"\nModel Statistics:")
    print(f"  Training samples: {surrogate.n_training_samples}")
    print(f"  CV Mass MAPE: {surrogate.cv_mass_mape:.1f}% +/- {surrogate.cv_mass_std:.1f}%")
    print(f"  CV Freq MAPE: {surrogate.cv_freq_mape:.1f}% +/- {surrogate.cv_freq_std:.1f}%")
    print(f"  Design variables: {surrogate.design_var_names}")

    # Verify model variables match config variables
    config_vars = set(bounds.keys())
    model_vars = set(surrogate.design_var_names)
    if config_vars != model_vars:
        print(f"\nWARNING: Variable mismatch!")
        print(f"  Config has: {config_vars}")
        print(f"  Model has: {model_vars}")
        print(f"  Missing from model: {config_vars - model_vars}")
        print(f"  Extra in model: {model_vars - config_vars}")

    # Run optimization
    print("\n" + "="*70)
    print("Running Multi-Objective Optimization")
    print("="*70)

    study, stats = run_optimization(surrogate, bounds, n_trials=args.trials)

    print(f"\nOptimization Statistics:")
    print(f"  Total time: {stats['total_time']:.1f}s")
    print(f"  Avg trial time: {stats['avg_trial_time_ms']:.2f}ms")
    print(f"  Trials per second: {stats['trials_per_second']:.1f}")
    print(f"  Extrapolating trials: {stats['extrapolation_count']} ({stats['extrapolation_pct']:.1f}%)")

    # Analyze Pareto front
    print("\n" + "="*70)
    print("Pareto Front Analysis")
    print("="*70)

    pareto_analysis = analyze_pareto_front(study, surrogate)

    print(f"\nPareto Front Summary:")
    print(f"  Total Pareto-optimal designs: {pareto_analysis['total_pareto_designs']}")
    print(f"  High confidence designs: {pareto_analysis['confident_designs']}")
    print(f"  Needs FEA validation: {pareto_analysis['needs_fea_designs']}")

    print(f"\nTop 5 Lowest Mass Designs:")
    for i, d in enumerate(pareto_analysis['designs'][:5]):
        status = "[NEEDS FEA]" if d['needs_fea'] else "[OK]"
        print(f"  {i+1}. Mass={d['mass']:.1f}g, Freq={d['frequency']:.2f}Hz {status}")
        print(f"      Params: {d['params']}")

    # Generate plots
    plot_path = project_root / "validated_nn_optimization_results.png"
    plot_results(study, pareto_analysis, surrogate, str(plot_path))

    # Save results
    results_path = project_root / "validated_nn_optimization_results.json"
    with open(results_path, 'w') as f:
        json.dump({
            'stats': stats,
            'pareto_summary': {
                'total': pareto_analysis['total_pareto_designs'],
                'confident': pareto_analysis['confident_designs'],
                'needs_fea': pareto_analysis['needs_fea_designs']
            },
            'top_designs': pareto_analysis['designs'][:20],
            'cv_metrics': {
                'mass_mape': surrogate.cv_mass_mape,
                'freq_mape': surrogate.cv_freq_mape
            }
        }, f, indent=2)
    print(f"\nSaved results: {results_path}")

    print("\n" + "="*70)
    print("NEXT STEPS")
    print("="*70)

    if pareto_analysis['confident_designs'] > 0:
        print("1. Review the confident Pareto-optimal designs above")
        print("2. Consider validating top 3-5 designs with actual FEA")
        print("3. If FEA matches predictions, use for manufacturing")
    else:
        print("1. All Pareto designs are extrapolating - need more FEA data!")
        print("2. Run FEA on designs marked [NEEDS FEA]")
        print("3. Retrain the model with new data")


if __name__ == '__main__':
    main()