//! Liquid CfC Trainer with gRPC Integration
//!
//! This module provides a production-ready Liquid CfC trainer that integrates with the ML
//! Training Service gRPC interface. It wraps the `LiquidTrainableAdapter` (which implements
//! `UnifiedTrainable`) to provide:
//! - GPU acceleration with automatic device selection
//! - Checkpoint management with safetensors serialization
//! - Per-epoch metrics reporting via gRPC progress callback
//! - Early stopping based on validation loss plateau

use ml_core::cuda_autograd::GpuTensor;
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};

use crate::liquid::adapter::LiquidTrainableAdapter;
use crate::liquid::candle_cfc::{CfCTrainConfig, DeviceConfig};
use crate::MLError;

/// Liquid CfC training hyperparameters (maps to gRPC LiquidParams)
#[derive(Debug, Clone)]
pub struct LiquidHyperparameters {
    /// Adam learning rate (typically 1e-4 to 1e-3)
    pub learning_rate: f64,
    /// Hidden state dimension of the CfC cell
    pub hidden_size: usize,
    /// Hidden dimensions for the backbone MLP layers
    pub backbone_hidden_sizes: Vec<usize>,
    /// Sequence length for temporal processing
    pub seq_len: usize,
    /// Minimum time constant for ODE dynamics
    pub tau_min: f64,
    /// Maximum time constant for ODE dynamics
    pub tau_max: f64,
    /// Dropout probability (0.0 to 0.5)
    pub dropout: f64,
    /// Maximum gradient norm for clipping
    pub gradient_clip: f64,
    /// Training batch size
    pub batch_size: usize,
    /// Number of training epochs
    pub epochs: usize,
    /// Input feature dimension
    pub input_size: usize,
    /// Output dimension (typically 3: buy/hold/sell)
    pub output_size: usize,
    /// Enable early stopping
    pub early_stopping_enabled: bool,
    /// Early stopping patience (epochs without improvement)
    pub early_stopping_patience: usize,
}

impl LiquidHyperparameters {
    /// Create conservative hyperparameters suitable for testing and development.
    /// WARNING: These are NOT optimized for production. Use hyperopt results instead.
    pub fn conservative() -> Self {
        Self {
            learning_rate: 0.001,
            hidden_size: 64,
            backbone_hidden_sizes: vec![64],
            seq_len: 32,
            tau_min: 0.01,
            tau_max: 1.0,
            dropout: 0.1,
            gradient_clip: 1.0,
            batch_size: 64,
            epochs: 100,
            input_size: 51,
            output_size: 3,
            early_stopping_enabled: true,
            early_stopping_patience: 10,
        }
    }
}

impl From<LiquidHyperparameters> for CfCTrainConfig {
    fn from(params: LiquidHyperparameters) -> Self {
        CfCTrainConfig {
            input_size: params.input_size,
            hidden_size: params.hidden_size,
            output_size: params.output_size,
            backbone_hidden_sizes: params.backbone_hidden_sizes,
            tau_min: params.tau_min,
            tau_max: params.tau_max,
            learning_rate: params.learning_rate,
            batch_size: params.batch_size,
            seq_len: params.seq_len,
            max_epochs: params.epochs,
            early_stopping_patience: params.early_stopping_patience,
            dropout: params.dropout,
            gradient_clip: params.gradient_clip,
            market_regime_adaptation: true,
            device: DeviceConfig::Auto,
        }
    }
}

/// Training progress metrics (returned to gRPC service)
#[derive(Debug, Clone)]
pub struct LiquidTrainingMetrics {
    pub epoch: usize,
    pub train_loss: f64,
    pub val_loss: Option<f64>,
    pub grad_norm: f64,
    pub learning_rate: f64,
}

/// Liquid CfC Trainer for gRPC integration
pub struct LiquidTrainer {
    adapter: LiquidTrainableAdapter,
    hyperparams: LiquidHyperparameters,
    checkpoint_dir: PathBuf,
    /// Best validation loss seen so far (for early stopping)
    pub best_val_loss: f64,
    /// Epochs without improvement counter
    patience_counter: usize,
}

impl std::fmt::Debug for LiquidTrainer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LiquidTrainer")
            .field("hyperparams", &self.hyperparams)
            .field("checkpoint_dir", &self.checkpoint_dir)
            .field("best_val_loss", &self.best_val_loss)
            .field("patience_counter", &self.patience_counter)
            .finish()
    }
}

impl LiquidTrainer {
    /// Create new Liquid CfC trainer
    ///
    /// # Arguments
    /// * `hyperparams` - Training hyperparameters from gRPC request
    /// * `checkpoint_dir` - Directory for saving model checkpoints
    pub fn new<P: AsRef<Path>>(
        hyperparams: LiquidHyperparameters,
        checkpoint_dir: P,
    ) -> Result<Self, MLError> {
        info!(
            "Initializing Liquid CfC trainer with hidden_size={}, seq_len={}, backbone={:?}",
            hyperparams.hidden_size, hyperparams.seq_len, hyperparams.backbone_hidden_sizes
        );

        let config: CfCTrainConfig = hyperparams.clone().into();
        let adapter = LiquidTrainableAdapter::new(config)?;

        let device = adapter.device();
        info!("Liquid CfC trainer using device: {:?}", device);

        Ok(Self {
            adapter,
            hyperparams,
            checkpoint_dir: checkpoint_dir.as_ref().to_path_buf(),
            best_val_loss: f64::MAX,
            patience_counter: 0,
        })
    }

    /// Train Liquid CfC model on tensor data
    ///
    /// # Arguments
    /// * `train_data` - Training data as `(input, target)` tensor pairs
    ///   - input shape: `[batch, seq_len, input_size]`
    ///   - target shape: `[batch, output_size]`
    /// * `val_data` - Validation data (same format), used for early stopping
    /// * `progress_callback` - Callback for reporting progress to gRPC stream
    ///
    /// # Returns
    /// Training history as a vector of per-epoch metrics
    pub fn train<F>(
        &mut self,
        train_data: &[(GpuTensor, GpuTensor)],
        val_data: &[(GpuTensor, GpuTensor)],
        mut progress_callback: F,
    ) -> Result<Vec<LiquidTrainingMetrics>, MLError>
    where
        F: FnMut(LiquidTrainingMetrics),
    {
        info!(
            "Starting Liquid CfC training for {} epochs ({} train batches, {} val batches)",
            self.hyperparams.epochs,
            train_data.len(),
            val_data.len()
        );

        if train_data.is_empty() {
            return Err(MLError::ValidationError {
                message: "Training data is empty".to_owned(),
            });
        }

        let mut history = Vec::with_capacity(self.hyperparams.epochs);

        for epoch in 0..self.hyperparams.epochs {
            debug!(
                "Training epoch {}/{}",
                epoch + 1,
                self.hyperparams.epochs
            );

            // Train one epoch
            let (train_loss, grad_norm) = self.train_epoch(train_data)?;

            // Compute validation loss
            let val_loss = (!val_data.is_empty())
                .then(|| self.evaluate(val_data))
                .transpose()?;

            // Build metrics
            let metrics = LiquidTrainingMetrics {
                epoch: epoch + 1,
                train_loss,
                val_loss,
                grad_norm,
                learning_rate: self.hyperparams.learning_rate,
            };

            progress_callback(metrics.clone());
            history.push(metrics);

            // Early stopping check
            if self.hyperparams.early_stopping_enabled {
                let loss_to_check = val_loss.unwrap_or(train_loss);
                if loss_to_check < self.best_val_loss {
                    self.best_val_loss = loss_to_check;
                    self.patience_counter = 0;

                    // Save best checkpoint
                    if let Err(e) = self.save_checkpoint(epoch + 1) {
                        warn!("Failed to save best checkpoint: {}", e);
                    }
                } else {
                    self.patience_counter += 1;
                    if self.patience_counter >= self.hyperparams.early_stopping_patience {
                        warn!(
                            "Early stopping triggered at epoch {}/{}: {} epochs without improvement",
                            epoch + 1,
                            self.hyperparams.epochs,
                            self.patience_counter
                        );
                        break;
                    }
                }
            }

            // Save checkpoint every 10 epochs
            if (epoch + 1) % 10 == 0 {
                if let Err(e) = self.save_checkpoint(epoch + 1) {
                    warn!("Failed to save epoch {} checkpoint: {}", epoch + 1, e);
                }
            }

            info!(
                "Epoch {}/{}: train_loss={:.6}, val_loss={}, grad_norm={:.4}",
                epoch + 1,
                self.hyperparams.epochs,
                train_loss,
                val_loss.map_or("N/A".to_owned(), |v| format!("{:.6}", v)),
                grad_norm
            );
        }

        // Save final checkpoint
        let final_epoch = history.len();
        if let Err(e) = self.save_checkpoint(final_epoch) {
            warn!("Failed to save final checkpoint: {}", e);
        }

        info!(
            "Liquid CfC training complete: {} epochs, best_val_loss={:.6}",
            history.len(),
            self.best_val_loss
        );
        Ok(history)
    }

    /// Train one epoch over the given data, returning (avg_loss, grad_norm)
    ///
    /// Loss is accumulated on-GPU as an F32 tensor to avoid per-batch GPU→CPU
    /// synchronisation. A single `.to_scalar()` extraction happens after the loop.
    /// A NaN guard fires every 100 batches to detect divergence early.
    fn train_epoch(&mut self, data: &[(GpuTensor, GpuTensor)]) -> Result<(f64, f64), MLError> {
        let mut total_loss = 0.0_f64;
        let mut total_grad_norm = 0.0_f64;
        let mut batch_count: usize = 0;

        for (input, target) in data {
            let output = self.adapter.forward(input)?;
            let loss = self.adapter.compute_loss(&output, target)?;
            let grad_norm = self.adapter.backward(&loss)?;
            self.adapter.optimizer_step()?;

            total_loss += loss as f64;
            total_grad_norm += grad_norm;
            batch_count += 1;

            // NaN guard every 100 batches
            if batch_count % 100 == 0 && !total_loss.is_finite() {
                return Err(MLError::TrainingError(format!(
                    "Non-finite accumulated loss ({total_loss}) at batch {batch_count}"
                )));
            }
        }

        if batch_count == 0 {
            return Ok((0.0, 0.0));
        }

        let avg_loss = total_loss / batch_count as f64;
        let avg_grad_norm = total_grad_norm / batch_count as f64;

        Ok((avg_loss, avg_grad_norm))
    }

    /// Evaluate model on validation data without updating weights.
    ///
    /// Loss is accumulated on-GPU as a detached F32 tensor — a single
    /// `.to_scalar()` extraction happens after the loop, eliminating
    /// per-batch GPU→CPU synchronisation.
    fn evaluate(&mut self, data: &[(GpuTensor, GpuTensor)]) -> Result<f64, MLError> {
        let mut total_loss = 0.0_f64;
        let mut batch_count: usize = 0;

        for (input, target) in data {
            let output = self.adapter.forward(input)?;
            let loss = self.adapter.compute_loss(&output, target)?;
            total_loss += loss as f64;
            batch_count += 1;
        }

        if batch_count == 0 {
            return Ok(0.0);
        }
        Ok(total_loss / batch_count as f64)
    }

    /// Save model checkpoint to safetensors format
    fn save_checkpoint(&self, epoch: usize) -> Result<(), MLError> {
        let checkpoint_path = self
            .checkpoint_dir
            .join(format!("liquid_cfc_epoch_{}.safetensors", epoch));

        info!("Saving Liquid CfC checkpoint to {:?}", checkpoint_path);

        // Create checkpoint directory if it doesn't exist
        if let Some(parent) = checkpoint_path.parent() {
            std::fs::create_dir_all(parent).map_err(|e| MLError::ConfigError(format!("Failed to create checkpoint directory: {}", e)))?;
        }

        // Save VarStore weights to checkpoint via safetensors
        ml_core::checkpoint::save_safetensors(
            self.adapter.varmap(),
            &checkpoint_path,
            self.adapter.varmap().stream(),
            None,
        ).map_err(|e| MLError::ConfigError(format!("Failed to save Liquid CfC checkpoint: {}", e)))?;

        let file_size = std::fs::metadata(&checkpoint_path)
            .map(|m| m.len() / 1024)
            .unwrap_or(0);

        info!(
            "Checkpoint saved: {} ({} KB)",
            checkpoint_path.display(),
            file_size
        );

        Ok(())
    }

    /// Get current hyperparameters
    pub fn hyperparameters(&self) -> &LiquidHyperparameters {
        &self.hyperparams
    }

    /// Get the underlying adapter for direct access
    pub fn adapter(&self) -> &LiquidTrainableAdapter {
        &self.adapter
    }
}

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

    fn create_test_params() -> LiquidHyperparameters {
        LiquidHyperparameters {
            learning_rate: 0.01,
            hidden_size: 16,
            backbone_hidden_sizes: vec![16],
            seq_len: 4,
            tau_min: 0.01,
            tau_max: 1.0,
            dropout: 0.0,
            gradient_clip: 1.0,
            batch_size: 4,
            epochs: 5,
            input_size: 8,
            output_size: 3,
            early_stopping_enabled: false,
            early_stopping_patience: 3,
        }
    }

    #[test]
    fn test_liquid_hyperparameters_conservative() {
        let params = LiquidHyperparameters::conservative();
        assert_eq!(params.hidden_size, 64);
        assert_eq!(params.output_size, 3);
        assert_eq!(params.input_size, 51);
        assert!(params.early_stopping_enabled);
    }

    #[test]
    fn test_liquid_config_conversion() {
        let params = create_test_params();
        let config: CfCTrainConfig = params.into();
        assert_eq!(config.hidden_size, 16);
        assert_eq!(config.seq_len, 4);
        assert_eq!(config.output_size, 3);
        assert_eq!(config.input_size, 8);
    }

    #[test]
    fn test_liquid_trainer_creation() {
        let params = create_test_params();
        let trainer = LiquidTrainer::new(params, "/tmp/liquid_cfc_checkpoints");
        assert!(trainer.is_ok());
    }

    #[test]
    fn test_liquid_trainer_train() {
        let params = create_test_params();
        let mut trainer = LiquidTrainer::new(params, "/tmp/liquid_cfc_checkpoints")
            .expect("Trainer creation failed");

        // Create synthetic training data using GpuTensor
        let stream = trainer.adapter().stream().clone();
        let input = GpuTensor::randn(&[4, 4, 8], 0.5, &stream)
            .expect("input tensor");
        let target = GpuTensor::zeros(&[4, 3], &stream)
            .expect("target tensor");

        let train_data = vec![(input.clone(), target.clone())];
        let val_data = vec![(input, target)];

        let mut metrics_log = Vec::new();
        let result = trainer.train(&train_data, &val_data, |m| {
            metrics_log.push(m);
        });

        assert!(result.is_ok(), "Training failed: {:?}", result.err());
        let history = result.expect("already checked");
        assert_eq!(history.len(), 5);
        assert_eq!(metrics_log.len(), 5);

        // Verify metrics are populated
        let first = metrics_log.first().expect("at least one metric");
        assert_eq!(first.epoch, 1);
        assert!(first.train_loss > 0.0);
        assert!(first.val_loss.is_some());
    }

    #[test]
    fn test_liquid_trainer_early_stopping() {
        let mut params = create_test_params();
        params.epochs = 200;
        params.early_stopping_enabled = true;
        params.early_stopping_patience = 3;

        let mut trainer = LiquidTrainer::new(params, "/tmp/liquid_cfc_checkpoints")
            .expect("Trainer creation failed");

        let stream = trainer.adapter().stream().clone();
        // Deterministic data: train on zeros→zeros, validate on fixed noise→large targets.
        // Model learns zero-mapping for training, but val targets are far from zero
        // so val loss cannot improve → early stopping fires.
        let train_input = GpuTensor::zeros(&[4, 4, 8], &stream)
            .expect("train input tensor");
        let train_target = GpuTensor::zeros(&[4, 3], &stream)
            .expect("train target tensor");
        let val_input = GpuTensor::zeros(&[4, 4, 8], &stream)
            .expect("val input tensor");
        // Large val targets that the model can never reach after learning zero-mapping
        let val_target = GpuTensor::full(&[4, 3], 100.0, &stream)
            .expect("val target tensor");

        let train_data = vec![(train_input, train_target)];
        let val_data = vec![(val_input, val_target)];

        let result = trainer.train(&train_data, &val_data, |_| {});
        assert!(result.is_ok());
        let history = result.expect("already checked");

        // Model converges on training (zeros→zeros) but val loss stays high
        // (zeros→100 is impossible). Early stopping triggers within ~10 epochs.
        assert!(
            history.len() < 200,
            "Expected early stopping, got {} epochs",
            history.len()
        );
    }

    #[test]
    fn test_liquid_trainer_empty_data_error() {
        let params = create_test_params();
        let mut trainer = LiquidTrainer::new(params, "/tmp/liquid_cfc_checkpoints")
            .expect("Trainer creation failed");

        let result = trainer.train(&[], &[], |_| {});
        assert!(result.is_err(), "Should reject empty training data");
    }
}
