"""
Active Learning Calibration Loop

This script implements the iterative calibration workflow:
1. Train initial NN on existing FEA data
2. Run NN optimization to find promising designs
3. Select high-uncertainty designs for FEA validation
4. Run FEA on selected designs (simulated here, needs real FEA integration)
5. Retrain NN with new data
6. Repeat until confidence threshold reached

Usage:
    python run_calibration_loop.py --study uav_arm_optimization --iterations 5

Note: For actual FEA integration, replace the simulate_fea() function with real NX calls.
"""
import sys
from pathlib import Path
import argparse
import json
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))

from optimization_engine.active_learning_surrogate import (
    ActiveLearningSurrogate,
    extract_training_data_from_study
)


def simulate_fea(design: dict, surrogate: ActiveLearningSurrogate) -> dict:
    """
    PLACEHOLDER: Simulate FEA results.

    In production, this would:
    1. Update NX model parameters
    2. Run FEA solve
    3. Extract results (mass, frequency, displacement, stress)

    For now, we use the ensemble mean + noise to simulate "ground truth"
    with systematic differences to test calibration.
    """
    # Get NN prediction
    pred = surrogate.predict(design)

    # Add systematic bias + noise to simulate FEA
    # This simulates the case where NN is systematically off
    fea_mass = pred['mass'] * 0.95 + np.random.normal(0, 50)  # NN overestimates mass by ~5%
    fea_freq = pred['frequency'] * 0.6 + np.random.normal(0, 2)  # NN overestimates freq significantly

    return {
        'mass': max(fea_mass, 1000),  # Ensure positive
        'frequency': max(fea_freq, 1),
        'max_displacement': pred.get('max_displacement', 0),
        'max_stress': pred.get('max_stress', 0)
    }


def run_nn_optimization(
    surrogate: ActiveLearningSurrogate,
    bounds: dict,
    n_trials: int = 500
) -> list:
    """Run NN-only optimization to generate candidate designs."""

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

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

        pred = surrogate.predict(params)

        # Store uncertainty in user_attrs
        trial.set_user_attr('uncertainty', pred['total_uncertainty'])
        trial.set_user_attr('params', params)

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

    study.optimize(objective, n_trials=n_trials, show_progress_bar=False)

    # Extract Pareto front designs with their uncertainty
    pareto_designs = []
    for trial in study.best_trials:
        pareto_designs.append({
            'params': trial.user_attrs['params'],
            'uncertainty': trial.user_attrs['uncertainty'],
            'mass': trial.values[0],
            'frequency': -trial.values[1]
        })

    return pareto_designs


def plot_calibration_progress(history: list, save_path: str):
    """Plot calibration progress over iterations."""
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))

    iterations = [h['iteration'] for h in history]

    # 1. Confidence score
    ax = axes[0, 0]
    confidence = [h['confidence_score'] for h in history]
    ax.plot(iterations, confidence, 'b-o', linewidth=2)
    ax.axhline(y=0.7, color='g', linestyle='--', label='Target (0.7)')
    ax.set_xlabel('Iteration')
    ax.set_ylabel('Confidence Score')
    ax.set_title('Model Confidence Over Iterations')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # 2. MAPE
    ax = axes[0, 1]
    mass_mape = [h['mass_mape'] for h in history]
    freq_mape = [h['freq_mape'] for h in history]
    ax.plot(iterations, mass_mape, 'b-o', label='Mass MAPE')
    ax.plot(iterations, freq_mape, 'r-s', label='Frequency MAPE')
    ax.axhline(y=10, color='g', linestyle='--', label='Target (10%)')
    ax.set_xlabel('Iteration')
    ax.set_ylabel('MAPE (%)')
    ax.set_title('Prediction Error Over Iterations')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # 3. Training samples
    ax = axes[1, 0]
    n_samples = [h['n_training_samples'] for h in history]
    ax.plot(iterations, n_samples, 'g-o', linewidth=2)
    ax.set_xlabel('Iteration')
    ax.set_ylabel('Training Samples')
    ax.set_title('Training Data Growth')
    ax.grid(True, alpha=0.3)

    # 4. Average uncertainty of selected designs
    ax = axes[1, 1]
    avg_uncertainty = [h['avg_selected_uncertainty'] for h in history]
    ax.plot(iterations, avg_uncertainty, 'm-o', linewidth=2)
    ax.set_xlabel('Iteration')
    ax.set_ylabel('Average Uncertainty')
    ax.set_title('Uncertainty of Selected Designs')
    ax.grid(True, alpha=0.3)

    plt.suptitle('Active Learning Calibration Progress', fontsize=14)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    print(f"Saved calibration progress plot: {save_path}")


def main():
    parser = argparse.ArgumentParser(description='Run Active Learning Calibration Loop')
    parser.add_argument('--study', default='uav_arm_optimization', help='Study name')
    parser.add_argument('--iterations', type=int, default=5, help='Number of calibration iterations')
    parser.add_argument('--fea-per-iter', type=int, default=10, help='FEA evaluations per iteration')
    parser.add_argument('--confidence-target', type=float, default=0.7, help='Target confidence')
    parser.add_argument('--simulate', action='store_true', default=True, help='Simulate FEA (for testing)')
    args = parser.parse_args()

    print("="*70)
    print("Active Learning Calibration Loop")
    print("="*70)
    print(f"Study: {args.study}")
    print(f"Max iterations: {args.iterations}")
    print(f"FEA per iteration: {args.fea_per_iter}")
    print(f"Confidence target: {args.confidence_target}")

    # Find database
    db_path = project_root / f"studies/{args.study}/2_results/study.db"
    study_name = args.study

    if not db_path.exists():
        db_path = project_root / "studies/uav_arm_atomizerfield_test/2_results/study.db"
        study_name = "uav_arm_atomizerfield_test"

    if not db_path.exists():
        print(f"ERROR: Database not found: {db_path}")
        return

    # Design bounds (from UAV arm study)
    bounds = {
        'beam_half_core_thickness': (1.0, 10.0),
        'beam_face_thickness': (0.5, 3.0),
        'holes_diameter': (0.5, 50.0),
        'hole_count': (6, 14)
    }

    # Load initial training data
    print(f"\n[1] Loading initial training data from {db_path}")
    design_params, objectives, design_var_names = extract_training_data_from_study(
        str(db_path), study_name
    )
    print(f"    Initial samples: {len(design_params)}")

    # Calibration history
    calibration_history = []

    # Track accumulated training data
    all_design_params = design_params.copy()
    all_objectives = objectives.copy()

    for iteration in range(args.iterations):
        print(f"\n{'='*70}")
        print(f"ITERATION {iteration + 1}/{args.iterations}")
        print("="*70)

        # Train ensemble surrogate
        print(f"\n[2.{iteration+1}] Training ensemble surrogate...")
        surrogate = ActiveLearningSurrogate(n_ensemble=5)
        surrogate.train(
            all_design_params, all_objectives, design_var_names,
            epochs=200
        )

        # Run NN optimization to find candidate designs
        print(f"\n[3.{iteration+1}] Running NN optimization (500 trials)...")
        pareto_designs = run_nn_optimization(surrogate, bounds, n_trials=500)
        print(f"    Found {len(pareto_designs)} Pareto designs")

        # Select designs for FEA validation (highest uncertainty)
        print(f"\n[4.{iteration+1}] Selecting designs for FEA validation...")
        candidate_params = [d['params'] for d in pareto_designs]
        selected = surrogate.select_designs_for_validation(
            candidate_params,
            n_select=args.fea_per_iter,
            strategy='diverse'  # Mix of high uncertainty + diversity
        )

        print(f"    Selected {len(selected)} designs:")
        avg_uncertainty = np.mean([s[2] for s in selected])
        for i, (idx, params, uncertainty) in enumerate(selected[:5]):
            print(f"      {i+1}. Uncertainty={uncertainty:.3f}, params={params}")

        # Run FEA (simulated)
        print(f"\n[5.{iteration+1}] Running FEA validation...")
        new_params = []
        new_objectives = []

        for idx, params, uncertainty in selected:
            if args.simulate:
                fea_result = simulate_fea(params, surrogate)
            else:
                # TODO: Call actual FEA here
                # fea_result = run_actual_fea(params)
                raise NotImplementedError("Real FEA not implemented")

            # Record for retraining
            param_array = [params.get(name, 0.0) for name in design_var_names]
            new_params.append(param_array)
            new_objectives.append([
                fea_result['mass'],
                fea_result['frequency'],
                fea_result.get('max_displacement', 0),
                fea_result.get('max_stress', 0)
            ])

            # Update validation tracking
            surrogate.update_with_validation([params], [fea_result])

        # Add new data to training set
        all_design_params = np.vstack([all_design_params, np.array(new_params, dtype=np.float32)])
        all_objectives = np.vstack([all_objectives, np.array(new_objectives, dtype=np.float32)])

        # Get confidence report
        report = surrogate.get_confidence_report()
        print(f"\n[6.{iteration+1}] Confidence Report:")
        print(f"    Confidence Score: {report['confidence_score']:.3f}")
        print(f"    Mass MAPE: {report['mass_mape']:.1f}%")
        print(f"    Freq MAPE: {report['freq_mape']:.1f}%")
        print(f"    Status: {report['status']}")
        print(f"    Recommendation: {report['recommendation']}")

        # Record history
        calibration_history.append({
            'iteration': iteration + 1,
            'n_training_samples': len(all_design_params),
            'confidence_score': report['confidence_score'],
            'mass_mape': report['mass_mape'],
            'freq_mape': report['freq_mape'],
            'avg_selected_uncertainty': avg_uncertainty,
            'status': report['status']
        })

        # Check if we've reached target confidence
        if report['confidence_score'] >= args.confidence_target:
            print(f"\n*** TARGET CONFIDENCE REACHED ({report['confidence_score']:.3f} >= {args.confidence_target}) ***")
            break

    # Save final model
    print("\n" + "="*70)
    print("CALIBRATION COMPLETE")
    print("="*70)

    model_path = project_root / "calibrated_surrogate.pt"
    surrogate.save(str(model_path))
    print(f"Saved calibrated model to: {model_path}")

    # Save calibration history
    history_path = project_root / "calibration_history.json"
    with open(history_path, 'w') as f:
        json.dump(calibration_history, f, indent=2)
    print(f"Saved calibration history to: {history_path}")

    # Plot progress
    plot_calibration_progress(calibration_history, str(project_root / "calibration_progress.png"))

    # Final summary
    final_report = surrogate.get_confidence_report()
    print(f"\nFinal Results:")
    print(f"  Training samples: {len(all_design_params)}")
    print(f"  Confidence score: {final_report['confidence_score']:.3f}")
    print(f"  Mass MAPE: {final_report['mass_mape']:.1f}%")
    print(f"  Freq MAPE: {final_report['freq_mape']:.1f}%")
    print(f"  Ready for optimization: {surrogate.is_ready_for_optimization()}")


if __name__ == '__main__':
    main()