//! Pre-defined processor profiles for common hardware.

use super::capabilities::{
    ComputeThroughput, MemorySpecs, MemoryType, PowerCharacteristics, ProcessorCapabilities,
    WorkloadCharacteristic,
};
use super::operation::OperationType;
use super::types::PowerTier;
use super::TpuVersion;
use std::collections::HashSet;

/// Pre-defined processor profiles.
pub struct ProcessorProfiles;

impl ProcessorProfiles {
    // ═══════════════════════════════════════════════════════════════
    // CPU PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// Default CPU profile.
    pub fn cpu_default() -> ProcessorCapabilities {
        ProcessorCapabilities::cpu(8, 3.5, false)
    }

    /// AMD EPYC 9654 (96 cores).
    pub fn amd_epyc_9654() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 2.7,
                fp32_tflops: 5.4,
                fp16_tflops: 10.8,
                bf16_tflops: 10.8,
                int8_tops: 21.6,
                int4_tops: 43.2,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 6 * 1024 * 1024 * 1024 * 1024, // 6 TB max
                bandwidth_gbps: 460,
                type_: MemoryType::Ddr5,
            },
            operations: ProcessorCapabilities::cpu(96, 2.4, false)
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 360,
                efficiency: 0.85,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::Sequential,
                WorkloadCharacteristic::MemoryBound,
            ],
        }
    }

    /// Intel Xeon w9-3595X (56 cores).
    pub fn intel_xeon_w9_3595x() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 3.2,
                fp32_tflops: 6.4,
                fp16_tflops: 12.8,
                bf16_tflops: 12.8,
                int8_tops: 25.6,
                int4_tops: 51.2,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 4 * 1024 * 1024 * 1024 * 1024, // 4 TB max
                bandwidth_gbps: 307,
                type_: MemoryType::Ddr5,
            },
            operations: ProcessorCapabilities::cpu(56, 2.9, true)
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 350,
                efficiency: 0.80,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::Sequential,
                WorkloadCharacteristic::MemoryBound,
            ],
        }
    }

    /// Apple M3 Max CPU cores.
    pub fn apple_m3_max_cpu() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 0.3,
                fp32_tflops: 0.6,
                fp16_tflops: 1.2,
                bf16_tflops: 1.2,
                int8_tops: 2.4,
                int4_tops: 4.8,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 128 * 1024 * 1024 * 1024, // 128 GB unified
                bandwidth_gbps: 400,
                type_: MemoryType::Unified,
            },
            operations: ProcessorCapabilities::cpu(16, 4.0, false)
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 40,
                efficiency: 0.95,
                power_tier: PowerTier::Low,
            },
            optimal_for: vec![
                WorkloadCharacteristic::Sequential,
                WorkloadCharacteristic::LowPower,
            ],
        }
    }

    // ═══════════════════════════════════════════════════════════════
    // NVIDIA GPU PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// Default NVIDIA GPU profile.
    pub fn nvidia_default() -> ProcessorCapabilities {
        ProcessorCapabilities::nvidia_gpu(8192, 256, 12, 600, (8, 0))
    }

    /// NVIDIA H100 SXM (80GB).
    pub fn nvidia_h100() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 67.0,
                fp32_tflops: 67.0,
                fp16_tflops: 1979.0, // With sparsity
                bf16_tflops: 1979.0,
                int8_tops: 3958.0,
                int4_tops: 7916.0,
                sparsity_speedup: 2.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 80 * 1024 * 1024 * 1024,
                bandwidth_gbps: 3350,
                type_: MemoryType::Hbm3,
            },
            operations: ProcessorCapabilities::nvidia_gpu(16896, 528, 80, 3350, (9, 0))
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 700,
                efficiency: 0.90,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
                WorkloadCharacteristic::LargeBatch,
                WorkloadCharacteristic::ComputeBound,
            ],
        }
    }

    /// NVIDIA A100 (80GB).
    pub fn nvidia_a100() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 19.5,
                fp32_tflops: 19.5,
                fp16_tflops: 624.0, // With sparsity
                bf16_tflops: 624.0,
                int8_tops: 1248.0,
                int4_tops: 2496.0,
                sparsity_speedup: 2.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 80 * 1024 * 1024 * 1024,
                bandwidth_gbps: 2039,
                type_: MemoryType::Hbm2e,
            },
            operations: ProcessorCapabilities::nvidia_gpu(6912, 432, 80, 2039, (8, 0))
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 400,
                efficiency: 0.88,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
                WorkloadCharacteristic::LargeBatch,
                WorkloadCharacteristic::ComputeBound,
            ],
        }
    }

    /// NVIDIA RTX 4090.
    pub fn nvidia_rtx_4090() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 1.3,
                fp32_tflops: 82.6,
                fp16_tflops: 330.4, // With sparsity
                bf16_tflops: 330.4,
                int8_tops: 660.8,
                int4_tops: 1321.6,
                sparsity_speedup: 2.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 24 * 1024 * 1024 * 1024,
                bandwidth_gbps: 1008,
                type_: MemoryType::Gddr6,
            },
            operations: ProcessorCapabilities::nvidia_gpu(16384, 512, 24, 1008, (8, 9))
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 450,
                efficiency: 0.85,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
                WorkloadCharacteristic::LargeBatch,
            ],
        }
    }

    /// NVIDIA RTX 3090.
    pub fn nvidia_rtx_3090() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 0.6,
                fp32_tflops: 35.6,
                fp16_tflops: 71.2,
                bf16_tflops: 71.2,
                int8_tops: 142.4,
                int4_tops: 284.8,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 24 * 1024 * 1024 * 1024,
                bandwidth_gbps: 936,
                type_: MemoryType::Gddr6,
            },
            operations: ProcessorCapabilities::nvidia_gpu(10496, 328, 24, 936, (8, 6))
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 350,
                efficiency: 0.82,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
                WorkloadCharacteristic::LargeBatch,
            ],
        }
    }

    // ═══════════════════════════════════════════════════════════════
    // AMD GPU PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// AMD MI300X.
    pub fn amd_mi300x() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 163.4,
                fp32_tflops: 163.4,
                fp16_tflops: 1307.0,
                bf16_tflops: 1307.0,
                int8_tops: 2614.0,
                int4_tops: 5228.0,
                sparsity_speedup: 2.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 192 * 1024 * 1024 * 1024, // 192 GB HBM3
                bandwidth_gbps: 5300,
                type_: MemoryType::Hbm3,
            },
            operations: {
                let mut ops = ProcessorCapabilities::nvidia_gpu(16384, 512, 80, 5300, (9, 0))
                    .operations;
                ops.remove(&OperationType::FlashAttention); // Different implementation
                ops
            },
            power: PowerCharacteristics {
                tdp_watts: 750,
                efficiency: 0.88,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
                WorkloadCharacteristic::LargeBatch,
                WorkloadCharacteristic::MemoryBound, // High memory bandwidth
            ],
        }
    }

    /// AMD RX 7900 XTX.
    pub fn amd_rx_7900_xtx() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 1.9,
                fp32_tflops: 61.0,
                fp16_tflops: 122.0,
                bf16_tflops: 122.0,
                int8_tops: 244.0,
                int4_tops: 488.0,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 24 * 1024 * 1024 * 1024,
                bandwidth_gbps: 960,
                type_: MemoryType::Gddr6,
            },
            operations: {
                let mut ops = ProcessorCapabilities::nvidia_gpu(6144, 0, 24, 960, (8, 0))
                    .operations;
                ops.remove(&OperationType::FlashAttention);
                ops
            },
            power: PowerCharacteristics {
                tdp_watts: 355,
                efficiency: 0.80,
                power_tier: PowerTier::High,
            },
            optimal_for: vec![
                WorkloadCharacteristic::HighlyParallel,
            ],
        }
    }

    // ═══════════════════════════════════════════════════════════════
    // GOOGLE TPU PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// Default TPU profile.
    pub fn google_tpu_default() -> ProcessorCapabilities {
        ProcessorCapabilities::tpu(TpuVersion::V4)
    }

    /// Google TPU v5p.
    pub fn google_tpu_v5p() -> ProcessorCapabilities {
        ProcessorCapabilities::tpu(TpuVersion::V5p)
    }

    /// Google TPU v4.
    pub fn google_tpu_v4() -> ProcessorCapabilities {
        ProcessorCapabilities::tpu(TpuVersion::V4)
    }

    /// Google Edge TPU.
    pub fn google_edge_tpu() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 0.0,
                fp32_tflops: 0.0,
                fp16_tflops: 0.0,
                bf16_tflops: 0.0,
                int8_tops: 4.0,
                int4_tops: 8.0,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 0, // Uses host memory
                bandwidth_gbps: 0,
                type_: MemoryType::Unified,
            },
            operations: {
                let mut ops = HashSet::new();
                ops.insert(OperationType::MatMul);
                ops.insert(OperationType::Conv2d);
                ops.insert(OperationType::DepthwiseConv);
                ops.insert(OperationType::Add);
                ops.insert(OperationType::Mul);
                ops.insert(OperationType::ReLU);
                ops.insert(OperationType::Softmax);
                ops
            },
            power: PowerCharacteristics {
                tdp_watts: 2,
                efficiency: 0.95,
                power_tier: PowerTier::UltraLow,
            },
            optimal_for: vec![
                WorkloadCharacteristic::LowPower,
                WorkloadCharacteristic::LowLatency,
                WorkloadCharacteristic::SmallBatch,
            ],
        }
    }

    // ═══════════════════════════════════════════════════════════════
    // GROQ LPU PROFILE
    // ═══════════════════════════════════════════════════════════════

    /// Groq LPU.
    pub fn groq_lpu() -> ProcessorCapabilities {
        ProcessorCapabilities::lpu()
    }

    // ═══════════════════════════════════════════════════════════════
    // APPLE NEURAL ENGINE PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// Apple Neural Engine (generic).
    pub fn apple_neural_engine(cores: u32) -> ProcessorCapabilities {
        ProcessorCapabilities::apple_neural_engine(cores)
    }

    /// Apple M3 Neural Engine (16 cores).
    pub fn apple_m3_neural_engine() -> ProcessorCapabilities {
        ProcessorCapabilities::apple_neural_engine(16)
    }

    /// Apple M3 Max Neural Engine (16 cores).
    pub fn apple_m3_max_neural_engine() -> ProcessorCapabilities {
        ProcessorCapabilities::apple_neural_engine(16) // Same as M3
    }

    /// Apple A17 Pro Neural Engine (35 TOPS).
    pub fn apple_a17_pro_neural_engine() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 0.0,
                fp32_tflops: 4.4,
                fp16_tflops: 8.8,
                bf16_tflops: 8.8,
                int8_tops: 35.0,
                int4_tops: 70.0,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 0, // Uses unified memory
                bandwidth_gbps: 200,
                type_: MemoryType::Unified,
            },
            operations: ProcessorCapabilities::apple_neural_engine(16)
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 8,
                efficiency: 0.98,
                power_tier: PowerTier::UltraLow,
            },
            optimal_for: vec![
                WorkloadCharacteristic::LowPower,
                WorkloadCharacteristic::LowLatency,
                WorkloadCharacteristic::SmallBatch,
            ],
        }
    }

    // ═══════════════════════════════════════════════════════════════
    // QUALCOMM NPU PROFILES
    // ═══════════════════════════════════════════════════════════════

    /// Qualcomm Hexagon NPU (Snapdragon 8 Gen 3).
    pub fn qualcomm_hexagon_8g3() -> ProcessorCapabilities {
        ProcessorCapabilities {
            compute: ComputeThroughput {
                fp64_tflops: 0.0,
                fp32_tflops: 3.0,
                fp16_tflops: 6.0,
                bf16_tflops: 6.0,
                int8_tops: 73.0, // 73 TOPS
                int4_tops: 146.0,
                sparsity_speedup: 1.0,
            },
            memory: MemorySpecs {
                capacity_bytes: 0, // Uses system memory
                bandwidth_gbps: 77,
                type_: MemoryType::Lpddr,
            },
            operations: ProcessorCapabilities::apple_neural_engine(16)
                .operations,
            power: PowerCharacteristics {
                tdp_watts: 10,
                efficiency: 0.95,
                power_tier: PowerTier::UltraLow,
            },
            optimal_for: vec![
                WorkloadCharacteristic::LowPower,
                WorkloadCharacteristic::LowLatency,
                WorkloadCharacteristic::SmallBatch,
            ],
        }
    }
}

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

    #[test]
    fn test_h100_profile() {
        let h100 = ProcessorProfiles::nvidia_h100();
        assert!(h100.compute.fp16_tflops > 1000.0);
        assert_eq!(h100.memory.capacity_bytes, 80 * 1024 * 1024 * 1024);
    }

    #[test]
    fn test_tpu_v5p_profile() {
        let tpu = ProcessorProfiles::google_tpu_v5p();
        assert!(tpu.compute.bf16_tflops > 900.0);
    }

    #[test]
    fn test_groq_lpu_profile() {
        let lpu = ProcessorProfiles::groq_lpu();
        assert!(lpu.memory.bandwidth_gbps > 50000); // Very high internal bandwidth
    }

    #[test]
    fn test_apple_ane_profile() {
        let ane = ProcessorProfiles::apple_m3_neural_engine();
        assert!(ane.power.tdp_watts < 20);
        assert!(ane.optimal_for.contains(&WorkloadCharacteristic::LowPower));
    }
}
