//! UnifiedTrainable adapter for KAN (Kolmogorov-Arnold Network).
//!
//! The KAN network in ml-supervised uses cuBLAS-backed GpuTensor for forward.
//! This adapter bridges to the unified training interface used by the
//! training orchestrator.
//!
//! Forward/backward use GPU-native operations via the cuda_autograd system.
//! The GpuVarStore holds all trainable parameters and the GpuAdamW optimizer
//! runs the update step entirely on GPU.

use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::{
    ActivationKernels, AdamWConfig, GpuAdamW, GpuLinear, GpuTensor, GpuVarStore, LossKernels,
    linear::LinearActivations,
};

use super::config::KANConfig;
use super::network::KANNetwork;

use crate::training::unified_trainer::{
    checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable,
};
use crate::MLError;

/// Activations saved during forward pass, needed for backward pass.
struct CachedForward {
    acts_input: LinearActivations,
    relu_mask: GpuTensor,
    acts_output: LinearActivations,
    loss_grad: GpuTensor,
}

/// KAN trainable adapter implementing UnifiedTrainable.
///
/// Owns a GpuVarStore with projection layers, a GpuAdamW optimizer,
/// and the KAN network for forward inference.
pub struct KANTrainableAdapter {
    config: KANConfig,
    network: KANNetwork,
    stream: Arc<CudaStream>,
    cublas: CudaBlas,
    var_store: GpuVarStore,
    input_linear: GpuLinear,
    output_linear: GpuLinear,
    optimizer: GpuAdamW,
    activation_kernels: ActivationKernels,
    loss_kernels: LossKernels,
    learning_rate: f64,
    step: usize,
    latest_metrics: TrainingMetrics,
    loss_history: Vec<f64>,
    last_grad_norm: f64,
    /// Cached forward-pass activations for backpropagation (GPU-resident).
    cached_fwd: Option<CachedForward>,
}

impl std::fmt::Debug for KANTrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("KANTrainableAdapter")
            .field("config", &self.config)
            .field("learning_rate", &self.learning_rate)
            .field("step", &self.step)
            .field("loss_history_len", &self.loss_history.len())
            .field("last_grad_norm", &self.last_grad_norm)
            .finish_non_exhaustive()
    }
}

impl KANTrainableAdapter {
    /// Create a new KAN trainable adapter.
    pub fn new(config: KANConfig) -> Result<Self, MLError> {
        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 cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| MLError::DeviceError(format!("cuBLAS: {e}")))?;

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

        // Build projection layers in the var store
        let input_dim = config.layer_widths.first().copied().unwrap_or(8);
        let output_dim = config.layer_widths.last().copied().unwrap_or(1);

        let mut var_store = GpuVarStore::new(Arc::clone(&stream));
        let input_linear = var_store.linear("input", input_dim, output_dim)?;
        let output_linear = var_store.linear("output", output_dim, 1)?;

        let optimizer = GpuAdamW::new(
            AdamWConfig {
                lr: config.learning_rate as f32,
                weight_decay: config.weight_decay as f32,
                ..AdamWConfig::default()
            },
            Arc::clone(&stream),
        )?;

        let activation_kernels = ActivationKernels::new(&stream)?;
        let loss_kernels = LossKernels::new(&stream)?;

        let learning_rate = config.learning_rate;

        Ok(Self {
            config,
            network,
            stream,
            cublas,
            var_store,
            input_linear,
            output_linear,
            optimizer,
            activation_kernels,
            loss_kernels,
            learning_rate,
            step: 0,
            latest_metrics: TrainingMetrics::default(),
            loss_history: Vec::new(),
            last_grad_norm: 0.0,
            cached_fwd: None,
        })
    }

    /// Access the KAN configuration.
    pub fn kan_config(&self) -> &KANConfig {
        &self.config
    }
}

impl UnifiedTrainable for KANTrainableAdapter {
    fn model_type(&self) -> &str {
        "KAN"
    }

    fn device_name(&self) -> String {
        "cuda:0".to_owned()
    }

    fn forward_loss(&mut self, input: &[f32], target: &[f32]) -> Result<f64, MLError> {
        let input_dim = self.input_linear.in_dim;
        let batch = input.len() / input_dim;
        if batch == 0 {
            return Err(MLError::InvalidInput("Empty input".to_owned()));
        }

        // Upload input and target to GPU
        let x = GpuTensor::from_host(input, vec![batch, input_dim], &self.stream)?;
        let t = GpuTensor::from_host(target, vec![target.len()], &self.stream)?;

        // input_linear (save activations for backward)
        let (hidden, acts_input) =
            self.input_linear
                .forward(&x, &self.var_store, &self.cublas, &self.stream)?;

        // ReLU (save mask for backward)
        let (activated, relu_mask) = self.activation_kernels.relu_fwd(&hidden, &self.stream)?;

        // output_linear (save activations for backward)
        let (pred, acts_output) =
            self.output_linear
                .forward(&activated, &self.var_store, &self.cublas, &self.stream)?;

        // MSE loss (returns loss scalar + dL/d_pred gradient)
        let result = self.loss_kernels.mse(&pred, &t, &self.stream)?;

        let loss_host = result.loss.to_host(&self.stream)?; // gpu-exit: 1 scalar loss
        let loss_val = loss_host.first().copied().unwrap_or(0.0) as f64;

        // Cache forward activations + loss gradient for backward
        self.cached_fwd = Some(CachedForward {
            acts_input,
            relu_mask,
            acts_output,
            loss_grad: result.grad,
        });

        self.loss_history.push(loss_val);
        self.latest_metrics.loss = loss_val;

        Ok(loss_val)
    }

    fn backward(&mut self, _loss_value: f64) -> Result<f64, MLError> {
        let fwd = self.cached_fwd.take().ok_or_else(|| {
            MLError::TrainingError("backward() called without prior forward_loss()".to_owned())
        })?;

        // Backprop through output_linear
        let output_grads = self.output_linear.backward(
            &fwd.loss_grad, &fwd.acts_output, &self.var_store, &self.cublas, &self.stream,
        )?;

        // Backprop through ReLU
        let d_hidden = self.activation_kernels.relu_bwd(
            &output_grads.dx, &fwd.relu_mask, &self.stream,
        )?;

        // Backprop through input_linear
        let input_grads = self.input_linear.backward(
            &d_hidden, &fwd.acts_input, &self.var_store, &self.cublas, &self.stream,
        )?;

        // Collect gradients for optimizer
        let mut gradients = BTreeMap::new();
        gradients.insert("output.weight".to_owned(), output_grads.dw);
        gradients.insert("output.bias".to_owned(), output_grads.db);
        gradients.insert("input.weight".to_owned(), input_grads.dw);
        gradients.insert("input.bias".to_owned(), input_grads.db);

        // Apply GpuAdamW (includes gradient clipping)
        let grad_norm = self.optimizer.step(&mut self.var_store, &gradients)
            .map_err(|e| MLError::TrainingError(format!("AdamW step failed: {e}")))?;

        self.last_grad_norm = grad_norm as f64;
        self.latest_metrics.grad_norm = Some(self.last_grad_norm);
        self.step += 1;

        Ok(self.last_grad_norm)
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        // Optimizer step is fused into backward() for efficiency (single pass).
        Ok(())
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        self.cached_fwd = None;
        Ok(())
    }

    fn get_learning_rate(&self) -> f64 {
        self.learning_rate
    }

    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        self.learning_rate = lr;
        self.optimizer.set_learning_rate(lr as f32);
        Ok(())
    }

    fn get_step(&self) -> usize {
        self.step
    }

    fn collect_metrics(&self) -> TrainingMetrics {
        let mut metrics = self.latest_metrics.clone();

        if !self.loss_history.is_empty() {
            let recent: Vec<f64> = self.loss_history.iter().rev().take(100).copied().collect();
            metrics.loss = recent.iter().sum::<f64>() / recent.len() as f64;
        }

        metrics.learning_rate = self.learning_rate;
        metrics
            .custom_metrics
            .insert("training_steps".to_owned(), self.step as f64);
        metrics
            .custom_metrics
            .insert("loss_history_len".to_owned(), self.loss_history.len() as f64);
        metrics
            .custom_metrics
            .insert("grid_size".to_owned(), self.config.grid_size as f64);
        metrics
            .custom_metrics
            .insert("spline_order".to_owned(), self.config.spline_order as f64);
        metrics.custom_metrics.insert(
            "num_layers".to_owned(),
            (self.config.layer_widths.len().saturating_sub(1)) as f64,
        );

        metrics
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        let metadata = CheckpointMetadata {
            model_type: "KAN".to_owned(),
            version: "1.0.0".to_owned(),
            epoch: 0,
            step: self.step,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&self.config).map_err(|e| {
                MLError::SerializationError {
                    reason: format!("Failed to serialize config: {}", e),
                }
            })?,
            metrics: self.collect_metrics(),
        };

        checkpoint::save_metadata(&metadata, checkpoint_path)?;

        // Save weights from var store
        let weights_path = format!("{}.weights.json", checkpoint_path);
        let exported = self.var_store.export_to_host()?; // gpu-exit: checkpoint save
        let mut all_weights: HashMap<String, Vec<f32>> = HashMap::new();
        for (name, (_shape, data)) in &exported {
            all_weights.insert(name.clone(), data.clone());
        }
        let json = serde_json::to_string(&all_weights).map_err(|e| {
            MLError::CheckpointError(format!("Failed to serialize weights: {}", e))
        })?;
        std::fs::write(&weights_path, json).map_err(|e| {
            MLError::CheckpointError(format!("Failed to write weights: {}", e))
        })?;

        tracing::info!(
            "Saved KAN checkpoint to {} (step {})",
            checkpoint_path,
            self.step
        );

        Ok(checkpoint_path.to_string())
    }

    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        let metadata = checkpoint::load_metadata(checkpoint_path)?;

        if metadata.model_type != "KAN" {
            return Err(MLError::CheckpointError(format!(
                "Invalid model type in checkpoint: expected KAN, got {}",
                metadata.model_type
            )));
        }

        self.step = metadata.step;
        self.latest_metrics = metadata.metrics.clone();

        tracing::info!(
            "Loaded KAN checkpoint from {} (step {})",
            checkpoint_path,
            metadata.step
        );

        Ok(metadata)
    }
}

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

    fn make_config() -> KANConfig {
        KANConfig {
            grid_size: 3,
            spline_order: 3,
            layer_widths: vec![8, 4, 1],
            learning_rate: 1e-3,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }

    #[test]
    fn test_model_type() {
        let cfg = make_config();
        let adapter = KANTrainableAdapter::new(cfg);
        if let Ok(a) = adapter {
            assert_eq!(a.model_type(), "KAN");
        }
    }

    #[test]
    fn test_device_name() {
        let cfg = make_config();
        if let Ok(a) = KANTrainableAdapter::new(cfg) {
            assert_eq!(a.device_name(), "cuda:0");
        }
    }

    #[test]
    fn test_learning_rate_get_set() {
        let cfg = make_config();
        if let Ok(mut adapter) = KANTrainableAdapter::new(cfg) {
            let original_lr = adapter.get_learning_rate();
            assert!((original_lr - 1e-3).abs() < 1e-10);

            adapter.set_learning_rate(5e-4).ok();
            assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10);
        }
    }

    #[test]
    fn test_collect_metrics() {
        let cfg = make_config();
        if let Ok(adapter) = KANTrainableAdapter::new(cfg) {
            let metrics = adapter.collect_metrics();
            assert!(metrics.custom_metrics.contains_key("training_steps"));
            assert!(metrics.custom_metrics.contains_key("grid_size"));
            assert!(metrics.custom_metrics.contains_key("spline_order"));
            assert!(metrics.custom_metrics.contains_key("num_layers"));
            assert_eq!(metrics.custom_metrics.get("grid_size").copied(), Some(3.0));
        }
    }
}
