//! KAN (Kolmogorov-Arnold Network) inference adapter for ensemble prediction
//!
//! Wraps a `KANNetwork` and normalizes its scalar output into a directional
//! signal + confidence via sigmoid for ensemble aggregation.

use std::sync::{Arc, Mutex};

use cudarc::driver::CudaStream;

use crate::ensemble::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
};
use crate::kan::config::KANConfig;
use crate::kan::network::KANNetwork;
use crate::{MLError, MLResult};

use ml_supervised::gpu_tensor::GpuTensor;

/// Inference adapter that wraps a KAN (B-spline activation) network.
#[allow(missing_debug_implementations)]
pub struct KanInferenceAdapter {
    model: Mutex<KANNetwork>,
    stream: Arc<CudaStream>,
    input_dim: usize,
}

// SAFETY: KANNetwork uses GpuTensor/CudaSlice which are Send+Sync.
#[allow(unsafe_code)]
unsafe impl Send for KanInferenceAdapter {}
#[allow(unsafe_code)]
unsafe impl Sync for KanInferenceAdapter {}

impl KanInferenceAdapter {
    /// Create a new KAN inference adapter from configuration.
    pub fn new(config: KANConfig) -> MLResult<Self> {
        let input_dim = config.layer_widths.first().copied().unwrap_or(51);

        let ctx = cudarc::driver::CudaContext::new(0)
            .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
        let stream = ctx
            .new_stream()
            .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;

        let network = KANNetwork::new(&config, &stream)?;

        Ok(Self {
            model: Mutex::new(network),
            stream,
            input_dim,
        })
    }

    /// Load from checkpoint (placeholder -- weight loading TBD).
    pub fn from_checkpoint(config: KANConfig, _path: &str) -> MLResult<Self> {
        Self::new(config)
    }

    /// Pad or truncate feature values to match input_dim.
    fn pad_features(&self, values: &[f64]) -> Vec<f32> {
        let mut padded = vec![0.0_f32; self.input_dim];
        let copy_len = values.len().min(self.input_dim);
        for i in 0..copy_len {
            if let Some(v) = values.get(i) {
                if let Some(slot) = padded.get_mut(i) {
                    *slot = *v as f32;
                }
            }
        }
        padded
    }
}

impl ModelInferenceAdapter for KanInferenceAdapter {
    fn model_name(&self) -> &str {
        "KAN"
    }

    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
        let start = std::time::Instant::now();

        let padded = self.pad_features(&features.values);
        let input = GpuTensor::from_vec(padded, &[1, self.input_dim], &self.stream)?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("KAN model lock poisoned: {e}")))?;
        let output = model.forward(&input)?;
        let output_host = output.to_vec()?; // gpu-exit: scalar prediction

        let raw_val = output_host.first().copied().unwrap_or(0.0) as f64;

        // Sigmoid -> direction + confidence
        let prob = 1.0 / (1.0 + (-raw_val).exp());
        let direction = (2.0_f64 * prob - 1.0).clamp(-1.0, 1.0);
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        let latency_us = start.elapsed().as_micros() as u64;

        Ok(EnsemblePrediction {
            model_name: "KAN".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: None,
            },
        })
    }

    fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
        let padded = self.pad_features(&features.values);
        let input = GpuTensor::from_vec(padded, &[1, self.input_dim], &self.stream)?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("KAN model lock poisoned: {e}")))?;
        let output = model.forward(&input)?;
        let output_host = output.to_vec()?; // gpu-exit: scalar prediction

        let raw_f32 = output_host.first().copied().unwrap_or(0.0);
        let prob = 1.0 / (1.0 + (-raw_f32 as f64).exp());
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        Ok(RawPrediction {
            direction_scalar: 0.0,
            confidence,
            logits: Some(vec![raw_f32]),
        })
    }

    fn predict_batch(&self, batch: &[FeatureVector]) -> MLResult<Vec<EnsemblePrediction>> {
        if batch.is_empty() {
            return Ok(Vec::new());
        }
        let start = std::time::Instant::now();
        let n = batch.len();

        let mut flat = Vec::with_capacity(n * self.input_dim);
        for fv in batch {
            flat.extend_from_slice(&self.pad_features(&fv.values));
        }

        let input = GpuTensor::from_vec(flat, &[n, self.input_dim], &self.stream)?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("KAN lock: {e}")))?;
        let output = model.forward(&input)?;
        drop(model);

        let output_host = output.to_vec()?; // gpu-exit: scalar prediction
        let out_dim = if output.shape.len() == 2 {
            output.dim(1).unwrap_or(1)
        } else {
            1
        };
        let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64;

        let mut results = Vec::with_capacity(n);
        for i in 0..n {
            let raw_val = output_host
                .get(i * out_dim)
                .copied()
                .unwrap_or(0.0) as f64;
            let prob = 1.0 / (1.0 + (-raw_val).exp());
            let direction = (2.0_f64 * prob - 1.0).clamp(-1.0, 1.0);
            let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

            results.push(EnsemblePrediction {
                model_name: "KAN".to_owned(),
                direction,
                confidence,
                metadata: PredictionMeta {
                    latency_us,
                    quantiles: None,
                    attention_weights: None,
                    q_values: None,
                },
            });
        }

        Ok(results)
    }

    fn is_ready(&self) -> bool {
        self.model.lock().is_ok()
    }
}

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

    fn test_config() -> KANConfig {
        KANConfig {
            layer_widths: vec![51, 32, 16, 1],
            grid_size: 3,
            spline_order: 3,
            ..Default::default()
        }
    }

    #[test]
    fn test_kan_adapter_creation() {
        let adapter = KanInferenceAdapter::new(test_config());
        assert!(
            adapter.is_ok(),
            "KAN adapter creation failed: {:?}",
            adapter.err()
        );
        let adapter = adapter.unwrap();
        assert_eq!(adapter.model_name(), "KAN");
        assert!(adapter.is_ready());
    }

    #[test]
    fn test_kan_adapter_predict_direction_range() {
        let adapter = KanInferenceAdapter::new(test_config()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000_000_000,
        };
        let pred = adapter.predict(&fv).unwrap();
        assert!(
            pred.direction >= -1.0 && pred.direction <= 1.0,
            "direction {} out of [-1,1]",
            pred.direction
        );
        assert!(
            pred.confidence >= 0.0 && pred.confidence <= 1.0,
            "confidence {} out of [0,1]",
            pred.confidence
        );
    }

    #[test]
    fn test_kan_adapter_deterministic() {
        let adapter = KanInferenceAdapter::new(test_config()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000_000_000,
        };
        let pred1 = adapter.predict(&fv).unwrap();
        let pred2 = adapter.predict(&fv).unwrap();
        assert_eq!(
            pred1.direction, pred2.direction,
            "Deterministic predictions should have same direction"
        );
    }
}
