//! Profiled model implementations.
//!
//! This module profiles REAL CPU tensor compute (matmul + softmax) via
//! `rtx-tensor`, with matrix dimensions scaled proportionally to represent
//! the relative cost of different model classes (e.g. `ResNet18` uses a much
//! smaller matmul than `ViTL16`). These are NOT literal architectural
//! replicas of ResNet/ViT/ConvNeXt — the goal is proportional, measurable
//! FLOPs, not model fidelity.
//!
//! `DeviceType::CUDA` and `DeviceType::Metal` are accepted for API
//! compatibility with the existing config surface, but this demo does not
//! pull in the `cuda`/`metal` features of `rtx-tensor`. All device variants
//! currently execute the exact same CPU compute path on `Device::cpu()` — no
//! GPU numbers are fabricated or claimed. If GPU profiling is desired, wire
//! in the `cuda`/`metal` features of `rtx-tensor` and dispatch accordingly.

use inference_profiler_shared::{DeviceType, ModelType};
use rtx_tensor::{Device, Tensor};
use std::time::{Duration, Instant};

/// Approximate matmul dimension (`n`) used to scale FLOPs for a given model
/// type. The matmul performed is roughly `[n, n] x [n, n]` (see
/// [`ProfiledModel::matmul_dim`]), chosen so relative timings track the
/// documented relative cost ordering of these architectures:
/// `ResNet18` < `ResNet50` < `ViTB16` < `ViTL16`, and
/// `ConvNeXtTiny` < `ConvNeXtBase`.
fn base_matmul_dim(model_type: ModelType) -> usize {
    match model_type {
        ModelType::ResNet18 => 96,
        ModelType::ResNet50 => 160,
        ModelType::ViTB16 => 256,
        ModelType::ViTL16 => 448,
        ModelType::ConvNeXtTiny => 192,
        ModelType::ConvNeXtBase => 288,
        ModelType::Custom => 128,
    }
}

/// A model whose `forward()` performs real CPU tensor compute via
/// `rtx-tensor` (matmul -> softmax -> matmul), scaled proportionally to the
/// configured [`ModelType`] and `batch_size`.
///
/// See module docs for the important caveat about `DeviceType`.
pub struct ProfiledModel {
    model_type: ModelType,
    device: DeviceType,
    batch_size: usize,
}

impl ProfiledModel {
    /// Creates a new profiled model.
    #[must_use]
    pub fn new(model_type: ModelType, device: DeviceType, batch_size: usize) -> Self {
        Self {
            model_type,
            device,
            batch_size,
        }
    }

    /// Matmul dimension used for this model's forward pass, scaled by
    /// `batch_size` on the leading dimension only (so time grows with batch
    /// size while the "model size" dimension stays fixed per model type).
    fn matmul_dim(&self) -> usize {
        base_matmul_dim(self.model_type)
    }

    /// Runs a real forward pass: `matmul -> softmax -> matmul` on
    /// CPU-resident `rtx_tensor::Tensor`s, and returns the wall-clock
    /// duration of the REAL computation (no fabricated sleeps).
    ///
    /// # Panics
    /// Panics if tensor allocation or the tensor ops fail (e.g. OOM), which
    /// would indicate an environment problem rather than a profiling
    /// concern.
    #[must_use]
    pub fn forward(&self) -> Duration {
        let device = Device::cpu();
        let n = self.matmul_dim();
        let rows = self.batch_size * n;

        let a = Tensor::randn(&[rows, n], &device).expect("failed to allocate input tensor a");
        let b = Tensor::randn(&[n, n], &device).expect("failed to allocate input tensor b");
        let c = Tensor::randn(&[n, n], &device).expect("failed to allocate input tensor c");

        let start = Instant::now();

        let attn_scores = a.matmul(&b).expect("matmul 1 failed");
        let attn_weights = attn_scores.softmax(1).expect("softmax failed");
        let _out = attn_weights.matmul(&c).expect("matmul 2 failed");

        start.elapsed()
    }

    /// Real allocated tensor memory (bytes -> MB) for the tensors involved
    /// in one forward pass at this model's configured batch size, assuming
    /// f32 (4 bytes/element) storage. This sums the three input tensors
    /// (`a`, `b`, `c`) plus the two intermediate outputs (`attn_scores`,
    /// `attn_weights`) and the final output tensor -- i.e. all real tensor
    /// allocations `forward()` performs.
    #[must_use]
    pub fn memory_usage_mb(&self) -> f64 {
        let n = self.matmul_dim();
        let rows = self.batch_size * n;

        let a_elems = rows * n; // [rows, n]
        let b_elems = n * n; // [n, n]
        let c_elems = n * n; // [n, n]
        let attn_scores_elems = rows * n; // a.matmul(b) -> [rows, n]
        let attn_weights_elems = rows * n; // softmax same shape
        let out_elems = rows * n; // attn_weights.matmul(c) -> [rows, n]

        let total_elems =
            a_elems + b_elems + c_elems + attn_scores_elems + attn_weights_elems + out_elems;
        let bytes = total_elems as f64 * 4.0;
        bytes / (1024.0 * 1024.0)
    }

    /// The configured device. Note that on this demo all `DeviceType`
    /// variants execute identical CPU compute (see module docs).
    #[must_use]
    pub fn device(&self) -> DeviceType {
        self.device
    }

    /// The configured model type.
    #[must_use]
    pub fn model_type(&self) -> ModelType {
        self.model_type
    }

    /// The configured batch size.
    #[must_use]
    pub fn batch_size(&self) -> usize {
        self.batch_size
    }
}

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

    #[test]
    fn test_profiled_model_creation() {
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 8);
        assert_eq!(model.model_type(), ModelType::ResNet18);
        assert_eq!(model.device(), DeviceType::CPU);
        assert_eq!(model.batch_size(), 8);
    }

    #[test]
    fn test_forward_returns_duration() {
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        let duration = model.forward();
        assert!(duration > Duration::ZERO);
    }

    #[test]
    fn test_forward_scales_with_batch_size() {
        // Real compute: larger batch should take at least as long, generally more,
        // though CPU jitter means we only assert a loose directional trend across
        // several samples rather than a single noisy measurement.
        let small = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        let large = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 16);

        let small_total: Duration = (0..3).map(|_| small.forward()).sum();
        let large_total: Duration = (0..3).map(|_| large.forward()).sum();

        assert!(large_total > small_total);
    }

    #[test]
    fn test_resnet18_faster_than_resnet50() {
        let resnet18 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 4);
        let resnet50 = ProfiledModel::new(ModelType::ResNet50, DeviceType::CPU, 4);

        assert!(resnet18.matmul_dim() < resnet50.matmul_dim());
    }

    #[test]
    fn test_vit_b16_slower_than_resnet50() {
        let resnet50 = ProfiledModel::new(ModelType::ResNet50, DeviceType::CPU, 4);
        let vit_b16 = ProfiledModel::new(ModelType::ViTB16, DeviceType::CPU, 4);

        assert!(vit_b16.matmul_dim() > resnet50.matmul_dim());
    }

    #[test]
    fn test_vit_l16_slower_than_vit_b16() {
        let vit_b16 = ProfiledModel::new(ModelType::ViTB16, DeviceType::CPU, 4);
        let vit_l16 = ProfiledModel::new(ModelType::ViTL16, DeviceType::CPU, 4);

        assert!(vit_l16.matmul_dim() > vit_b16.matmul_dim());
    }

    #[test]
    fn test_convnext_base_slower_than_convnext_tiny() {
        let tiny = ProfiledModel::new(ModelType::ConvNeXtTiny, DeviceType::CPU, 4);
        let base = ProfiledModel::new(ModelType::ConvNeXtBase, DeviceType::CPU, 4);

        assert!(base.matmul_dim() > tiny.matmul_dim());
    }

    #[test]
    fn test_memory_usage_scales_with_batch_size() {
        let model1 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        let model8 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 8);

        let mem1 = model1.memory_usage_mb();
        let mem8 = model8.memory_usage_mb();

        // Directional check only: real allocation sizes, no fabricated constant.
        assert!(mem8 > mem1);
    }

    #[test]
    fn test_memory_usage_positive() {
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        assert!(model.memory_usage_mb() > 0.0);
    }

    #[test]
    fn test_memory_usage_larger_models() {
        let resnet18 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        let vit_l16 = ProfiledModel::new(ModelType::ViTL16, DeviceType::CPU, 1);

        assert!(vit_l16.memory_usage_mb() > resnet18.memory_usage_mb());
    }

    #[test]
    fn test_forward_has_measurable_jitter_or_stability() {
        // Real CPU compute may or may not show visible jitter at this scale;
        // just assert it's always a positive, finite duration.
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);

        let durations: Vec<Duration> = (0..5).map(|_| model.forward()).collect();
        for d in durations {
            assert!(d > Duration::ZERO);
        }
    }

    #[test]
    fn test_cuda_and_metal_execute_same_cpu_path() {
        // Documented behavior: without gpu features wired in, CUDA/Metal
        // DeviceType variants run identical CPU compute to CPU.
        let cuda_model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CUDA, 1);
        let metal_model = ProfiledModel::new(ModelType::ResNet18, DeviceType::Metal, 1);

        assert_eq!(cuda_model.matmul_dim(), metal_model.matmul_dim());
    }

    #[test]
    fn test_custom_model_has_reasonable_dim() {
        let model = ProfiledModel::new(ModelType::Custom, DeviceType::CPU, 1);
        assert!(model.matmul_dim() > 0);
    }

    #[test]
    fn test_custom_model_has_reasonable_memory() {
        let model = ProfiledModel::new(ModelType::Custom, DeviceType::CPU, 1);
        let memory = model.memory_usage_mb();

        assert!(memory > 0.0);
        assert!(memory < 1000.0); // Reasonable range at batch_size=1
    }
}