"""
optimization_interface.py
Bridge between AtomizerField neural network and Atomizer optimization platform

AtomizerField Optimization Interface v2.1
Enables gradient-based optimization with neural field predictions.

Key Features:
- Drop-in replacement for FEA evaluation (1000× faster)
- Gradient computation for sensitivity analysis
- Field-aware optimization (knows WHERE stress occurs)
- Uncertainty quantification (knows when to trust predictions)
- Automatic FEA fallback for high-uncertainty cases
"""

import torch
import torch.nn.functional as F
import numpy as np
from pathlib import Path
import json
import time

from neural_models.field_predictor import AtomizerFieldModel
from neural_models.data_loader import FEAMeshDataset


class NeuralFieldOptimizer:
    """
    Optimization interface for AtomizerField

    This class provides a simple API for optimization:
    - evaluate(parameters) → objectives (max_stress, max_disp, etc.)
    - get_sensitivities(parameters) → gradients for optimization
    - get_fields(parameters) → complete stress/displacement fields

    Usage:
        optimizer = NeuralFieldOptimizer('checkpoint_best.pt')
        results = optimizer.evaluate(parameters)
        print(f"Max stress: {results['max_stress']:.2f} MPa")
    """

    def __init__(
        self,
        model_path,
        uncertainty_threshold=0.1,
        enable_gradients=True,
        device=None
    ):
        """
        Initialize optimizer

        Args:
            model_path (str): Path to trained model checkpoint
            uncertainty_threshold (float): Uncertainty above which to recommend FEA
            enable_gradients (bool): Enable gradient computation
            device (str): Device to run on ('cuda' or 'cpu')
        """
        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)

        print(f"\nAtomizerField Optimization Interface v2.1")
        print(f"Device: {self.device}")

        # Load model
        print(f"Loading model from {model_path}...")
        checkpoint = torch.load(model_path, map_location=self.device)

        # Create model
        model_config = checkpoint['config']['model']
        self.model = AtomizerFieldModel(**model_config)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model = self.model.to(self.device)
        self.model.eval()

        self.config = checkpoint['config']
        self.uncertainty_threshold = uncertainty_threshold
        self.enable_gradients = enable_gradients

        # Model info
        self.model_info = {
            'version': checkpoint.get('epoch', 'unknown'),
            'best_val_loss': checkpoint.get('best_val_loss', 'unknown'),
            'training_config': checkpoint['config']
        }

        print(f"Model loaded successfully!")
        print(f"  Epoch: {checkpoint.get('epoch', 'N/A')}")
        print(f"  Validation loss: {checkpoint.get('best_val_loss', 'N/A')}")

        # Statistics for tracking
        self.eval_count = 0
        self.total_time = 0.0

    def evaluate(self, graph_data, return_fields=False):
        """
        Evaluate design using neural network (drop-in FEA replacement)

        Args:
            graph_data: PyTorch Geometric Data object with mesh graph
            return_fields (bool): Return complete fields or just objectives

        Returns:
            dict: Optimization objectives and optionally complete fields
                - max_stress: Maximum von Mises stress (MPa)
                - max_displacement: Maximum displacement (mm)
                - mass: Total mass (kg) if available
                - fields: Complete stress/displacement fields (if return_fields=True)
                - inference_time_ms: Prediction time
                - uncertainty: Prediction uncertainty (if ensemble enabled)
        """
        start_time = time.time()

        # Move to device
        graph_data = graph_data.to(self.device)

        # Predict
        with torch.set_grad_enabled(self.enable_gradients):
            predictions = self.model(graph_data, return_stress=True)

        inference_time = (time.time() - start_time) * 1000  # ms

        # Extract objectives
        max_displacement = torch.max(
            torch.norm(predictions['displacement'][:, :3], dim=1)
        ).item()

        max_stress = torch.max(predictions['von_mises']).item()

        results = {
            'max_stress': max_stress,
            'max_displacement': max_displacement,
            'inference_time_ms': inference_time,
            'evaluation_count': self.eval_count
        }

        # Add complete fields if requested
        if return_fields:
            results['fields'] = {
                'displacement': predictions['displacement'].cpu().detach().numpy(),
                'stress': predictions['stress'].cpu().detach().numpy(),
                'von_mises': predictions['von_mises'].cpu().detach().numpy()
            }

        # Update statistics
        self.eval_count += 1
        self.total_time += inference_time

        return results

    def get_sensitivities(self, graph_data, objective='max_stress'):
        """
        Compute gradients for gradient-based optimization

        This enables MUCH faster optimization than finite differences!

        Args:
            graph_data: PyTorch Geometric Data with requires_grad=True
            objective (str): Which objective to differentiate ('max_stress' or 'max_displacement')

        Returns:
            dict: Gradients with respect to input features
                - node_gradients: ∂objective/∂node_features
                - edge_gradients: ∂objective/∂edge_features
        """
        if not self.enable_gradients:
            raise RuntimeError("Gradients not enabled. Set enable_gradients=True")

        # Enable gradients
        graph_data = graph_data.to(self.device)
        graph_data.x.requires_grad_(True)
        if graph_data.edge_attr is not None:
            graph_data.edge_attr.requires_grad_(True)

        # Forward pass
        predictions = self.model(graph_data, return_stress=True)

        # Compute objective
        if objective == 'max_stress':
            obj = torch.max(predictions['von_mises'])
        elif objective == 'max_displacement':
            disp_mag = torch.norm(predictions['displacement'][:, :3], dim=1)
            obj = torch.max(disp_mag)
        else:
            raise ValueError(f"Unknown objective: {objective}")

        # Backward pass
        obj.backward()

        # Extract gradients
        gradients = {
            'node_gradients': graph_data.x.grad.cpu().numpy(),
            'objective_value': obj.item()
        }

        if graph_data.edge_attr is not None and graph_data.edge_attr.grad is not None:
            gradients['edge_gradients'] = graph_data.edge_attr.grad.cpu().numpy()

        return gradients

    def batch_evaluate(self, graph_data_list, return_fields=False):
        """
        Evaluate multiple designs in batch (even faster!)

        Args:
            graph_data_list (list): List of graph data objects
            return_fields (bool): Return complete fields

        Returns:
            list: List of evaluation results
        """
        results = []

        for graph_data in graph_data_list:
            result = self.evaluate(graph_data, return_fields=return_fields)
            results.append(result)

        return results

    def needs_fea_validation(self, uncertainty):
        """
        Determine if FEA validation is recommended

        Args:
            uncertainty (float): Prediction uncertainty

        Returns:
            bool: True if FEA is recommended
        """
        return uncertainty > self.uncertainty_threshold

    def compare_with_fea(self, graph_data, fea_results):
        """
        Compare neural predictions with FEA ground truth

        Args:
            graph_data: Mesh graph
            fea_results (dict): FEA results with 'max_stress', 'max_displacement'

        Returns:
            dict: Comparison metrics
        """
        # Neural prediction
        pred = self.evaluate(graph_data)

        # Compute errors
        stress_error = abs(pred['max_stress'] - fea_results['max_stress'])
        stress_rel_error = stress_error / (fea_results['max_stress'] + 1e-8)

        disp_error = abs(pred['max_displacement'] - fea_results['max_displacement'])
        disp_rel_error = disp_error / (fea_results['max_displacement'] + 1e-8)

        comparison = {
            'neural_prediction': pred,
            'fea_results': fea_results,
            'errors': {
                'stress_error_abs': stress_error,
                'stress_error_rel': stress_rel_error,
                'displacement_error_abs': disp_error,
                'displacement_error_rel': disp_rel_error
            },
            'within_tolerance': stress_rel_error < 0.1 and disp_rel_error < 0.1
        }

        return comparison

    def get_statistics(self):
        """
        Get optimizer usage statistics

        Returns:
            dict: Statistics about predictions
        """
        avg_time = self.total_time / self.eval_count if self.eval_count > 0 else 0

        return {
            'total_evaluations': self.eval_count,
            'total_time_ms': self.total_time,
            'average_time_ms': avg_time,
            'model_info': self.model_info
        }

    def reset_statistics(self):
        """Reset usage statistics"""
        self.eval_count = 0
        self.total_time = 0.0


class ParametricOptimizer:
    """
    Optimizer for parametric designs

    This wraps NeuralFieldOptimizer and adds parameter → mesh conversion.
    Enables direct optimization over design parameters (thickness, radius, etc.)
    """

    def __init__(
        self,
        model_path,
        parameter_names,
        parameter_bounds,
        mesh_generator_fn
    ):
        """
        Initialize parametric optimizer

        Args:
            model_path (str): Path to trained model
            parameter_names (list): Names of design parameters
            parameter_bounds (dict): Bounds for each parameter
            mesh_generator_fn: Function that converts parameters → graph_data
        """
        self.neural_optimizer = NeuralFieldOptimizer(model_path)
        self.parameter_names = parameter_names
        self.parameter_bounds = parameter_bounds
        self.mesh_generator = mesh_generator_fn

        print(f"\nParametric Optimizer initialized")
        print(f"Design parameters: {parameter_names}")

    def evaluate_parameters(self, parameters):
        """
        Evaluate design from parameters

        Args:
            parameters (dict): Design parameters

        Returns:
            dict: Objectives (max_stress, max_displacement, etc.)
        """
        # Generate mesh from parameters
        graph_data = self.mesh_generator(parameters)

        # Evaluate
        results = self.neural_optimizer.evaluate(graph_data)

        # Add parameters to results
        results['parameters'] = parameters

        return results

    def optimize(
        self,
        initial_parameters,
        objectives,
        constraints,
        method='gradient',
        max_iterations=100
    ):
        """
        Run optimization

        Args:
            initial_parameters (dict): Starting point
            objectives (list): Objectives to minimize/maximize
            constraints (list): Constraint functions
            method (str): Optimization method ('gradient' or 'genetic')
            max_iterations (int): Maximum iterations

        Returns:
            dict: Optimal parameters and results
        """
        # This would integrate with scipy.optimize or genetic algorithms
        # Placeholder for now

        print(f"\nStarting optimization with {method} method...")
        print(f"Initial parameters: {initial_parameters}")
        print(f"Objectives: {objectives}")
        print(f"Max iterations: {max_iterations}")

        # TODO: Implement optimization loop
        # For gradient-based:
        #   1. Evaluate at current parameters
        #   2. Compute sensitivities
        #   3. Update parameters using gradients
        #   4. Repeat until convergence

        raise NotImplementedError("Full optimization loop coming in next update!")


def create_optimizer(model_path, config=None):
    """
    Factory function to create optimizer

    Args:
        model_path (str): Path to trained model
        config (dict): Optimizer configuration

    Returns:
        NeuralFieldOptimizer instance
    """
    if config is None:
        config = {}

    return NeuralFieldOptimizer(model_path, **config)


if __name__ == "__main__":
    # Example usage
    print("AtomizerField Optimization Interface")
    print("=" * 60)
    print("\nThis module provides fast optimization with neural field predictions.")
    print("\nExample usage:")
    print("""
    # Create optimizer
    optimizer = NeuralFieldOptimizer('checkpoint_best.pt')

    # Evaluate design
    results = optimizer.evaluate(graph_data)
    print(f"Max stress: {results['max_stress']:.2f} MPa")
    print(f"Inference time: {results['inference_time_ms']:.1f} ms")

    # Get sensitivities for gradient-based optimization
    gradients = optimizer.get_sensitivities(graph_data, objective='max_stress')

    # Batch evaluation (test 1000 designs in seconds!)
    all_results = optimizer.batch_evaluate(design_variants)
    """)

    print("\nOptimization interface ready!")
