//! UnifiedTrainable adapter for TGGN (Temporal Graph Gated Network)
//!
//! Wraps a GpuLinear-based projection network to provide the UnifiedTrainable
//! interface for TGGN. The projection network maps flattened graph features
//! through hidden layers, enabling gradient-based training via the unified
//! training orchestrator.
//!
//! Architecture: input_linear(node_dim -> hidden_dim) -> ReLU -> output_linear(hidden_dim -> 1)
//!
//! Forward pass uses cuBLAS-backed GpuLinear from cuda_autograd for GPU-native
//! inference. Backward/optimizer use GpuAdamW for GPU-native parameter updates.

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::TGGNConfig;

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,
}

/// Adapter wrapping TGGN with a GpuLinear-based projection network for unified training.
///
/// The projection network has two linear layers:
///   - `input_linear`: projects from `node_dim` to `hidden_dim`
///   - `output_linear`: projects from `hidden_dim` to 1 (scalar prediction)
///
/// Forward inference runs through cuBLAS-backed GpuLinear.
/// Training uses GpuAdamW for GPU-native parameter updates.
pub struct TGGNTrainableAdapter {
    config: TGGNConfig,
    var_store: GpuVarStore,
    input_linear: GpuLinear,
    output_linear: GpuLinear,
    stream: Arc<CudaStream>,
    cublas: CudaBlas,
    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 TGGNTrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TGGNTrainableAdapter")
            .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 TGGNTrainableAdapter {
    /// Create a new TGGN trainable adapter with projection network.
    pub fn new(config: TGGNConfig) -> Result<Self, MLError> {
        if config.node_dim == 0 {
            return Err(MLError::ConfigError(
                "TGGN requires node_dim > 0".to_owned(),
            ));
        }
        if config.hidden_dim == 0 {
            return Err(MLError::ConfigError(
                "TGGN requires hidden_dim > 0".to_owned(),
            ));
        }

        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 mut var_store = GpuVarStore::new(Arc::clone(&stream));
        let input_linear = var_store.linear("input", config.node_dim, config.hidden_dim)?;
        let output_linear = var_store.linear("output", config.hidden_dim, 1)?;

        let optimizer = GpuAdamW::new(
            AdamWConfig {
                lr: 1e-3,
                weight_decay: 1e-4,
                ..AdamWConfig::default()
            },
            Arc::clone(&stream),
        )?;

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

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

    /// Access the underlying TGGN configuration.
    pub fn tggn_config(&self) -> &TGGNConfig {
        &self.config
    }
}

impl UnifiedTrainable for TGGNTrainableAdapter {
    fn model_type(&self) -> &str {
        "TGGN"
    }

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

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

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

        // input_linear: node_dim -> hidden_dim (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: hidden_dim -> 1 (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("node_dim".to_owned(), self.config.node_dim as f64);
        metrics
            .custom_metrics
            .insert("hidden_dim".to_owned(), self.config.hidden_dim as f64);
        metrics
            .custom_metrics
            .insert("num_layers".to_owned(), self.config.num_layers as f64);

        metrics
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        let metadata = CheckpointMetadata {
            model_type: "TGGN".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 TGGN 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 != "TGGN" {
            return Err(MLError::CheckpointError(format!(
                "Invalid model type in checkpoint: expected TGGN, got {}",
                metadata.model_type
            )));
        }

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

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

        Ok(metadata)
    }
}

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

    fn make_config() -> TGGNConfig {
        TGGNConfig {
            max_nodes: 16,
            max_edges: 32,
            node_dim: 8,
            edge_dim: 4,
            hidden_dim: 16,
            num_layers: 2,
            temporal_decay: 0.99,
            update_frequency_ns: 1_000_000,
            use_simd: false,
        }
    }

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

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

    #[test]
    fn test_learning_rate_get_set() {
        if let Ok(mut adapter) = TGGNTrainableAdapter::new(make_config()) {
            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() {
        if let Ok(adapter) = TGGNTrainableAdapter::new(make_config()) {
            let metrics = adapter.collect_metrics();
            assert!(metrics.custom_metrics.contains_key("training_steps"));
            assert!(metrics.custom_metrics.contains_key("node_dim"));
            assert!(metrics.custom_metrics.contains_key("hidden_dim"));
            assert!(metrics.custom_metrics.contains_key("num_layers"));
            assert_eq!(metrics.custom_metrics.get("node_dim").copied(), Some(8.0));
        }
    }
}
