//! Neural network architecture for source localization PINN.
//!
//! Implements a physics-informed network with Fourier feature encoding
//! for learning neural source distributions from MEG/EEG data.

use crate::error::{PinnError, PinnResult};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Fourier feature encoding for positional inputs
///
/// Maps 3D coordinates to higher-dimensional space using:
/// γ(x) = [sin(2πBx), cos(2πBx)]
///
/// where B is a matrix of learnable or random frequencies.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FourierFeatures {
    /// Frequency matrix [n_features, 3]
    frequencies: Array2<f64>,
    /// Scale factor for frequencies
    scale: f64,
}

impl FourierFeatures {
    /// Create new Fourier features with random frequencies
    pub fn new(n_features: usize, scale: f64) -> Self {
        // Initialize with random frequencies (normally distributed)
        let mut frequencies = Array2::zeros((n_features, 3));
        for i in 0..n_features {
            for j in 0..3 {
                // Simple pseudo-random initialization
                let seed = (i * 3 + j) as f64;
                frequencies[[i, j]] = (seed * 0.618033988749895).fract() * 2.0 - 1.0;
            }
        }

        Self {
            frequencies: frequencies * scale,
            scale,
        }
    }

    /// Create with specific frequencies
    pub fn with_frequencies(frequencies: Array2<f64>, scale: f64) -> Self {
        Self { frequencies, scale }
    }

    /// Number of output features (2x input for sin and cos)
    pub fn output_dim(&self) -> usize {
        self.frequencies.nrows() * 2
    }

    /// Encode 3D coordinates
    pub fn encode(&self, coords: &Array2<f64>) -> Array2<f64> {
        let n_points = coords.nrows();
        let n_freqs = self.frequencies.nrows();
        let mut output = Array2::zeros((n_points, n_freqs * 2));

        for p in 0..n_points {
            for f in 0..n_freqs {
                // Compute dot product: B[f] · x[p]
                let mut dot = 0.0;
                for d in 0..3 {
                    dot += self.frequencies[[f, d]] * coords[[p, d]];
                }
                let angle = 2.0 * std::f64::consts::PI * dot;
                output[[p, f]] = angle.sin();
                output[[p, f + n_freqs]] = angle.cos();
            }
        }

        output
    }

    /// Encode a single point
    pub fn encode_point(&self, point: &[f64; 3]) -> Array1<f64> {
        let n_freqs = self.frequencies.nrows();
        let mut output = Array1::zeros(n_freqs * 2);

        for f in 0..n_freqs {
            let mut dot = 0.0;
            for d in 0..3 {
                dot += self.frequencies[[f, d]] * point[d];
            }
            let angle = 2.0 * std::f64::consts::PI * dot;
            output[f] = angle.sin();
            output[f + n_freqs] = angle.cos();
        }

        output
    }
}

/// Activation function types
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum Activation {
    /// Hyperbolic tangent
    Tanh,
    /// Rectified Linear Unit
    ReLU,
    /// Sigmoid
    Sigmoid,
    /// Swish (x * sigmoid(x))
    Swish,
    /// GELU (Gaussian Error Linear Unit)
    GELU,
}

impl Activation {
    /// Apply activation function
    pub fn apply(&self, x: f64) -> f64 {
        match self {
            Activation::Tanh => x.tanh(),
            Activation::ReLU => x.max(0.0),
            Activation::Sigmoid => 1.0 / (1.0 + (-x).exp()),
            Activation::Swish => x / (1.0 + (-x).exp()),
            Activation::GELU => {
                let sqrt2 = std::f64::consts::SQRT_2;
                0.5 * x * (1.0 + libm::erf(x / sqrt2))
            }
        }
    }

    /// Apply derivative for backpropagation
    pub fn derivative(&self, x: f64) -> f64 {
        match self {
            Activation::Tanh => {
                let t = x.tanh();
                1.0 - t * t
            }
            Activation::ReLU => {
                if x > 0.0 {
                    1.0
                } else {
                    0.0
                }
            }
            Activation::Sigmoid => {
                let s = 1.0 / (1.0 + (-x).exp());
                s * (1.0 - s)
            }
            Activation::Swish => {
                let s = 1.0 / (1.0 + (-x).exp());
                s + x * s * (1.0 - s)
            }
            Activation::GELU => {
                let sqrt2 = std::f64::consts::SQRT_2;
                let sqrt2pi = (2.0 * std::f64::consts::PI).sqrt();
                let erf_term = libm::erf(x / sqrt2);
                let exp_term = (-0.5 * x * x).exp();
                0.5 * (1.0 + erf_term) + x * exp_term / sqrt2pi
            }
        }
    }
}

/// Dense (fully connected) layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DenseLayer {
    /// Weight matrix [in_features, out_features]
    weights: Array2<f64>,
    /// Bias vector [out_features]
    biases: Array1<f64>,
    /// Activation function
    activation: Option<Activation>,
}

impl DenseLayer {
    /// Create a new dense layer with Xavier initialization
    pub fn new(in_features: usize, out_features: usize, activation: Option<Activation>) -> Self {
        // Xavier/Glorot initialization
        let scale = (2.0 / (in_features + out_features) as f64).sqrt();
        let mut weights = Array2::zeros((in_features, out_features));

        // Simple deterministic initialization for reproducibility
        for i in 0..in_features {
            for j in 0..out_features {
                let seed = (i * out_features + j) as f64;
                weights[[i, j]] = ((seed * 0.618033988749895).fract() * 2.0 - 1.0) * scale;
            }
        }

        let biases = Array1::zeros(out_features);

        Self {
            weights,
            biases,
            activation,
        }
    }

    /// Forward pass
    pub fn forward(&self, input: &Array2<f64>) -> Array2<f64> {
        let linear = input.dot(&self.weights) + &self.biases;

        if let Some(act) = self.activation {
            linear.mapv(|x| act.apply(x))
        } else {
            linear
        }
    }

    /// Forward pass for single input
    pub fn forward_single(&self, input: &Array1<f64>) -> Array1<f64> {
        let linear = input.dot(&self.weights) + &self.biases;

        if let Some(act) = self.activation {
            linear.mapv(|x| act.apply(x))
        } else {
            linear
        }
    }

    /// Get weights (for updates)
    pub fn weights(&self) -> &Array2<f64> {
        &self.weights
    }

    /// Get mutable weights
    pub fn weights_mut(&mut self) -> &mut Array2<f64> {
        &mut self.weights
    }

    /// Get biases
    pub fn biases(&self) -> &Array1<f64> {
        &self.biases
    }

    /// Get mutable biases
    pub fn biases_mut(&mut self) -> &mut Array1<f64> {
        &mut self.biases
    }
}

/// Source network configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceNetworkConfig {
    /// Hidden layer sizes
    pub hidden_layers: Vec<usize>,
    /// Number of Fourier features
    pub n_fourier_features: usize,
    /// Fourier feature scale
    pub fourier_scale: f64,
    /// Activation function
    pub activation: Activation,
    /// Output current density components (typically 3 for Jx, Jy, Jz)
    pub output_dim: usize,
    /// Whether to also output conductivity
    pub output_conductivity: bool,
}

impl Default for SourceNetworkConfig {
    fn default() -> Self {
        Self {
            hidden_layers: vec![64, 128, 128, 64],
            n_fourier_features: 32,
            fourier_scale: 10.0,
            activation: Activation::Tanh,
            output_dim: 3,
            output_conductivity: false,
        }
    }
}

/// Neural network for source localization
///
/// Takes 3D coordinates as input and outputs:
/// - Current density components (Jx, Jy, Jz)
/// - Optionally, tissue conductivity σ(x)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceNetwork {
    /// Configuration
    config: SourceNetworkConfig,
    /// Fourier feature encoder
    fourier: FourierFeatures,
    /// Hidden layers
    layers: Vec<DenseLayer>,
    /// Current density output layer
    current_output: DenseLayer,
    /// Conductivity output layer (optional)
    conductivity_output: Option<DenseLayer>,
}

impl SourceNetwork {
    /// Create a new source network
    pub fn new(config: SourceNetworkConfig) -> Self {
        let fourier = FourierFeatures::new(config.n_fourier_features, config.fourier_scale);

        // Build hidden layers
        let input_dim = fourier.output_dim();
        let mut layers = Vec::new();
        let mut prev_dim = input_dim;

        for &hidden_dim in &config.hidden_layers {
            layers.push(DenseLayer::new(
                prev_dim,
                hidden_dim,
                Some(config.activation),
            ));
            prev_dim = hidden_dim;
        }

        // Output layer for current density (no activation - unbounded output)
        let current_output = DenseLayer::new(prev_dim, config.output_dim, None);

        // Optional conductivity output (sigmoid for bounded positive output)
        let conductivity_output = if config.output_conductivity {
            Some(DenseLayer::new(prev_dim, 1, Some(Activation::Sigmoid)))
        } else {
            None
        };

        Self {
            config,
            fourier,
            layers,
            current_output,
            conductivity_output,
        }
    }

    /// Forward pass for multiple points
    ///
    /// # Arguments
    /// * `coords` - Input coordinates [n_points, 3]
    ///
    /// # Returns
    /// * Current density [n_points, 3] and optionally conductivity [n_points]
    pub fn forward(&self, coords: &Array2<f64>) -> (Array2<f64>, Option<Array1<f64>>) {
        // Fourier encoding
        let mut x = self.fourier.encode(coords);

        // Hidden layers
        for layer in &self.layers {
            x = layer.forward(&x);
        }

        // Current density output
        let current = self.current_output.forward(&x);

        // Conductivity output (if enabled)
        let conductivity = self.conductivity_output.as_ref().map(|layer| {
            let sigma = layer.forward(&x);
            // Convert to 1D array and scale to reasonable conductivity range (0.01 - 2.0 S/m)
            sigma.column(0).mapv(|s| 0.01 + s * 1.99).to_owned()
        });

        (current, conductivity)
    }

    /// Forward pass for single point
    pub fn forward_single(&self, point: &[f64; 3]) -> (Array1<f64>, Option<f64>) {
        // Fourier encoding
        let mut x = self.fourier.encode_point(point);

        // Hidden layers
        for layer in &self.layers {
            x = layer.forward_single(&x);
        }

        // Current density output
        let current = self.current_output.forward_single(&x);

        // Conductivity output
        let conductivity = self.conductivity_output.as_ref().map(|layer| {
            let sigma = layer.forward_single(&x);
            0.01 + sigma[0] * 1.99
        });

        (current, conductivity)
    }

    /// Get network configuration
    pub fn config(&self) -> &SourceNetworkConfig {
        &self.config
    }

    /// Total number of parameters
    pub fn n_parameters(&self) -> usize {
        let mut count = 0;

        for layer in &self.layers {
            count += layer.weights().len();
            count += layer.biases().len();
        }

        count += self.current_output.weights().len();
        count += self.current_output.biases().len();

        if let Some(ref layer) = self.conductivity_output {
            count += layer.weights().len();
            count += layer.biases().len();
        }

        count
    }

    /// Get all layer weights as flat vector (for optimization)
    pub fn get_parameters(&self) -> Vec<f64> {
        let mut params = Vec::with_capacity(self.n_parameters());

        for layer in &self.layers {
            params.extend(layer.weights().iter());
            params.extend(layer.biases().iter());
        }

        params.extend(self.current_output.weights().iter());
        params.extend(self.current_output.biases().iter());

        if let Some(ref layer) = self.conductivity_output {
            params.extend(layer.weights().iter());
            params.extend(layer.biases().iter());
        }

        params
    }

    /// Set all layer weights from flat vector
    pub fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
        if params.len() != self.n_parameters() {
            return Err(PinnError::DimensionMismatch(format!(
                "Expected {} parameters, got {}",
                self.n_parameters(),
                params.len()
            )));
        }

        let mut idx = 0;

        for layer in &mut self.layers {
            let w_len = layer.weights().len();
            let b_len = layer.biases().len();

            for (i, w) in layer.weights_mut().iter_mut().enumerate() {
                *w = params[idx + i];
            }
            idx += w_len;

            for (i, b) in layer.biases_mut().iter_mut().enumerate() {
                *b = params[idx + i];
            }
            idx += b_len;
        }

        // Current output layer
        let w_len = self.current_output.weights().len();
        let b_len = self.current_output.biases().len();

        for (i, w) in self.current_output.weights_mut().iter_mut().enumerate() {
            *w = params[idx + i];
        }
        idx += w_len;

        for (i, b) in self.current_output.biases_mut().iter_mut().enumerate() {
            *b = params[idx + i];
        }
        idx += b_len;

        // Conductivity layer (if present)
        if let Some(ref mut layer) = self.conductivity_output {
            let w_len = layer.weights().len();

            for (i, w) in layer.weights_mut().iter_mut().enumerate() {
                *w = params[idx + i];
            }
            idx += w_len;

            for b in layer.biases_mut().iter_mut() {
                *b = params[idx];
                idx += 1;
            }
        }

        Ok(())
    }
}

/// Sensor network for mapping sensor data to latent representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SensorEncoder {
    /// Encoder layers
    layers: Vec<DenseLayer>,
    /// Output dimension
    latent_dim: usize,
}

impl SensorEncoder {
    /// Create a new sensor encoder
    pub fn new(n_sensors: usize, hidden_dims: &[usize], latent_dim: usize) -> Self {
        let mut layers = Vec::new();
        let mut prev_dim = n_sensors;

        for &dim in hidden_dims {
            layers.push(DenseLayer::new(prev_dim, dim, Some(Activation::Tanh)));
            prev_dim = dim;
        }

        // Final layer to latent space
        layers.push(DenseLayer::new(prev_dim, latent_dim, None));

        Self { layers, latent_dim }
    }

    /// Encode sensor measurements
    pub fn encode(&self, measurements: &Array1<f64>) -> Array1<f64> {
        let mut x = measurements.clone();
        for layer in &self.layers {
            x = layer.forward_single(&x);
        }
        x
    }

    /// Encode batch of measurements
    pub fn encode_batch(&self, measurements: &Array2<f64>) -> Array2<f64> {
        let mut x = measurements.clone();
        for layer in &self.layers {
            x = layer.forward(&x);
        }
        x
    }

    /// Get latent dimension
    pub fn latent_dim(&self) -> usize {
        self.latent_dim
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn test_fourier_features() {
        let ff = FourierFeatures::new(16, 10.0);
        assert_eq!(ff.output_dim(), 32);

        let coords = Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 0.05, 0.01, 0.0, 0.05]).unwrap();

        let encoded = ff.encode(&coords);
        assert_eq!(encoded.shape(), &[2, 32]);
    }

    #[test]
    fn test_activation_functions() {
        assert!((Activation::Tanh.apply(0.0)).abs() < 1e-10);
        assert!((Activation::ReLU.apply(-1.0)).abs() < 1e-10);
        assert!((Activation::Sigmoid.apply(0.0) - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_dense_layer() {
        let layer = DenseLayer::new(4, 8, Some(Activation::Tanh));
        let input = Array2::zeros((3, 4));
        let output = layer.forward(&input);
        assert_eq!(output.shape(), &[3, 8]);
    }

    #[test]
    fn test_source_network() {
        let config = SourceNetworkConfig {
            hidden_layers: vec![32, 64, 32],
            n_fourier_features: 16,
            fourier_scale: 10.0,
            activation: Activation::Tanh,
            output_dim: 3,
            output_conductivity: true,
        };

        let network = SourceNetwork::new(config);

        let coords = Array2::from_shape_vec(
            (5, 3),
            vec![
                0.0, 0.0, 0.05, 0.01, 0.0, 0.05, -0.01, 0.0, 0.05, 0.0, 0.01, 0.05, 0.0, -0.01,
                0.05,
            ],
        )
        .unwrap();

        let (current, conductivity) = network.forward(&coords);

        assert_eq!(current.shape(), &[5, 3]);
        assert!(conductivity.is_some());
        assert_eq!(conductivity.unwrap().len(), 5);
    }

    #[test]
    fn test_parameter_io() {
        let config = SourceNetworkConfig::default();
        let mut network = SourceNetwork::new(config);

        let params = network.get_parameters();
        let n_params = network.n_parameters();

        assert_eq!(params.len(), n_params);

        // Modify and restore
        let mut new_params = params.clone();
        new_params[0] = 999.0;

        network.set_parameters(&new_params).unwrap();
        let retrieved = network.get_parameters();

        assert!((retrieved[0] - 999.0).abs() < 1e-10);
    }

    #[test]
    fn test_sensor_encoder() {
        let encoder = SensorEncoder::new(64, &[32, 16], 8);

        let measurements = Array1::zeros(64);
        let latent = encoder.encode(&measurements);

        assert_eq!(latent.len(), 8);
    }
}