//! Sample data and configurations for FoundationForge demo.

use forge_shared::{
    BenchmarkConfig, CompressionStep, DistillLoss, DistillationConfig, ExportConfig, ExportFormat,
    ModelArchitecture, ModelInfo, PipelineConfig, PruneMethod, PruneSchedule, PruningConfig,
    QuantMethod, QuantPrecision, QuantizationConfig, StudentConfig, TargetDevice,
};

// ============================================================================
// Model Configurations
// ============================================================================

/// Create a sample LLaMA 7B model info.
#[must_use]
pub fn llama_7b() -> ModelInfo {
    ModelInfo {
        name: "llama-7b".to_string(),
        architecture: ModelArchitecture::Transformer,
        num_parameters: 7_000_000_000,
        num_layers: 32,
        hidden_dim: 4096,
        vocab_size: Some(32000),
        image_size: None,
        precision_bits: 16,
        size_bytes: 14_000_000_000,
    }
}

/// Create a sample LLaMA 13B model info.
#[must_use]
pub fn llama_13b() -> ModelInfo {
    ModelInfo {
        name: "llama-13b".to_string(),
        architecture: ModelArchitecture::Transformer,
        num_parameters: 13_000_000_000,
        num_layers: 40,
        hidden_dim: 5120,
        vocab_size: Some(32000),
        image_size: None,
        precision_bits: 16,
        size_bytes: 26_000_000_000,
    }
}

/// Create a sample Mistral 7B model info.
#[must_use]
pub fn mistral_7b() -> ModelInfo {
    ModelInfo {
        name: "mistral-7b".to_string(),
        architecture: ModelArchitecture::Transformer,
        num_parameters: 7_300_000_000,
        num_layers: 32,
        hidden_dim: 4096,
        vocab_size: Some(32000),
        image_size: None,
        precision_bits: 16,
        size_bytes: 14_600_000_000,
    }
}

/// Create a sample ViT-Large model info.
#[must_use]
pub fn vit_large() -> ModelInfo {
    ModelInfo {
        name: "vit-large".to_string(),
        architecture: ModelArchitecture::ViT,
        num_parameters: 307_000_000,
        num_layers: 24,
        hidden_dim: 1024,
        vocab_size: None,
        image_size: Some((224, 224)),
        precision_bits: 32,
        size_bytes: 1_228_000_000,
    }
}

/// Create a sample BERT-Base model info.
#[must_use]
pub fn bert_base() -> ModelInfo {
    ModelInfo {
        name: "bert-base".to_string(),
        architecture: ModelArchitecture::Transformer,
        num_parameters: 110_000_000,
        num_layers: 12,
        hidden_dim: 768,
        vocab_size: Some(30522),
        image_size: None,
        precision_bits: 32,
        size_bytes: 440_000_000,
    }
}

// ============================================================================
// Quantization Configurations
// ============================================================================

/// Create INT4 GPTQ quantization config.
#[must_use]
pub fn gptq_int4_config() -> QuantizationConfig {
    QuantizationConfig {
        method: QuantMethod::GPTQ,
        precision: QuantPrecision::Int4,
        group_size: Some(128),
        quantize_activations: false,
        quantize_embeddings: false,
        calibration_size: 512,
        symmetric: false,
        per_channel: true,
        skip_layers: vec![],
    }
}

/// Create INT8 AWQ quantization config.
#[must_use]
pub fn awq_int8_config() -> QuantizationConfig {
    QuantizationConfig {
        method: QuantMethod::AWQ,
        precision: QuantPrecision::Int8,
        group_size: None,
        quantize_activations: true,
        quantize_embeddings: false,
        calibration_size: 256,
        symmetric: true,
        per_channel: true,
        skip_layers: vec![],
    }
}

/// Create GGML quantization config.
#[must_use]
pub fn ggml_q4_k_m_config() -> QuantizationConfig {
    QuantizationConfig {
        method: QuantMethod::GGML,
        precision: QuantPrecision::Int4,
        group_size: Some(32),
        quantize_activations: false,
        quantize_embeddings: true,
        calibration_size: 0,
        symmetric: false,
        per_channel: true,
        skip_layers: vec![],
    }
}

/// Create dynamic quantization config.
#[must_use]
pub fn dynamic_int8_config() -> QuantizationConfig {
    QuantizationConfig {
        method: QuantMethod::Dynamic,
        precision: QuantPrecision::Int8,
        group_size: None,
        quantize_activations: true,
        quantize_embeddings: false,
        calibration_size: 0,
        symmetric: true,
        per_channel: false,
        skip_layers: vec!["lm_head".to_string()],
    }
}

// ============================================================================
// Pruning Configurations
// ============================================================================

/// Create 50% magnitude pruning config.
#[must_use]
pub fn magnitude_50_config() -> PruningConfig {
    PruningConfig {
        method: PruneMethod::Magnitude,
        target_sparsity: 0.5,
        schedule: PruneSchedule::Gradual,
        num_steps: 10,
        initial_sparsity: 0.0,
        final_sparsity: 0.5,
        retrain_epochs: 3,
        n: None,
        m: None,
        exclude_layers: vec!["embedding".to_string(), "lm_head".to_string()],
    }
}

/// Create 2:4 structured sparsity config.
#[must_use]
pub fn nm_2_4_config() -> PruningConfig {
    PruningConfig {
        method: PruneMethod::NM,
        target_sparsity: 0.5,
        schedule: PruneSchedule::OneShot,
        num_steps: 1,
        initial_sparsity: 0.5,
        final_sparsity: 0.5,
        retrain_epochs: 5,
        n: Some(2),
        m: Some(4),
        exclude_layers: vec![],
    }
}

/// Create aggressive 90% pruning config.
#[must_use]
pub fn aggressive_prune_config() -> PruningConfig {
    PruningConfig {
        method: PruneMethod::Movement,
        target_sparsity: 0.9,
        schedule: PruneSchedule::Cubic,
        num_steps: 20,
        initial_sparsity: 0.3,
        final_sparsity: 0.9,
        retrain_epochs: 10,
        n: None,
        m: None,
        exclude_layers: vec!["embedding".to_string()],
    }
}

// ============================================================================
// Distillation Configurations
// ============================================================================

/// Create standard KL distillation config.
#[must_use]
pub fn kl_distill_config() -> DistillationConfig {
    DistillationConfig {
        loss_type: DistillLoss::KL,
        temperature: 4.0,
        alpha: 0.5,
        intermediate_matching: true,
        layer_mapping: vec![(0, 0), (10, 1), (21, 2), (31, 3)],
        epochs: 10,
        batch_size: 32,
        learning_rate: 5e-5,
        progressive: false,
    }
}

/// Create TinyBERT-style distillation config.
#[must_use]
pub fn tinybert_distill_config() -> DistillationConfig {
    DistillationConfig {
        loss_type: DistillLoss::Combined,
        temperature: 1.0,
        alpha: 0.7,
        intermediate_matching: true,
        layer_mapping: vec![(0, 0), (3, 1), (6, 2), (9, 3), (11, 4), (11, 5)],
        epochs: 20,
        batch_size: 64,
        learning_rate: 3e-5,
        progressive: true,
    }
}

/// Create attention transfer distillation config.
#[must_use]
pub fn attention_distill_config() -> DistillationConfig {
    DistillationConfig {
        loss_type: DistillLoss::AttentionTransfer,
        temperature: 2.0,
        alpha: 0.3,
        intermediate_matching: true,
        layer_mapping: vec![],
        epochs: 15,
        batch_size: 16,
        learning_rate: 1e-4,
        progressive: false,
    }
}

// ============================================================================
// Student Model Configurations
// ============================================================================

/// Create a 6-layer student config.
#[must_use]
pub fn student_6_layer() -> StudentConfig {
    StudentConfig {
        num_layers: 6,
        hidden_dim: 768,
        num_heads: 12,
        intermediate_dim: 3072,
        init_from_teacher: true,
        copy_layers: vec![0, 2, 5, 8, 10, 11],
    }
}

/// Create a 4-layer student config.
#[must_use]
pub fn student_4_layer() -> StudentConfig {
    StudentConfig {
        num_layers: 4,
        hidden_dim: 512,
        num_heads: 8,
        intermediate_dim: 2048,
        init_from_teacher: true,
        copy_layers: vec![0, 4, 8, 11],
    }
}

// ============================================================================
// Pipeline Configurations
// ============================================================================

/// Create a mobile deployment pipeline.
#[must_use]
pub fn mobile_pipeline() -> PipelineConfig {
    PipelineConfig {
        name: "mobile-optimized".to_string(),
        steps: vec![
            CompressionStep::Quantize(gptq_int4_config()),
            CompressionStep::PruneHeads { target_heads: 8 },
        ],
        eval_dataset: Some("wikitext".to_string()),
        target_size: Some(2_000_000_000),
        max_accuracy_loss: 0.05,
        export: ExportConfig {
            format: ExportFormat::CoreML,
            optimize: true,
            include_tokenizer: true,
            target_device: TargetDevice::Mobile,
            output_path: "output/mobile_model".to_string(),
        },
    }
}

/// Create a server deployment pipeline.
#[must_use]
pub fn server_pipeline() -> PipelineConfig {
    PipelineConfig {
        name: "server-optimized".to_string(),
        steps: vec![CompressionStep::Quantize(awq_int8_config())],
        eval_dataset: Some("wikitext".to_string()),
        target_size: None,
        max_accuracy_loss: 0.01,
        export: ExportConfig {
            format: ExportFormat::SafeTensors,
            optimize: true,
            include_tokenizer: true,
            target_device: TargetDevice::CUDA,
            output_path: "output/server_model".to_string(),
        },
    }
}

/// Create an edge deployment pipeline.
#[must_use]
pub fn edge_pipeline() -> PipelineConfig {
    PipelineConfig {
        name: "edge-optimized".to_string(),
        steps: vec![
            CompressionStep::Prune(nm_2_4_config()),
            CompressionStep::Quantize(gptq_int4_config()),
            CompressionStep::PruneVocab { keep_tokens: 16000 },
        ],
        eval_dataset: Some("wikitext".to_string()),
        target_size: Some(1_000_000_000),
        max_accuracy_loss: 0.08,
        export: ExportConfig {
            format: ExportFormat::TFLite,
            optimize: true,
            include_tokenizer: true,
            target_device: TargetDevice::Edge,
            output_path: "output/edge_model".to_string(),
        },
    }
}

/// Create a maximum compression pipeline.
#[must_use]
pub fn max_compression_pipeline() -> PipelineConfig {
    PipelineConfig {
        name: "max-compression".to_string(),
        steps: vec![
            CompressionStep::Distill(kl_distill_config()),
            CompressionStep::Prune(aggressive_prune_config()),
            CompressionStep::Quantize(ggml_q4_k_m_config()),
        ],
        eval_dataset: Some("wikitext".to_string()),
        target_size: Some(500_000_000),
        max_accuracy_loss: 0.15,
        export: ExportConfig {
            format: ExportFormat::GGUF,
            optimize: true,
            include_tokenizer: true,
            target_device: TargetDevice::CPU,
            output_path: "output/tiny_model".to_string(),
        },
    }
}

// ============================================================================
// Benchmark Configurations
// ============================================================================

/// Create a quick benchmark config.
#[must_use]
pub fn quick_benchmark() -> BenchmarkConfig {
    BenchmarkConfig {
        dataset: "wikitext".to_string(),
        num_samples: 100,
        batch_size: 1,
        measure_latency: true,
        measure_memory: true,
        warmup_runs: 2,
        benchmark_runs: 5,
    }
}

/// Create a full benchmark config.
#[must_use]
pub fn full_benchmark() -> BenchmarkConfig {
    BenchmarkConfig {
        dataset: "wikitext".to_string(),
        num_samples: 1000,
        batch_size: 1,
        measure_latency: true,
        measure_memory: true,
        warmup_runs: 5,
        benchmark_runs: 20,
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_model_infos() {
        let llama7 = llama_7b();
        let llama13 = llama_13b();

        assert!(llama13.num_parameters > llama7.num_parameters);
        assert!(llama13.num_layers > llama7.num_layers);
    }

    #[test]
    fn test_quant_configs() {
        let gptq = gptq_int4_config();
        let awq = awq_int8_config();

        assert_eq!(gptq.precision, QuantPrecision::Int4);
        assert_eq!(awq.precision, QuantPrecision::Int8);
    }

    #[test]
    fn test_prune_configs() {
        let mag = magnitude_50_config();
        let nm = nm_2_4_config();

        assert_eq!(mag.target_sparsity, 0.5);
        assert_eq!(nm.n, Some(2));
        assert_eq!(nm.m, Some(4));
    }

    #[test]
    fn test_distill_configs() {
        let kl = kl_distill_config();
        let tinybert = tinybert_distill_config();

        assert_eq!(kl.loss_type, DistillLoss::KL);
        assert_eq!(tinybert.loss_type, DistillLoss::Combined);
    }

    #[test]
    fn test_pipelines() {
        let mobile = mobile_pipeline();
        let edge = edge_pipeline();

        assert!(mobile.target_size.unwrap() > edge.target_size.unwrap());
        assert_eq!(mobile.export.target_device, TargetDevice::Mobile);
    }

    #[test]
    fn test_benchmarks() {
        let quick = quick_benchmark();
        let full = full_benchmark();

        assert!(quick.num_samples < full.num_samples);
        assert!(quick.benchmark_runs < full.benchmark_runs);
    }
}