//! UnifiedTrainable adapter for TLOB (Time Limit Order Book)
//!
//! Wraps a GpuLinear-based projection network to provide the UnifiedTrainable
//! interface for TLOB. The projection network maps flattened order book features
//! through hidden layers, enabling gradient-based training via the unified
//! training orchestrator.
//!
//! Architecture: input_linear(seq_len*feature_dim -> d_model) -> ReLU -> output_linear(d_model -> 1)
//!
//! Forward pass uses cuBLAS-backed GpuLinear from cuda_autograd for GPU-native
//! inference. Backward/optimizer use GpuAdamW for GPU-native parameter updates.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

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

use std::collections::BTreeMap;

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

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

/// Configuration for the TLOB trainable adapter projection network.
///
/// This is separate from the ONNX-oriented `TLOBConfig` in `transformer.rs`.
/// It defines the architecture for the projection network used
/// by the unified training orchestrator.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TLOBAdapterConfig {
    /// Model / hidden dimension for projection layers
    pub d_model: usize,
    /// Number of attention heads (reserved for future transformer layers)
    pub num_heads: usize,
    /// Number of projection layers (reserved for future deeper networks)
    pub num_layers: usize,
    /// Sequence length (number of time steps in LOB snapshot)
    pub seq_len: usize,
    /// Feature dimension per time step (e.g. 51 for TLOB features)
    pub feature_dim: usize,
}

impl Default for TLOBAdapterConfig {
    fn default() -> Self {
        Self {
            d_model: 128,
            num_heads: 4,
            num_layers: 2,
            seq_len: 128,
            feature_dim: 51,
        }
    }
}

/// Adapter wrapping TLOB with a GpuLinear-based projection network for unified training.
///
/// The projection network has two linear layers:
///   - `input_linear`: projects from `seq_len * feature_dim` to `d_model`
///   - `output_linear`: projects from `d_model` to 1 (scalar prediction)
///
/// Forward inference runs through cuBLAS-backed GpuLinear.
/// Training uses GpuAdamW for GPU-native parameter updates.
pub struct TLOBTrainableAdapter {
    config: TLOBAdapterConfig,
    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>,
}

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

impl std::fmt::Debug for TLOBTrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TLOBTrainableAdapter")
            .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 TLOBTrainableAdapter {
    /// Create a new TLOB trainable adapter with projection network.
    pub fn new(config: TLOBAdapterConfig) -> Result<Self, MLError> {
        if config.seq_len == 0 || config.feature_dim == 0 {
            return Err(MLError::ConfigError(format!(
                "TLOB requires seq_len > 0 and feature_dim > 0 (got {}x{})",
                config.seq_len, config.feature_dim
            )));
        }

        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 input_dim = config.seq_len * config.feature_dim;

        let mut var_store = GpuVarStore::new(Arc::clone(&stream));
        let input_linear = var_store.linear("input", input_dim, config.d_model)?;
        let output_linear = var_store.linear("output", config.d_model, 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 TLOB adapter configuration.
    pub fn tlob_config(&self) -> &TLOBAdapterConfig {
        &self.config
    }
}

impl UnifiedTrainable for TLOBTrainableAdapter {
    fn model_type(&self) -> &str {
        "TLOB"
    }

    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.config.seq_len * self.config.feature_dim;
        let batch = input.len() / input_dim;
        if batch == 0 {
            return Err(MLError::InvalidInput("Empty input".to_owned()));
        }

        // Upload 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: seq_len*feature_dim -> d_model (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: d_model -> 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: dL/d_activated
        let output_grads = self.output_linear.backward(
            &fwd.loss_grad,
            &fwd.acts_output,
            &self.var_store,
            &self.cublas,
            &self.stream,
        )?;

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

        // Backprop through input_linear: dL/d_input (not needed) + dW, dB
        let input_grads = self.input_linear.backward(
            &d_hidden,
            &fwd.acts_input,
            &self.var_store,
            &self.cublas,
            &self.stream,
        )?;

        // Collect all parameter 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 update (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).
        // This is called by the orchestrator after backward(), but the work is
        // already done. Just ensure consistency.
        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("d_model".to_owned(), self.config.d_model as f64);
        metrics
            .custom_metrics
            .insert("seq_len".to_owned(), self.config.seq_len as f64);
        metrics
            .custom_metrics
            .insert("feature_dim".to_owned(), self.config.feature_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: "TLOB".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 TLOB 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 != "TLOB" {
            return Err(MLError::CheckpointError(format!(
                "Invalid model type in checkpoint: expected TLOB, got {}",
                metadata.model_type
            )));
        }

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

        // Load weights from JSON sidecar
        let weights_path = format!("{}.weights.json", checkpoint_path);
        let weights_file = std::path::Path::new(&weights_path);
        if weights_file.exists() {
            let json = std::fs::read_to_string(weights_file).map_err(|e| {
                MLError::CheckpointError(format!("Failed to read weights: {}", e))
            })?;
            let all_weights: std::collections::HashMap<String, Vec<f32>> =
                serde_json::from_str(&json).map_err(|e| {
                    MLError::CheckpointError(format!("Failed to deserialize weights: {}", e))
                })?;

            // Convert HashMap to BTreeMap with shapes from current var_store
            let mut import_data = std::collections::BTreeMap::new();
            for (name, values) in &all_weights {
                let shape = self
                    .var_store
                    .get(name)
                    .map(|p| p.shape.clone())
                    .unwrap_or_else(|| vec![values.len()]);
                import_data.insert(name.clone(), (shape, values.clone()));
            }

            self.var_store.import_from_host(&import_data)?; // gpu-entry: checkpoint weight import
            tracing::info!(
                "Loaded TLOB weights from {} ({} parameters)",
                weights_path,
                all_weights.len()
            );
        } else {
            tracing::warn!(
                "No weights file at {}, using current (potentially Xavier-init) weights",
                weights_path
            );
        }

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

        Ok(metadata)
    }
}

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

    fn make_config() -> TLOBAdapterConfig {
        TLOBAdapterConfig {
            d_model: 32,
            num_heads: 2,
            num_layers: 2,
            seq_len: 32,
            feature_dim: 51,
        }
    }

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

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

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