//! Knowledge distillation module for model compression.

use forge_shared::{DistillLoss, DistillationConfig, ModelInfo};

/// Distiller for knowledge distillation.
#[derive(Debug)]
pub struct Distiller {
    /// RNG state.
    #[allow(dead_code)]
    rng_state: u64,
}

impl Default for Distiller {
    fn default() -> Self {
        Self::new()
    }
}

impl Distiller {
    /// Create a new distiller.
    pub fn new() -> Self {
        Self { rng_state: 42 }
    }

    /// Distill knowledge from teacher to student.
    pub fn distill(&mut self, teacher: &ModelInfo, config: &DistillationConfig) -> ModelInfo {
        // Create student model info
        let student_layers = if config.layer_mapping.is_empty() {
            teacher.num_layers / 2 // Default: half the layers
        } else {
            config.layer_mapping.len()
        };

        let student_hidden = teacher.hidden_dim;
        let student_params = self.estimate_student_params(teacher, student_layers);

        ModelInfo {
            name: format!("{}-distilled", teacher.name),
            architecture: teacher.architecture,
            num_parameters: student_params,
            num_layers: student_layers,
            hidden_dim: student_hidden,
            vocab_size: teacher.vocab_size,
            image_size: teacher.image_size,
            precision_bits: 16,
            size_bytes: student_params * 2, // FP16
        }
    }

    /// Estimate student model parameters.
    fn estimate_student_params(&self, teacher: &ModelInfo, student_layers: usize) -> u64 {
        let layer_ratio = student_layers as f64 / teacher.num_layers as f64;

        // Params scale roughly with number of layers (plus embeddings)
        let embedding_params =
            teacher.vocab_size.unwrap_or(32000) as u64 * teacher.hidden_dim as u64;
        let layer_params = (teacher.num_parameters - embedding_params * 2) as f64 * layer_ratio;

        (embedding_params * 2 + layer_params as u64).max(1_000_000)
    }

    /// Compute distillation loss.
    pub fn compute_loss(
        &self,
        teacher_logits: &[f64],
        student_logits: &[f64],
        config: &DistillationConfig,
    ) -> f64 {
        match config.loss_type {
            DistillLoss::KL => {
                self.kl_divergence(teacher_logits, student_logits, config.temperature)
            }
            DistillLoss::MSE => self.mse_loss(teacher_logits, student_logits),
            DistillLoss::Cosine => self.cosine_loss(teacher_logits, student_logits),
            DistillLoss::AttentionTransfer => {
                // Simplified: use MSE on attention patterns
                self.mse_loss(teacher_logits, student_logits) * 0.5
            }
            DistillLoss::FeatureMatching => {
                // Simplified: use MSE on features
                self.mse_loss(teacher_logits, student_logits) * 0.3
            }
            DistillLoss::Contrastive => {
                // Simplified contrastive loss
                self.cosine_loss(teacher_logits, student_logits) * 0.5
            }
            DistillLoss::Combined => {
                let kl = self.kl_divergence(teacher_logits, student_logits, config.temperature);
                let mse = self.mse_loss(teacher_logits, student_logits);
                config.alpha as f64 * kl + (1.0 - config.alpha as f64) * mse
            }
        }
    }

    /// KL divergence loss.
    fn kl_divergence(&self, teacher: &[f64], student: &[f64], temperature: f32) -> f64 {
        if teacher.is_empty() || student.is_empty() {
            return 0.0;
        }

        let t = temperature as f64;

        // Softmax with temperature
        let teacher_soft = self.softmax_with_temp(teacher, t);
        let student_soft = self.softmax_with_temp(student, t);

        // KL divergence
        let mut kl = 0.0;
        for (p, q) in teacher_soft.iter().zip(student_soft.iter()) {
            if *p > 1e-10 && *q > 1e-10 {
                kl += p * (p.ln() - q.ln());
            }
        }

        kl * t * t // Scale by T^2 as in original paper
    }

    /// Softmax with temperature.
    fn softmax_with_temp(&self, logits: &[f64], temperature: f64) -> Vec<f64> {
        if logits.is_empty() {
            return vec![];
        }

        let scaled: Vec<f64> = logits.iter().map(|&x| x / temperature).collect();
        let max = scaled.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        let exp_sum: f64 = scaled.iter().map(|&x| (x - max).exp()).sum();

        scaled.iter().map(|&x| (x - max).exp() / exp_sum).collect()
    }

    /// MSE loss.
    fn mse_loss(&self, teacher: &[f64], student: &[f64]) -> f64 {
        if teacher.is_empty() || student.is_empty() {
            return 0.0;
        }

        let n = teacher.len().min(student.len());
        let sum: f64 = teacher[..n]
            .iter()
            .zip(student[..n].iter())
            .map(|(t, s)| (t - s).powi(2))
            .sum();

        sum / n as f64
    }

    /// Cosine similarity loss.
    fn cosine_loss(&self, teacher: &[f64], student: &[f64]) -> f64 {
        if teacher.is_empty() || student.is_empty() {
            return 1.0;
        }

        let n = teacher.len().min(student.len());

        let dot: f64 = teacher[..n]
            .iter()
            .zip(student[..n].iter())
            .map(|(t, s)| t * s)
            .sum();

        let norm_t: f64 = teacher[..n].iter().map(|t| t * t).sum::<f64>().sqrt();
        let norm_s: f64 = student[..n].iter().map(|s| s * s).sum::<f64>().sqrt();

        if norm_t > 1e-10 && norm_s > 1e-10 {
            1.0 - dot / (norm_t * norm_s)
        } else {
            1.0
        }
    }

    /// Estimate accuracy degradation from distillation.
    pub fn estimate_degradation(&self, teacher: &ModelInfo, config: &DistillationConfig) -> f64 {
        let layer_ratio = if config.layer_mapping.is_empty() {
            0.5
        } else {
            config.layer_mapping.len() as f64 / teacher.num_layers as f64
        };

        // Base degradation from size reduction
        let base_degradation = (1.0 - layer_ratio) * 0.1;

        // Better distillation methods help
        let method_factor = match config.loss_type {
            DistillLoss::KL => 0.8,
            DistillLoss::Combined => 0.7,
            DistillLoss::AttentionTransfer => 0.6,
            _ => 1.0,
        };

        // More training helps
        let epoch_factor = 1.0 / (1.0 + config.epochs as f64 * 0.05);

        // Intermediate matching helps
        let intermediate_factor = if config.intermediate_matching {
            0.8
        } else {
            1.0
        };

        base_degradation * method_factor * epoch_factor * intermediate_factor
    }

    /// Create layer mapping for progressive distillation.
    pub fn create_layer_mapping(
        &self,
        teacher_layers: usize,
        student_layers: usize,
    ) -> Vec<(usize, usize)> {
        if student_layers >= teacher_layers {
            return (0..student_layers).map(|i| (i, i)).collect();
        }

        let step = teacher_layers as f64 / student_layers as f64;
        (0..student_layers)
            .map(|i| {
                let teacher_idx = (i as f64 * step) as usize;
                (teacher_idx.min(teacher_layers - 1), i)
            })
            .collect()
    }
}

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

    #[test]
    fn test_distiller_creation() {
        let distiller = Distiller::new();
        assert_eq!(distiller.rng_state, 42);
    }

    #[test]
    fn test_distill() {
        let mut distiller = Distiller::new();
        let teacher = forge_shared::sample_model_info();
        let config = DistillationConfig::default();

        let student = distiller.distill(&teacher, &config);

        assert!(student.num_layers < teacher.num_layers);
        assert!(student.num_parameters < teacher.num_parameters);
    }

    #[test]
    fn test_kl_divergence() {
        let distiller = Distiller::new();

        let teacher = vec![1.0, 2.0, 3.0];
        let student = vec![1.0, 2.0, 3.0];

        let loss = distiller.kl_divergence(&teacher, &student, 4.0);
        assert!(loss.abs() < 0.01); // Same distributions should have low KL
    }

    #[test]
    fn test_mse_loss() {
        let distiller = Distiller::new();

        let teacher = vec![1.0, 2.0, 3.0];
        let student = vec![1.0, 2.0, 3.0];

        let loss = distiller.mse_loss(&teacher, &student);
        assert_eq!(loss, 0.0);

        let student2 = vec![2.0, 3.0, 4.0];
        let loss2 = distiller.mse_loss(&teacher, &student2);
        assert_eq!(loss2, 1.0); // Each diff is 1, squared is 1, mean is 1
    }

    #[test]
    fn test_cosine_loss() {
        let distiller = Distiller::new();

        // Same direction
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        let loss = distiller.cosine_loss(&a, &b);
        assert!(loss.abs() < 0.01);

        // Orthogonal
        let c = vec![0.0, 1.0, 0.0];
        let loss2 = distiller.cosine_loss(&a, &c);
        assert!((loss2 - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_compute_loss_kl() {
        let distiller = Distiller::new();
        let config = DistillationConfig {
            loss_type: DistillLoss::KL,
            temperature: 4.0,
            ..Default::default()
        };

        let teacher = vec![1.0, 2.0, 3.0, 4.0];
        let student = vec![1.5, 2.5, 2.5, 4.0];

        let loss = distiller.compute_loss(&teacher, &student, &config);
        assert!(loss.is_finite());
        assert!(loss >= 0.0);
    }

    #[test]
    fn test_estimate_degradation() {
        let distiller = Distiller::new();
        let teacher = forge_shared::sample_model_info();

        let config_few_epochs = DistillationConfig {
            epochs: 1,
            ..Default::default()
        };
        let config_many_epochs = DistillationConfig {
            epochs: 100,
            ..Default::default()
        };

        let deg_few = distiller.estimate_degradation(&teacher, &config_few_epochs);
        let deg_many = distiller.estimate_degradation(&teacher, &config_many_epochs);

        // More epochs should reduce degradation
        assert!(deg_many < deg_few);
    }

    #[test]
    fn test_create_layer_mapping() {
        let distiller = Distiller::new();

        let mapping = distiller.create_layer_mapping(32, 6);
        assert_eq!(mapping.len(), 6);

        // First student layer should map to early teacher layer
        assert_eq!(mapping[0].1, 0);

        // Last student layer should map to late teacher layer
        assert!(mapping[5].0 > 20);
    }
}