//! Inter-process communication types for PINN benchmark

use serde::{Deserialize, Serialize};

use crate::config::{BenchmarkConfig, ProblemType};

/// Request messages from frontend to backend
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum PINNBenchmarkRequest {
    /// Initialize benchmark with configuration
    Initialize {
        /// Benchmark configuration
        config: BenchmarkConfig,
    },

    /// Start training
    StartTraining,

    /// Stop training
    StopTraining,

    /// Get current training status
    GetStatus,

    /// Run inference on test points
    RunInference {
        /// Test points for inference
        test_points: Vec<Vec<f64>>,
    },

    /// Get benchmark results
    GetResults,

    /// Reset benchmark state
    Reset,
}

/// Response messages from backend to frontend
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum PINNBenchmarkResponse {
    /// Initialization successful
    Initialized,

    /// Training started
    TrainingStarted,

    /// Training progress update
    TrainingProgress {
        /// Training progress data
        progress: TrainingProgress,
    },

    /// Training completed
    TrainingCompleted {
        /// Benchmark result
        result: BenchmarkResult,
    },

    /// Training stopped by user
    TrainingStopped,

    /// Inference results
    InferenceResults {
        /// Prediction values
        predictions: Vec<f64>,
    },

    /// Benchmark results
    Results {
        /// Comparison result
        comparison: ComparisonResult,
    },

    /// Error occurred
    Error {
        /// Error message
        message: String,
    },
}

/// Training progress information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current epoch
    pub epoch: usize,

    /// Total loss (combined)
    pub total_loss: f64,

    /// Physics loss (PDE residual)
    pub physics_loss: f64,

    /// Boundary condition loss
    pub boundary_loss: f64,

    /// Initial condition loss
    pub initial_loss: f64,

    /// Current learning rate
    pub learning_rate: f64,

    /// Elapsed time in milliseconds
    pub elapsed_ms: u64,
}

impl TrainingProgress {
    /// Creates a new training progress instance
    #[must_use]
    pub const fn new(
        epoch: usize,
        total_loss: f64,
        physics_loss: f64,
        boundary_loss: f64,
        initial_loss: f64,
        learning_rate: f64,
        elapsed_ms: u64,
    ) -> Self {
        Self {
            epoch,
            total_loss,
            physics_loss,
            boundary_loss,
            initial_loss,
            learning_rate,
            elapsed_ms,
        }
    }
}

/// Benchmark result for a single problem
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
    /// Problem type
    pub problem: ProblemType,

    /// Total training time in seconds
    pub total_training_time_s: f64,

    /// Final loss value
    pub final_loss: f64,

    /// L2 error vs analytical solution
    pub l2_error: f64,

    /// L-infinity error (max absolute error)
    pub linf_error: f64,

    /// Peak memory usage in MB
    pub memory_mb: f64,

    /// Training throughput (samples per second)
    pub throughput_samples_per_sec: f64,

    /// Convergence history (loss per epoch)
    pub convergence_history: Vec<f64>,
}

impl BenchmarkResult {
    /// Creates a new benchmark result
    #[must_use]
    pub fn new(
        problem: ProblemType,
        total_training_time_s: f64,
        final_loss: f64,
        l2_error: f64,
        linf_error: f64,
        memory_mb: f64,
        throughput_samples_per_sec: f64,
        convergence_history: Vec<f64>,
    ) -> Self {
        Self {
            problem,
            total_training_time_s,
            final_loss,
            l2_error,
            linf_error,
            memory_mb,
            throughput_samples_per_sec,
            convergence_history,
        }
    }
}

/// Accuracy metrics comparing PINN to reference solution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccuracyMetrics {
    /// Mean absolute error
    pub mae: f64,

    /// Root mean square error
    pub rmse: f64,

    /// Relative L2 error
    pub relative_l2: f64,

    /// Maximum absolute error
    pub max_error: f64,

    /// R-squared coefficient
    pub r_squared: f64,
}

impl AccuracyMetrics {
    /// Creates new accuracy metrics
    #[must_use]
    pub const fn new(
        mae: f64,
        rmse: f64,
        relative_l2: f64,
        max_error: f64,
        r_squared: f64,
    ) -> Self {
        Self {
            mae,
            rmse,
            relative_l2,
            max_error,
            r_squared,
        }
    }

    /// Computes accuracy metrics from predictions and reference values
    #[must_use]
    pub fn compute(predictions: &[f64], reference: &[f64]) -> Self {
        assert_eq!(
            predictions.len(),
            reference.len(),
            "Arrays must have same length"
        );

        let n = predictions.len() as f64;
        let mut sum_abs_error: f64 = 0.0;
        let mut sum_sq_error: f64 = 0.0;
        let mut max_error: f64 = 0.0;

        for (pred, ref_val) in predictions.iter().zip(reference.iter()) {
            let error = (pred - ref_val).abs();
            sum_abs_error += error;
            sum_sq_error += error * error;
            max_error = max_error.max(error);
        }

        let mae = sum_abs_error / n;
        let rmse = (sum_sq_error / n).sqrt();

        // Compute relative L2 error
        let l2_error = sum_sq_error.sqrt();
        let l2_reference = reference.iter().map(|x| x * x).sum::<f64>().sqrt();
        let relative_l2 = if l2_reference > 1e-10 {
            l2_error / l2_reference
        } else {
            l2_error
        };

        // Compute R-squared
        let mean_ref = reference.iter().sum::<f64>() / n;
        let ss_tot = reference
            .iter()
            .map(|x| (x - mean_ref).powi(2))
            .sum::<f64>();
        let ss_res = sum_sq_error;
        let r_squared = if ss_tot > 1e-10 {
            1.0 - (ss_res / ss_tot)
        } else {
            0.0
        };

        Self {
            mae,
            rmse,
            relative_l2,
            max_error,
            r_squared,
        }
    }
}

/// Comparison result between PINN and traditional solver
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonResult {
    /// PINN benchmark result
    pub pinn_result: BenchmarkResult,

    /// Time taken by traditional solver (FDM/FEM) in seconds
    pub reference_time_s: f64,

    /// Speedup factor for inference (PINN vs traditional)
    pub speedup_inference: f64,

    /// Accuracy comparison metrics
    pub accuracy_comparison: AccuracyMetrics,
}

impl ComparisonResult {
    /// Creates a new comparison result
    #[must_use]
    pub const fn new(
        pinn_result: BenchmarkResult,
        reference_time_s: f64,
        speedup_inference: f64,
        accuracy_comparison: AccuracyMetrics,
    ) -> Self {
        Self {
            pinn_result,
            reference_time_s,
            speedup_inference,
            accuracy_comparison,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::ProblemType;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_training_progress_creation() {
        let progress = TrainingProgress::new(100, 0.5, 0.3, 0.1, 0.1, 0.001, 5000);
        assert_eq!(progress.epoch, 100);
        assert_abs_diff_eq!(progress.total_loss, 0.5);
        assert_abs_diff_eq!(progress.learning_rate, 0.001);
    }

    #[test]
    fn test_benchmark_result_creation() {
        let result = BenchmarkResult::new(
            ProblemType::Burgers1D,
            10.5,
            0.001,
            0.01,
            0.05,
            256.0,
            1000.0,
            vec![1.0, 0.5, 0.1, 0.01],
        );

        assert_eq!(result.problem, ProblemType::Burgers1D);
        assert_abs_diff_eq!(result.total_training_time_s, 10.5);
        assert_eq!(result.convergence_history.len(), 4);
    }

    #[test]
    fn test_accuracy_metrics_compute_perfect() {
        let predictions = vec![1.0, 2.0, 3.0, 4.0];
        let reference = vec![1.0, 2.0, 3.0, 4.0];

        let metrics = AccuracyMetrics::compute(&predictions, &reference);

        assert_abs_diff_eq!(metrics.mae, 0.0, epsilon = 1e-10);
        assert_abs_diff_eq!(metrics.rmse, 0.0, epsilon = 1e-10);
        assert_abs_diff_eq!(metrics.max_error, 0.0, epsilon = 1e-10);
        assert!(metrics.r_squared >= 0.99); // Should be close to 1.0
    }

    #[test]
    fn test_accuracy_metrics_compute_with_errors() {
        let predictions = vec![1.1, 2.2, 2.9, 4.1];
        let reference = vec![1.0, 2.0, 3.0, 4.0];

        let metrics = AccuracyMetrics::compute(&predictions, &reference);

        assert!(metrics.mae > 0.0);
        assert!(metrics.rmse > 0.0);
        assert!(metrics.max_error >= metrics.mae);
        assert!(metrics.r_squared < 1.0);
        assert!(metrics.r_squared > 0.0);
    }

    #[test]
    fn test_accuracy_metrics_mae_calculation() {
        let predictions = vec![1.0, 2.0, 3.0];
        let reference = vec![1.1, 1.9, 3.2];

        let metrics = AccuracyMetrics::compute(&predictions, &reference);

        // MAE = (0.1 + 0.1 + 0.2) / 3 = 0.4 / 3 ≈ 0.1333
        assert_abs_diff_eq!(metrics.mae, 0.1333, epsilon = 0.001);
    }

    #[test]
    fn test_accuracy_metrics_rmse_calculation() {
        let predictions = vec![1.0, 2.0];
        let reference = vec![1.0, 3.0];

        let metrics = AccuracyMetrics::compute(&predictions, &reference);

        // RMSE = sqrt((0^2 + 1^2) / 2) = sqrt(0.5) ≈ 0.707
        assert_abs_diff_eq!(metrics.rmse, 0.707, epsilon = 0.001);
    }

    #[test]
    fn test_accuracy_metrics_max_error() {
        let predictions = vec![1.0, 2.0, 3.0, 4.0];
        let reference = vec![1.1, 2.05, 2.7, 4.0];

        let metrics = AccuracyMetrics::compute(&predictions, &reference);

        // Max error should be 0.3 (at index 2)
        assert_abs_diff_eq!(metrics.max_error, 0.3, epsilon = 1e-10);
    }

    #[test]
    fn test_serde_training_progress() {
        let progress = TrainingProgress::new(50, 0.25, 0.15, 0.05, 0.05, 0.0005, 2500);
        let json = serde_json::to_string(&progress).expect("Failed to serialize");
        let deserialized: TrainingProgress =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(progress.epoch, deserialized.epoch);
        assert_abs_diff_eq!(progress.total_loss, deserialized.total_loss);
    }

    #[test]
    fn test_serde_benchmark_result() {
        let result = BenchmarkResult::new(
            ProblemType::Heat1D,
            5.0,
            0.01,
            0.05,
            0.1,
            128.0,
            500.0,
            vec![1.0, 0.5, 0.1],
        );

        let json = serde_json::to_string(&result).expect("Failed to serialize");
        let deserialized: BenchmarkResult =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(result.problem, deserialized.problem);
        assert_eq!(
            result.convergence_history.len(),
            deserialized.convergence_history.len()
        );
    }

    #[test]
    fn test_serde_request() {
        let config = BenchmarkConfig::default_for_problem(ProblemType::Poisson2D);
        let request = PINNBenchmarkRequest::Initialize { config };

        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let _deserialized: PINNBenchmarkRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
    }

    #[test]
    fn test_serde_response() {
        let progress = TrainingProgress::new(100, 0.5, 0.3, 0.1, 0.1, 0.001, 5000);
        let response = PINNBenchmarkResponse::TrainingProgress { progress };

        let json = serde_json::to_string(&response).expect("Failed to serialize");
        let _deserialized: PINNBenchmarkResponse =
            serde_json::from_str(&json).expect("Failed to deserialize");
    }

    #[test]
    fn test_comparison_result_creation() {
        let pinn_result = BenchmarkResult::new(
            ProblemType::Burgers1D,
            10.0,
            0.01,
            0.05,
            0.1,
            256.0,
            1000.0,
            vec![1.0, 0.5, 0.1],
        );

        let accuracy = AccuracyMetrics::new(0.01, 0.02, 0.05, 0.1, 0.95);

        let comparison = ComparisonResult::new(pinn_result, 50.0, 100.0, accuracy);

        assert_abs_diff_eq!(comparison.reference_time_s, 50.0);
        assert_abs_diff_eq!(comparison.speedup_inference, 100.0);
    }
}