#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GPU Infrastructure Recommender for AI Models
===========================================

A comprehensive tool for estimating VRAM requirements and recommending optimal
GPU configurations for Large Language Model (LLM) deployment and training.

Author: Rudali Huidrom
Version: 4.0.0
First Written On: 08 December 2025
Last Updated: 18 December 2025

Overview
--------
This module provides a Gradio-based web interface that:
1. Calculates precise VRAM requirements for LLMs based on model specifications
2. Estimates throughput and time-to-completion for various GPU configurations
3. Recommends cost-effective hardware solutions for both inference and training
4. Supports multiple quantization methods, fine-tuning strategies, and frameworks

Key Features
-----------
- Automatic model resolution from HuggingFace Hub
- Support for inference (single/batched) and training (Full FT, LoRA, QLoRA)
- Empirical throughput benchmarks for major GPU families
- Multi-GPU configuration support with communication overhead modeling
- Framework-specific optimizations (vLLM, HuggingFace, TensorRT)
- Real-time cost estimation with multiple pricing tiers
- Throughput scaling with sequence length
- CSV export with detailed analysis and formulas
- Responsive UI with dark mode support

Technical Approach
-----------------
The recommender uses empirically-validated formulas derived from:
- MLPerf benchmarks and vendor specifications
- Production deployment data from real-world LLM serving
- Memory profiling of training workloads
- Community-contributed performance metrics

Accuracy: ±15-20% variance expected due to architecture-specific optimizations,
framework versions, and runtime conditions.

Dependencies
-----------
Required:
    - gradio>=3.0.0: Web interface framework
    - python>=3.8: Core language support

Optional:
    - transformers>=4.30.0: Automatic model config resolution
    - huggingface_hub>=0.16.0: HuggingFace API access for gated models

Usage
-----
    python app.py

Environment Variables:
    HF_TOKEN: HuggingFace API token for accessing gated models

License
-------
Copyright (c) 2025. All rights reserved.
"""

# =================================================================================================
# IMPORTS
# =================================================================================================

import math
import re
import os
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
from enum import Enum
import gradio as gr

# =================================================================================================
# CONFIGURATION AND CONSTANTS
# =================================================================================================

# Authentication token for HuggingFace API access
# Set via environment variable: export HF_TOKEN="your_token_here"
HF_TOKEN = os.getenv("HF_TOKEN", "")
# =================================================================================================
# UI STYLING
# =================================================================================================

# Custom CSS for visual differentiation of recommendation tiers
# - Budget tier: Green gradient for cost-effective options
# - Runner-up tier: Blue gradient for balanced options
# - Performance tier: Purple gradient for maximum performance
# - Dark/Light mode support
# - Mobile/Tablet responsive
CUSTOM_CSS = """
<style>
/* ============================================
   LIGHT MODE (Default)
   ============================================ */
.budget-box { 
    background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%) !important; 
    border: 2px solid #22c55e !important; 
    border-radius: 12px !important; 
    padding: 16px !important;
    box-shadow: 0 2px 8px rgba(34, 197, 94, 0.1) !important;
}
.runner-box { 
    background: linear-gradient(135deg, #eff6ff 0%, #dbeafe 100%) !important; 
    border: 2px solid #3b82f6 !important; 
    border-radius: 12px !important; 
    padding: 16px !important;
    box-shadow: 0 2px 8px rgba(59, 130, 246, 0.1) !important;
}
.perf-box { 
    background: linear-gradient(135deg, #faf5ff 0%, #f3e8ff 100%) !important; 
    border: 2px solid #a855f7 !important; 
    border-radius: 12px !important; 
    padding: 16px !important;
    box-shadow: 0 2px 8px rgba(168, 85, 247, 0.1) !important;
}
.warning-box {
    background-color: #fef3c7 !important;
    border: 1px solid #f59e0b !important;
    border-radius: 8px !important;
    padding: 12px !important;
    margin: 8px 0 !important;
}
.error-box {
    background-color: #fee2e2 !important;
    border: 1px solid #ef4444 !important;
    border-radius: 8px !important;
    padding: 12px !important;
    margin: 8px 0 !important;
}

/* ============================================
   DARK MODE
   ============================================ */
@media (prefers-color-scheme: dark) {
    .budget-box { 
        background: linear-gradient(135deg, #052e16 0%, #064e3b 100%) !important; 
        border: 2px solid #22c55e !important; 
        box-shadow: 0 2px 8px rgba(34, 197, 94, 0.25) !important;
    }
    .budget-box, .budget-box * {
        color: #dcfce7 !important;
    }
    
    .runner-box { 
        background: linear-gradient(135deg, #172554 0%, #1e3a8a 100%) !important; 
        border: 2px solid #3b82f6 !important; 
        box-shadow: 0 2px 8px rgba(59, 130, 246, 0.25) !important;
    }
    .runner-box, .runner-box * {
        color: #dbeafe !important;
    }
    
    .perf-box { 
        background: linear-gradient(135deg, #3b0764 0%, #4c1d95 100%) !important; 
        border: 2px solid #a855f7 !important; 
        box-shadow: 0 2px 8px rgba(168, 85, 247, 0.25) !important;
    }
    .perf-box, .perf-box * {
        color: #f3e8ff !important;
    }
    
    .warning-box {
        background-color: #451a03 !important;
        border: 1px solid #f59e0b !important;
        color: #fde68a !important;
    }
    
    .error-box {
        background-color: #450a0a !important;
        border: 1px solid #ef4444 !important;
        color: #fecaca !important;
    }
}

/* ============================================
   TABLET (max-width: 1024px)
   Stack main columns, keep boxes side by side
   ============================================ */
@media screen and (max-width: 1024px) {
    .gradio-container {
        padding: 12px !important;
    }
}

/* ============================================
   MOBILE (max-width: 768px)
   ============================================ */
@media screen and (max-width: 768px) {
    .gradio-container {
        padding: 8px !important;
    }
    
    /* Stack recommendation boxes vertically */
    .budget-box, .runner-box, .perf-box {
        margin-bottom: 12px !important;
    }
    
    /* Larger touch targets */
    button, .gradio-button {
        min-height: 44px !important;
        font-size: 16px !important;
    }
    
    input, select, textarea {
        min-height: 44px !important;
        font-size: 16px !important; /* Prevents iOS zoom */
    }
}

/* ============================================
   SMALL MOBILE (max-width: 480px)
   ============================================ */
@media screen and (max-width: 480px) {
    .gradio-container {
        padding: 4px !important;
    }
    
    .budget-box, .runner-box, .perf-box {
        padding: 12px !important;
        border-radius: 8px !important;
    }
    
    h1, .markdown h1 {
        font-size: 1.4rem !important;
    }
    
    h3, .markdown h3 {
        font-size: 1rem !important;
    }
}

/* ============================================
   TOUCH DEVICES
   ============================================ */
@media (pointer: coarse) {
    button, .gradio-button, 
    input, select, textarea {
        min-height: 44px !important;
    }
    
    input[type="checkbox"], input[type="radio"] {
        width: 20px !important;
        height: 20px !important;
    }
}

/* ============================================
   ACCESSIBILITY
   ============================================ */
@media (prefers-reduced-motion: reduce) {
    * {
        transition: none !important;
        animation: none !important;
    }
}

*:focus-visible {
    outline: 2px solid #3b82f6 !important;
    outline-offset: 2px !important;
}

/* ============================================
   PLOTLY CHART RESPONSIVE STYLES
   ============================================ */
/* Make Plotly chart container responsive */
.js-plotly-plot, .plotly {
    width: 100% !important;
}

.js-plotly-plot .plotly .main-svg {
    width: 100% !important;
}

/* Mobile: Ensure chart doesn't overflow */
@media screen and (max-width: 768px) {
    .js-plotly-plot, .plotly, .plot-container {
        width: 100% !important;
        overflow-x: auto !important;
    }
    
    /* Make Gradio plot component full width */
    .gradio-plot {
        width: 100% !important;
        min-height: 400px !important;
    }
}

/* Small mobile: Compact chart */
@media screen and (max-width: 480px) {
    .gradio-plot {
        min-height: 350px !important;
    }
}

/* Landscape phone: Give more width to chart */
@media screen and (max-width: 900px) and (orientation: landscape) {
    .gradio-plot {
        min-height: 300px !important;
    }
}
</style>
"""

# =================================================================================================
# DATA STRUCTURES AND TYPE DEFINITIONS
# =================================================================================================

class Task(Enum):
    """
    Enumeration of supported computational tasks.
    
    Attributes:
        INFERENCE: Model inference/serving workloads
        TRAINING: Model training/fine-tuning workloads
    """
    INFERENCE = "Inference"
    TRAINING = "Training"

class FineTuningMethod(Enum):
    """
    Enumeration of supported fine-tuning strategies.
    
    Attributes:
        FULL: Full fine-tuning (all parameters trainable)
        LORA: Low-Rank Adaptation (parameter-efficient, full precision base)
        QLORA: Quantized LoRA (parameter-efficient, quantized base)
    """
    FULL = "Full Fine-Tuning"
    LORA = "LoRA"
    QLORA = "QLoRA"

class Framework(Enum):
    """
    Enumeration of supported inference/training frameworks.
    
    Attributes:
        VLLM: vLLM (optimized for high-throughput inference)
        HUGGINGFACE: HuggingFace Transformers (general-purpose)
    """
    VLLM = "vllm"
    HUGGINGFACE = "huggingface"

@dataclass
class GPUConfig:
    """
    Configuration specification for GPU hardware.
    
    This dataclass encapsulates all relevant specifications and pricing
    information for a GPU configuration, supporting both single and
    multi-GPU setups.
    
    Attributes:
        name (str): Human-readable identifier (e.g., "Nvidia H100 SXM (8x)")
        vram (int): Total VRAM across all GPUs in GB
        count (int): Number of GPUs in this configuration
        tflops (float): Total TFLOPS (FP16) across all GPUs
        bandwidth (int): Total memory bandwidth in GB/s
        price_od (float): On-demand hourly rate in INR
        price_1m (float): 1-month reserved hourly rate in INR
        price_6m (float): 6-month reserved hourly rate in INR
        price_12m (float): 12-month reserved hourly rate in INR
    
    Properties:
        vram_per_gpu (float): VRAM per individual GPU
        
    Methods:
        get_price(tier): Returns price for specified tier
    """
    name: str
    vram: int  # Total VRAM in GB
    count: int  # Number of GPUs in config
    tflops: float
    bandwidth: int  # GB/s
    price_od: float  # On-demand price in INR/hour
    price_1m: float  # 1-month reserved
    price_6m: float  # 6-month reserved
    price_12m: float  # 12-month reserved
    
    @property
    def vram_per_gpu(self) -> float:
        """
        Calculate VRAM per individual GPU.
        
        Returns:
            float: VRAM in GB for a single GPU in this configuration
        """
        return self.vram / self.count
    
    def get_price(self, tier: str) -> float:
        """
        Retrieve price for specified pricing tier.
        
        Args:
            tier (str): Pricing tier ("On Demand", "1 Month Reserved", etc.)
            
        Returns:
            float: Hourly rate in INR for the specified tier
        """
        price_map = {
            "On Demand": self.price_od,
            "1 Month Reserved": self.price_1m,
            "6 Month Reserved": self.price_6m,
            "12 Month Reserved": self.price_12m,
        }
        return price_map.get(tier, self.price_od)

@dataclass
class ModelSpec:
    """
    Specification of transformer model architecture.
    
    Encapsulates key architectural parameters required for accurate
    memory and performance estimation.
    
    Attributes:
        params (int): Total number of model parameters
        layers (int): Number of transformer layers
        heads (int): Number of attention heads
        kv_heads (int): Number of key/value heads (for GQA/MQA)
        head_dim (int): Dimension of each attention head
        context (int): Maximum context length (position embeddings)
    
    Properties:
        params_bn (float): Parameters in billions
        
    Notes:
        For standard Multi-Head Attention: kv_heads = heads
        For Grouped Query Attention (GQA): kv_heads < heads
        Example: Llama 3 uses heads=32, kv_heads=8 (4:1 ratio)
    """
    params: int  # Total parameters
    layers: int
    heads: int
    kv_heads: int
    head_dim: int
    context: int
    
    @property
    def params_bn(self) -> float:
        """
        Convert parameter count to billions.
        
        Returns:
            float: Number of parameters in billions (1e9)
        """
        return self.params / 1e9

# =================================================================================================
# MODEL DATABASE AND CONSTANTS
# =================================================================================================

# Popular pre-trained models available in the dropdown selector
# Sourced from HuggingFace Hub's most-used instruction-tuned models
MODEL_CHOICES = [
    "meta-llama/Llama-3.3-70B-Instruct",
    "meta-llama/Llama-3.1-405B-Instruct",
    "meta-llama/Llama-3.1-70B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct",
    "meta-llama/Llama-3.2-3B-Instruct",
    "meta-llama/Llama-3.2-1B-Instruct",
    "Qwen/Qwen2.5-72B-Instruct",
    "Qwen/Qwen2.5-32B-Instruct",
    "Qwen/Qwen2.5-14B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct",
    "Qwen/Qwen2.5-3B-Instruct",
    "Qwen/Qwen2.5-1.5B-Instruct",
    "Qwen/Qwen2.5-Coder-32B-Instruct",
    "mistralai/Mistral-Large-Instruct-2411",
    "mistralai/Mistral-Small-Instruct-2409",
    "mistralai/Mistral-Nemo-Instruct-2407",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "mistralai/Mixtral-8x22B-Instruct-v0.1",
    "mistralai/Ministral-8B-Instruct-2410",
]

# =================================================================================================
# PRECISION AND QUANTIZATION SPECIFICATIONS
# =================================================================================================

# Mapping of precision formats to bytes per parameter
# Used for accurate memory footprint calculation across different quantization schemes
#
# Precision Format Categories:
#   Full Precision: fp32 (4 bytes) - Maximum accuracy, highest memory
#   Half Precision: fp16, bf16 (2 bytes) - Standard training/inference
#   Quantized: int8 (1 byte) - 4x compression, minimal quality loss
#   Low-bit: int4, nf4 (0.5-0.56 bytes) - 8x compression, some quality degradation
#   Compressed: awq, gptq (~0.52 bytes) - Advanced quantization with lookup tables
#
# Note: nf4 (NormalFloat4) is specifically designed for QLoRA and provides
# better quality than standard int4 at the same bitwidth
PRECISION_MAP = {
    # "float32": 4.0,
    # "fp32": 4.0,
    "bf16": 2.0,        # BFloat16 - preferred for training (better range than fp16)
    "fp16": 2.0,        # Float16 - standard for inference
    "nf4": 0.5625,      # NormalFloat4 - QLoRA's quantization format
    "4bit": 0.5625,
    "int4": 0.50,
    "int8": 1.0,
    "awq": 0.52,        # Activation-aware Weight Quantization (inference-only)
    "gptq": 0.52,       # GPTQ quantization (inference-only)
}

# Framework-specific memory overhead (in GB)
# Represents additional memory required by the framework runtime beyond model weights
#
# Factors contributing to overhead:
#   - Kernel workspace and temporary buffers
#   - Execution graph and operator metadata
#   - Memory pools and allocator overhead
#   - Framework-specific data structures
#
# These values are empirically determined from profiling real deployments
FRAMEWORK_OVERHEAD = {
    "vllm": 1.5,           # PagedAttention + continuous batching optimizations
    "huggingface": 3.5,    # Flexible abstractions + dynamic computation graph
    "tensorrt": 1.0,       # Highly optimized CUDA graphs + operator fusion
}

# Quantization methods that only support inference workloads
# These methods modify weight representation in ways incompatible with gradient computation
# Training requires full-precision gradients for optimizer updates
INFERENCE_ONLY_QUANT = ['awq', 'gptq', 'exl2']

# Throughput speedup factors for different quantization methods
# Values represent throughput multiplier relative to FP16 baseline
# Based on NVIDIA TensorRT-LLM, vLLM, and MLPerf benchmarks
# Conservative estimates to avoid over-promising
QUANTIZATION_SPEEDUP = {
    # "fp32": 0.8,     # Slightly slower than FP16 (more compute required)
    # "float32": 0.8,
    "fp16": 1.0,     # Baseline reference
    "bf16": 1.0,     # Same throughput as FP16
    "int8": 1.8,     # ~2x faster (INT8 Tensor Cores + less bandwidth)
    "int4": 3.0,     # ~3-4x faster (INT4 Tensor Cores + 4x less bandwidth)
    "4bit": 3.0,
    "nf4": 3.0,      # Similar to INT4
    "awq": 3.2,      # Optimized INT4 quantization
    "gptq": 3.2,     # Optimized INT4 quantization
}

# Framework efficiency multipliers relative to vLLM baseline
# Based on production benchmarks and community reports
# vLLM is set as baseline (1.0) as it's highly optimized for inference
FRAMEWORK_SPEEDUP = {
    "vllm": 1.0,           # Baseline (PagedAttention, continuous batching, optimized)
    "huggingface": 0.7,    # More flexible but less optimized (~30% slower)
    "tensorrt": 1.3,       # Most optimized for NVIDIA GPUs (~30% faster)
}

# =================================================================================================
# GPU HARDWARE DATABASE
# =================================================================================================

# Comprehensive database of available GPU configurations
# Each entry represents a specific hardware configuration with associated pricing
# Pricing is in Indian Rupees (INR) per hour for various reservation tiers
GPU_DATABASE = [
    # AMD MI300X
    GPUConfig('AMD MI300X (1x)', 192, 1, 1300.0, 5300, 168.224, 165.048, 161.88, 148.0),
    GPUConfig('AMD MI300X (2x)', 384, 2, 2600.0, 10600, 378.504, 371.358, 364.23, 333.0),
    GPUConfig('AMD MI300X (4x)', 768, 4, 5200.0, 21200, 757.008, 742.716, 728.46, 666.0),
    GPUConfig('AMD MI300X (8x)', 1536, 8, 10400.0, 42400, 1416.56, 1389.904, 1363.2, 1336.0),
    
    # AMD MI325X
    GPUConfig('AMD MI325X (1x)', 256, 1, 1300.0, 6000, 169.2, 123.3, 102.6, 85.5),
    GPUConfig('AMD MI325X (2x)', 512, 2, 2600.0, 12000, 338.4, 246.6, 205.2, 171.0),
    GPUConfig('AMD MI325X (4x)', 1024, 4, 5200.0, 24000, 676.8, 493.2, 410.4, 342.0),
    GPUConfig('AMD MI325X (8x)', 2048, 8, 10400.0, 48000, 1351.8, 990.0, 820.8, 684.0),
    
    # NVIDIA H100 SXM
    GPUConfig('Nvidia H100 SXM (1x)', 80, 1, 1979.0, 3350, 153.0, 134.1, 125.1, 117.0),
    GPUConfig('Nvidia H100 SXM (2x)', 160, 2, 3958.0, 6700, 306.0, 268.2, 250.2, 234.0),
    GPUConfig('Nvidia H100 SXM (4x)', 320, 4, 7916.0, 13400, 612.0, 536.4, 500.4, 468.0),
    GPUConfig('Nvidia H100 SXM (8x)', 640, 8, 15832.0, 26800, 1224.0, 1072.8, 1000.8, 936.0),
    
    # NVIDIA H100 NVL
    GPUConfig('Nvidia H100 NVL (1x)', 94, 1, 1671.0, 3900, 140.0, 135.0, 118.0, 100.0),
    GPUConfig('Nvidia H100 NVL (2x)', 188, 2, 3342.0, 7800, 337.48, 294.44, 274.04, 257.08),
    GPUConfig('Nvidia H100 NVL (4x)', 376, 4, 6684.0, 15600, 674.96, 588.88, 548.08, 514.16),
    GPUConfig('Nvidia H100 NVL (8x)', 752, 8, 13368.0, 31200, 1349.92, 1177.76, 1096.16, 1028.32),
    
    # NVIDIA H100 PCIe
    GPUConfig('Nvidia H100 PCIe (1x)', 80, 1, 1513.0, 2000, 252.0, 234.0, 209.0, 185.0),
    GPUConfig('Nvidia H100 PCIe (8x)', 640, 8, 12104.0, 16000, 2008.0, 1864.0, 1664.0, 1472.0),
    
    # NVIDIA H200 SXM
    GPUConfig('Nvidia H200 SXM (1x)', 141, 1, 1979.0, 4800, 140.0, 135.0, 118.0, 100.0),
    GPUConfig('Nvidia H200 SXM (2x)', 282, 2, 3958.0, 9600, 510.0, 448.0, 418.0, 390.0),
    GPUConfig('Nvidia H200 SXM (4x)', 564, 4, 7916.0, 19200, 1020.0, 896.0, 836.0, 780.0),
    GPUConfig('Nvidia H200 SXM (8x)', 1128, 8, 15832.0, 38400, 1125.0, 1100.0, 945.0, 785.0),
    
    # NVIDIA H200 NVL
    GPUConfig('Nvidia H200 NVL (1x)', 141, 1, 1671.0, 3900, 146.38, 143.61, 140.85, 138.09),
    GPUConfig('Nvidia H200 NVL (2x)', 282, 2, 3342.0, 7800, 292.75, 287.23, 281.7, 276.18),
    GPUConfig('Nvidia H200 NVL (4x)', 564, 4, 6684.0, 15600, 585.5, 574.45, 563.41, 552.36),
    GPUConfig('Nvidia H200 NVL (8x)', 1128, 8, 13368.0, 31200, 1171.0, 1148.91, 1126.81, 1104.72),
    
    # NVIDIA H200 PCIe
    GPUConfig('Nvidia H200 PCIe (8x)', 1128, 8, 13368.0, 31200, 3236.8, 2737.0, 2665.6, 2380.0),
    
    # NVIDIA B200 SXM
    GPUConfig('Nvidia B200 SXM (1x)', 180, 1, 4500.0, 8000, 323.0, 308.0, 293.0, 279.0),
    GPUConfig('Nvidia B200 SXM (2x)', 360, 2, 9000.0, 16000, 646.0, 616.0, 586.0, 558.0),
    GPUConfig('Nvidia B200 SXM (4x)', 720, 4, 18000.0, 32000, 1292.0, 1232.0, 1172.0, 1116.0),
    GPUConfig('Nvidia B200 SXM (8x)', 1440, 8, 36000.0, 64000, 2584.0, 2464.0, 2344.0, 2232.0),
    
    # NVIDIA A100 40GB
    GPUConfig('Nvidia A100 40GB (1x)', 40, 1, 312.0, 1935, 136.0, 89.0, 85.0, 81.0),
    GPUConfig('Nvidia A100 40GB (2x)', 80, 2, 624.0, 3870, 272.0, 178.0, 170.0, 162.0),
    GPUConfig('Nvidia A100 40GB (4x)', 160, 4, 1248.0, 7740, 544.0, 356.0, 340.0, 324.0),
    GPUConfig('Nvidia A100 40GB (8x)', 320, 8, 2496.0, 15480, 3175.66, 3175.66, 3175.66, 3175.66),
    
    # NVIDIA A100 80GB
    GPUConfig('Nvidia A100 80GB (1x)', 80, 1, 312.0, 1935, 135.9, 89.1, 85.5, 81.0),
    GPUConfig('Nvidia A100 80GB (2x)', 160, 2, 624.0, 3870, 271.8, 178.2, 171.0, 162.0),
    GPUConfig('Nvidia A100 80GB (4x)', 320, 4, 1248.0, 7740, 543.6, 356.4, 342.0, 324.0),
    GPUConfig('Nvidia A100 80GB (8x)', 640, 8, 2496.0, 15480, 1087.2, 712.8, 684.0, 648.0),
    
    # NVIDIA L40S
    GPUConfig('Nvidia L40S (1x)', 48, 1, 733.0, 864, 67.5, 49.5, 49.5, 45.0),
    GPUConfig('Nvidia L40S (2x)', 96, 2, 1466.0, 1728, 135.0, 99.0, 99.0, 90.0),
    GPUConfig('Nvidia L40S (4x)', 192, 4, 2932.0, 3456, 306.0, 198.0, 198.0, 180.0),
    GPUConfig('Nvidia L40S (8x)', 384, 8, 5864.0, 6912, 540.0, 396.0, 396.0, 360.0),
    
    # NVIDIA L4
    GPUConfig('Nvidia L4 (1x)', 24, 1, 242.0, 300, 45.07, 29.0, 26.75, 24.0),
    GPUConfig('Nvidia L4 (2x)', 48, 2, 484.0, 600, 98.84, 58.0, 54.0, 48.0),
    GPUConfig('Nvidia L4 (4x)', 96, 4, 968.0, 1200, 196.68, 116.0, 108.0, 96.0),
    GPUConfig('Nvidia L4 (8x)', 192, 8, 1936.0, 2400, 510.37, 495.06, 459.34, 302.51),
    
    # Intel Gaudi 2
    GPUConfig('Intel Gaudi 2 (1x)', 96, 1, 180.0, 600, 57.6, 46.8, 39.6, 34.2),
    GPUConfig('Intel Gaudi 2 (2x)', 192, 2, 360.0, 1200, 115.2, 93.6, 79.2, 68.4),
    GPUConfig('Intel Gaudi 2 (4x)', 384, 4, 720.0, 2400, 230.4, 187.2, 158.4, 136.8),
    GPUConfig('Intel Gaudi 2 (8x)', 768, 8, 1440.0, 4800, 460.8, 374.4, 316.8, 273.6),
    
    # Intel Gaudi 3
    GPUConfig('Intel Gaudi 3 (1x)', 128, 1, 459.0, 3600, 153.0, 134.1, 125.1, 117.0),
    GPUConfig('Intel Gaudi 3 (2x)', 256, 2, 918.0, 7200, 306.0, 268.2, 250.2, 234.0),
    GPUConfig('Intel Gaudi 3 (4x)', 512, 4, 1836.0, 14400, 612.0, 536.4, 500.4, 468.0),
    GPUConfig('Intel Gaudi 3 (8x)', 1024, 8, 3672.0, 28800, 1224.0, 1072.8, 1000.8, 936.0),
]

# =================================================================================================
# GPU Throughput Benchmarks (Empirical Data)
# =================================================================================================
# Based on real-world benchmarks from MLPerf, vendor data, and community testing
# Tokens per second per GPU for different model sizes
# 
# Last Updated: 18 December 2025
# Sources:
# - MLPerf Training v3.1 (November 2023)
# - NVIDIA TensorRT-LLM benchmarks (Q4 2024)
# - vLLM project benchmarks (Q4 2024)
# - Community benchmarks from HuggingFace, Anyscale
#
# Note: These are approximate values. Actual performance varies based on:
# - Specific model architecture
# - Sequence length
# - Batch size
# - Framework optimizations
# - Hardware configuration
# Expect ±15-20% variance in real-world usage

GPU_THROUGHPUT_BENCHMARKS = {
    # Format: GPU_name -> {model_size -> (inference_tps_single, inference_tps_batched, training_tps)}
    'H100': {
        7: (120, 1400, 1800),
        13: (80, 950, 1200),
        70: (10, 90, 360),
        405: (2, 25, 90),
    },
    'H200': {
        7: (130, 1500, 1950),
        13: (85, 1000, 1300),
        70: (11, 95, 390),
        405: (2, 27, 95),
    },
    'B200': {
        7: (160, 1800, 2400),
        13: (105, 1200, 1600),
        70: (13, 115, 480),
        405: (3, 32, 120),
    },
    'A100': {
        7: (80, 850, 900),
        13: (55, 580, 600),
        70: (6, 55, 180),
        405: (1, 5, 20),  
    },
    'L40S': {
        7: (45, 500, 700),
        13: (30, 330, 460),
        70: (2, 15, 80),  
        405: (0.5, 2, 10),  
    },
    'L4': {
        7: (25, 280, 300),
        13: (10, 120, 150),  
        70: (1, 8, 40),  
        405: (0.3, 1, 5),  
    },
    'MI300X': {
        7: (80, 850, 1000),
        13: (55, 580, 660),
        70: (6, 55, 200),
        405: (1, 5, 20),  
    },
    'MI325X': {
        7: (88, 935, 1100),
        13: (60, 640, 720),
        70: (7, 60, 220),
        405: (1, 6, 22),  
    },
    'Gaudi2': {
        7: (50, 650, 800),
        13: (35, 450, 540),
        20: (25, 330, 400),
        70: (6, 80, 160),
        405: (0.8, 4, 15),  
    },
    'Gaudi3': {
        7: (70, 900, 1120),
        13: (50, 630, 760),
        20: (35, 460, 560),
        70: (8, 110, 225),
        405: (1, 5, 18),  
    },
}

def get_gpu_family(gpu_name: str) -> str:
    """Extract GPU family from full GPU name."""
    if 'H200' in gpu_name:
        return 'H200'
    elif 'H100' in gpu_name:
        return 'H100'
    elif 'B200' in gpu_name:
        return 'B200'
    elif 'A100' in gpu_name:
        return 'A100'
    elif 'L40S' in gpu_name:
        return 'L40S'
    elif 'L4' in gpu_name:
        return 'L4'
    elif 'MI325X' in gpu_name:
        return 'MI325X'
    elif 'MI300X' in gpu_name:
        return 'MI300X'
    elif 'Gaudi 3' in gpu_name or 'Gaudi3' in gpu_name:
        return 'Gaudi3'
    elif 'Gaudi 2' in gpu_name or 'Gaudi2' in gpu_name:
        return 'Gaudi2'
    return 'H100'  # Default fallback

def interpolate_throughput(gpu_family: str, model_size_bn: float, task: str, batched: bool = False) -> float:
    """
    Interpolate throughput for a given GPU family and model size.
    
    Args:
        gpu_family: GPU family name (e.g., 'H100', 'A100')
        model_size_bn: Model size in billions of parameters
        task: 'Inference' or 'Training'
        batched: Whether to use batched inference numbers
        
    Returns:
        Estimated tokens per second per GPU
    """
    if gpu_family not in GPU_THROUGHPUT_BENCHMARKS:
        gpu_family = 'H100'  # Fallback
    
    benchmarks = GPU_THROUGHPUT_BENCHMARKS[gpu_family]
    
    # Get the right metric index: (single_inf, batched_inf, training)
    if task == "Inference":
        metric_idx = 1 if batched else 0
    else:
        metric_idx = 2
    
    # Create list of (size, throughput) tuples
    benchmark_points = [(size, values[metric_idx]) for size, values in benchmarks.items()]
    benchmark_points.sort()
    
    # Find surrounding points for interpolation
    for i in range(len(benchmark_points) - 1):
        size1, tps1 = benchmark_points[i]
        size2, tps2 = benchmark_points[i + 1]
        
        if size1 <= model_size_bn <= size2:
            # Log-linear interpolation (performance scales roughly inversely with size)
            log_size = math.log(model_size_bn)
            log_size1 = math.log(size1)
            log_size2 = math.log(size2)
            
            ratio = (log_size - log_size1) / (log_size2 - log_size1)
            log_tps = math.log(tps1) + ratio * (math.log(tps2) - math.log(tps1))
            
            return math.exp(log_tps)
    
    # Extrapolate if outside range
    if model_size_bn < benchmark_points[0][0]:
        size, tps = benchmark_points[0]
        return tps * (size / model_size_bn) ** 0.7
    else:
        size, tps = benchmark_points[-1]
        return tps * (size / model_size_bn) ** 0.7

def get_lora_overhead_factor(
    rank: Optional[int],
    ft_method: Optional[str],
    model_size_bn: float,
    spec: ModelSpec
) -> float:
    """
    Calculate throughput reduction factor due to LoRA adapters.
    
    LoRA adds computational overhead through extra matrix multiplications:
    - For each adapted layer: output = base_output + (B @ A @ input)
    - Where A is (hidden_dim × rank) and B is (rank × hidden_dim)
    - Higher rank = more computation = lower throughput
    
    This function uses empirically-measured overhead from:
    - QLoRA paper (Dettmers et al., 2023)
    - Community benchmarks (HuggingFace, Axolotl)
    - Production LoRA training deployments
    
    Args:
        rank: LoRA rank (None if not using LoRA/QLoRA)
        ft_method: Fine-tuning method
        model_size_bn: Model size in billions of parameters
        spec: Model specification (for hidden_dim and layers)
        
    Returns:
        Throughput multiplier relative to Full Fine-Tuning baseline.
        > 1.0 means faster (LoRA/QLoRA train fewer parameters)
        = 1.0 means baseline (Full FT)
        
    Real-world behavior:
        - Full FT: Updates ALL parameters = baseline (slowest)
        - LoRA: Updates only adapter params (~0.1-1%) = 2-4x faster
        - QLoRA: Same as LoRA but quant overhead reduces speedup slightly
        
    Example:
        >>> get_lora_overhead_factor(64, "QLoRA", 7.0, spec)
        2.5  # 2.5x faster than Full FT with rank=64 on 7B model
    """
    if ft_method not in ["LoRA", "QLoRA"] or rank is None:
        return 1.0  # Full FT baseline - no speedup
    
    # LoRA/QLoRA speedup factors based on empirical benchmarks
    # The speedup comes from:
    # 1. Training only adapter parameters (0.1-1% of model)
    # 2. Smaller gradient computations
    # 3. Reduced optimizer state updates
    #
    # However, there are overheads:
    # 1. Forward pass still processes full model
    # 2. Adapter computations add some latency
    # 3. Higher ranks = more adapter params = less speedup
    
    if model_size_bn <= 10:
        # 7B models - LoRA provides significant speedup
        speedup_map = {
            8: 3.5,     # Very small adapter = big speedup
            16: 3.2,
            32: 2.8,
            64: 2.4,
            128: 2.0,
            256: 1.6,
        }
    elif model_size_bn <= 20:
        # 13B models
        speedup_map = {
            8: 3.2,
            16: 2.9,
            32: 2.5,
            64: 2.2,
            128: 1.8,
            256: 1.5,
        }
    elif model_size_bn <= 100:
        # 70B models - LoRA speedup is proportionally larger
        # because adapter is even smaller relative to model
        speedup_map = {
            8: 4.0,
            16: 3.6,
            32: 3.0,
            64: 2.5,
            128: 2.0,
            256: 1.6,
        }
    else:
        # 405B+ models - largest relative speedup
        speedup_map = {
            8: 4.5,
            16: 4.0,
            32: 3.4,
            64: 2.8,
            128: 2.2,
            256: 1.8,
        }
    
    # Find or interpolate for the given rank
    if rank in speedup_map:
        return speedup_map[rank]
    
    # Interpolate for ranks not in the map
    ranks = sorted(speedup_map.keys())
    for i in range(len(ranks) - 1):
        if ranks[i] < rank < ranks[i+1]:
            r1, r2 = ranks[i], ranks[i+1]
            v1, v2 = speedup_map[r1], speedup_map[r2]
            # Linear interpolation in log-space for smoother scaling
            import math
            log_rank = math.log(rank)
            log_r1 = math.log(r1)
            log_r2 = math.log(r2)
            ratio = (log_rank - log_r1) / (log_r2 - log_r1)
            return v1 + ratio * (v2 - v1)
    
    # Extrapolate if beyond range
    if rank < ranks[0]:
        return speedup_map[ranks[0]]  # Use smallest rank speedup
    else:
        # For very high ranks (> 256), speedup diminishes toward 1.0
        return max(1.2, speedup_map[ranks[-1]] * 0.9)

def calculate_throughput(
    gpu_config: GPUConfig,
    spec: ModelSpec,
    task: str,
    batch_size: int,
    precision: str = "fp16",
    framework: str = "vllm",
    ft_method: Optional[str] = None,
    rank: Optional[int] = None,
    seq_len: int = 2048
) -> Tuple[float, float, str]:
    """
    Calculate estimated throughput for a GPU configuration.
    
    Throughput varies significantly by:
    1. Quantization (INT8/INT4 Tensor Cores provide 2-4x speedup)
    2. Framework (TensorRT-LLM > vLLM > HuggingFace)
    3. Batch size and GPU architecture
    4. Sequence length (longer sequences reduce throughput)
    5. LoRA rank (for training - higher rank = more overhead)
    
    Args:
        gpu_config: GPU configuration
        spec: Model specification
        task: 'Inference' or 'Training'
        batch_size: Batch size
        precision: Quantization/precision format (e.g., 'fp16', 'int8', 'int4')
        framework: Inference framework ('vllm', 'huggingface', 'tensorrt')
        ft_method: Fine-tuning method ('LoRA', 'QLoRA', 'Full Fine-Tuning')
        seq_len: Sequence length (affects KV cache access and attention compute)
        rank: LoRA rank (only used for LoRA/QLoRA training)
        
    Returns:
        Tuple of (tokens_per_second_per_gpu, total_tokens_per_second, description)
        
    Example:
        >>> # INT4 with TensorRT-LLM is ~4x faster than FP16 HuggingFace
        >>> calc_throughput(h100, llama7b, "Inference", 32, "int4", "tensorrt")
        (3900, 3900, "Batched inference (int4, tensorrt) - 4.2x speedup")
        
        >>> # LoRA rank affects training throughput
        >>> calc_throughput(h100, llama7b, "Training", 16, "nf4", "huggingface", "QLoRA", 64)
        (1600, 1600, "Training throughput (nf4, LoRA r=64, 89% efficiency)")
    """
    gpu_family = get_gpu_family(gpu_config.name)
    model_size_bn = spec.params_bn
    
    # Determine if we should use batched numbers
    use_batched = batch_size >= 8 if task == "Inference" else False
    
    # Get base throughput per GPU (assumes FP16 on vLLM baseline)
    tps_per_gpu = interpolate_throughput(gpu_family, model_size_bn, task, use_batched)
    
    # Apply quantization speedup multiplier (INFERENCE ONLY)
    # INT8/INT4 are significantly faster due to specialized Tensor Cores
    # For training, quantization provides memory savings, not speed improvements
    if task == "Inference":
        quant_speedup = QUANTIZATION_SPEEDUP.get(precision, 1.0)
        tps_per_gpu *= quant_speedup
    else:
        quant_speedup = 1.0  # No speedup for training
        # For QLoRA, there's actually a slight slowdown due to
        # quantization/dequantization overhead during forward pass
        if precision in ['nf4', '4bit', 'int4', 'int8'] and ft_method == "QLoRA":
            tps_per_gpu *= 0.85  # ~15% overhead for quantized training
    
    # Apply framework efficiency multiplier
    # TensorRT-LLM is more optimized than vLLM, HuggingFace is less optimized
    framework_speedup = FRAMEWORK_SPEEDUP.get(framework.lower(), 1.0)
    tps_per_gpu *= framework_speedup
    
    # Apply LoRA/QLoRA speedup if applicable (TRAINING ONLY)
    # LoRA/QLoRA train only adapter parameters = faster than Full FT
    lora_speedup = 1.0  # Default: Full FT baseline (no speedup)
    if task == "Training":
        lora_speedup = get_lora_overhead_factor(rank, ft_method, model_size_bn, spec)
        tps_per_gpu *= lora_speedup
    
    # Calculate combined speedup for description
    combined_speedup = quant_speedup * framework_speedup
    
    # Apply batch scaling for inference
    if task == "Inference" and batch_size > 1:
        if use_batched:
            # Already using batched numbers, apply efficiency factor
            # Remove cap to allow larger batches to increase throughput
            batch_efficiency = (batch_size / 32) ** 0.7
            tps_per_gpu *= batch_efficiency
        else:
            # Single stream numbers, scale by batch with diminishing returns
            batch_efficiency = min(1.0, (batch_size / 8) ** 0.6)
            tps_per_gpu *= batch_efficiency * batch_size
    
    # Apply batch scaling for training
    if task == "Training" and batch_size > 1:
        # Remove cap to allow larger batches to increase throughput
        batch_efficiency = (batch_size / 8) ** 0.7
        tps_per_gpu *= batch_efficiency
    
    # Apply sequence length scaling
    # Longer sequences reduce throughput due to:
    # 1. Increased KV cache memory bandwidth
    # 2. O(n²) attention complexity (though optimized with Flash Attention)
    # 3. More memory pressure reducing effective parallelism
    # Baseline: 2048 tokens, ~15% reduction per doubling of sequence length
    SEQ_LEN_BASELINE = 2048
    if seq_len != SEQ_LEN_BASELINE:
        seq_factor = (SEQ_LEN_BASELINE / seq_len) ** 0.15
        tps_per_gpu *= seq_factor
    
    # Apply multi-GPU communication overhead
    if gpu_config.count > 1:
        if gpu_config.count <= 4:
            comm_efficiency = 0.90
        elif gpu_config.count <= 8:
            comm_efficiency = 0.85
        else:
            comm_efficiency = 0.75
        tps_per_gpu *= comm_efficiency
    
    # Total throughput across all GPUs
    total_tps = tps_per_gpu * gpu_config.count
    
    # Generate description
    if task == "Inference":
        desc = f"{'Batched' if use_batched else 'Single-stream'} inference ({precision}, {framework})"
        if combined_speedup != 1.0:
            desc += f" - {combined_speedup:.1f}x speedup"
    else:
        desc = f"Training throughput ({precision})"
        if ft_method in ["LoRA", "QLoRA"] and rank:
            # Add LoRA rank info and speedup factor
            desc += f", LoRA r={rank}, {lora_speedup:.1f}x vs Full FT"
    
    if gpu_config.count > 1:
        desc += f" ({gpu_config.count}x GPUs, {comm_efficiency:.0%} efficiency)"
    
    return tps_per_gpu, total_tps, desc

def format_time_estimate(total_tokens: int, throughput_tps: float) -> str:
    """
    Format time estimate based on tokens and throughput.
    
    Args:
        total_tokens: Total tokens to process
        throughput_tps: Throughput in tokens per second
        
    Returns:
        Formatted time string
    """
    if throughput_tps <= 0:
        return "N/A"
    
    seconds = total_tokens / throughput_tps
    
    if seconds < 60:
        return f"{seconds:.1f}s"
    elif seconds < 3600:
        return f"{seconds/60:.1f}m"
    elif seconds < 86400:
        return f"{seconds/3600:.1f}h"
    else:
        return f"{seconds/86400:.1f}d"

# =================================================================================================
# Validation and Input Processing
# =================================================================================================

class ValidationError(Exception):
    """Custom exception for validation errors."""
    pass

def validate_inputs(
    seq_len: float,
    batch: float,
    rank: float,
    sample_count: float,
    input_tokens: float,
    output_tokens: float
) -> List[str]:
    """
    Validate user inputs and return list of warnings/errors.
    
    Returns:
        List of validation messages (empty if all valid)
    """
    warnings = []
    
    # Check reasonable ranges
    if seq_len < 128:
        warnings.append("WARNING: Context length < 128 may be too small for most models")
    if seq_len > 100000:
        warnings.append("WARNING: Very large context length will require significant VRAM")
    
    if batch < 1:
        warnings.append("ERROR: Batch size must be at least 1")
    if batch > 512:
        warnings.append("WARNING: Very large batch size may exceed VRAM limits")
    
    # Only validate rank for training with LoRA/QLoRA
    if rank is not None:
        if rank < 4:
            warnings.append("WARNING: LoRA rank < 4 may be too low for effective fine-tuning")
        if rank > 256:
            warnings.append("WARNING: LoRA rank > 256 may be inefficient (diminishing returns)")
    
    if sample_count < 1:
        warnings.append("ERROR: Sample count must be at least 1")
    
    if input_tokens < 1 or output_tokens < 1:
        warnings.append("ERROR: Token counts must be positive")
    
    return warnings

# =================================================================================================
# Model Resolution
# =================================================================================================

# Check if transformers library is available
try:
    from transformers import AutoConfig
    from huggingface_hub import HfApi
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    print("Tip: Install 'huggingface_hub' & 'transformers' for automatic model resolution")

def estimate_architecture(params_bn: float) -> ModelSpec:
    """
    Estimates model architecture based on parameter count.
    
    Args:
        params_bn: Number of parameters in billions
        
    Returns:
        ModelSpec with estimated architecture
    """
    # Rule-of-thumb estimates based on common architectures
    if params_bn < 1:
        layers, heads, kv_heads = 12, 12, 12
    elif params_bn < 4:
        layers, heads, kv_heads = 20, 16, 16
    elif params_bn < 10:
        layers, heads, kv_heads = 32, 32, 8
    elif params_bn < 20:
        layers, heads, kv_heads = 40, 40, 8
    elif params_bn < 50:
        layers, heads, kv_heads = 60, 64, 8
    else:
        layers, heads, kv_heads = 80, 80, 8
    
    return ModelSpec(
        params=int(params_bn * 1e9),
        layers=layers,
        heads=heads,
        kv_heads=kv_heads,
        head_dim=128,
        context=32768
    )

def fetch_hf_config(repo_id: str, token: Optional[str] = None) -> Tuple[ModelSpec, str]:
    """
    Fetches model configuration from Hugging Face Hub.
    
    Args:
        repo_id: Repository ID (e.g., 'meta-llama/Llama-2-7b')
        token: Optional HuggingFace authentication token
        
    Returns:
        Tuple of (ModelSpec, repo_id)
        
    Raises:
        PermissionError: If model is gated
        FileNotFoundError: If model doesn't exist
        Exception: Other errors
    """
    if not TRANSFORMERS_AVAILABLE:
        raise ImportError("transformers library not available")
    
    try:
        config = AutoConfig.from_pretrained(repo_id, trust_remote_code=True, token=token)
        
        # Get parameter count
        params = getattr(config, "num_parameters", None)
        if callable(params):
            params = params()
        
        # Estimate if not available
        if params is None:
            hidden = config.hidden_size
            layers = config.num_hidden_layers
            intermediate = getattr(config, "intermediate_size", hidden * 4)
            params = layers * (4 * hidden * hidden + 3 * hidden * intermediate)
        
        spec = ModelSpec(
            params=params,
            layers=config.num_hidden_layers,
            heads=config.num_attention_heads,
            kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
            head_dim=config.hidden_size // config.num_attention_heads,
            context=getattr(config, "max_position_embeddings", 8192)
        )
        
        return spec, repo_id
        
    except Exception as e:
        err = str(e).lower()
        if "401" in err or "403" in err or "gated" in err:
            raise PermissionError(f"Model '{repo_id}' is gated. Please provide HF token.")
        if "404" in err or "not found" in err:
            raise FileNotFoundError(f"Model '{repo_id}' not found on Hugging Face Hub")
        raise e

def resolve_model(model_name: str, token: Optional[str] = None) -> Tuple[ModelSpec, str, List[str]]:
    """
    Resolves model specification from name or parameter count.
    
    Args:
        model_name: Either HF repo ID or parameter count (e.g., "7B")
        token: Optional HuggingFace token
        
    Returns:
        Tuple of (ModelSpec, source_description, logs)
    """
    logs = []
    
    # Try to parse as parameter count (e.g., "7B", "70B")
    match = re.match(r'^(\d+\.?\d*)\s*[Bb]', model_name.strip())
    if match:
        params_bn = float(match.group(1))
        spec = estimate_architecture(params_bn)
        logs.append(f"Using estimated architecture for {params_bn}B parameters")
        return spec, f"Estimated {params_bn}B model", logs
    
    # Try to fetch from Hugging Face
    if TRANSFORMERS_AVAILABLE:
        try:
            spec, repo = fetch_hf_config(model_name, token)
            logs.append(f"Successfully loaded config from Hugging Face: {repo}")
            logs.append(f"  Parameters: {spec.params_bn:.2f}B, Layers: {spec.layers}, Context: {spec.context}")
            return spec, f"HuggingFace: {repo}", logs
        except PermissionError as e:
            logs.append(f"PERMISSION DENIED: {str(e)}")
            logs.append("   Falling back to estimation...")
        except FileNotFoundError as e:
            logs.append(f"NOT FOUND: {str(e)}")
            logs.append("   Falling back to estimation...")
        except Exception as e:
            logs.append(f"ERROR: Error loading config: {str(e)}")
            logs.append("   Falling back to estimation...")
    
    # Fallback: estimate from common model names
    name_lower = model_name.lower()
    if "405b" in name_lower:
        params_bn = 405
    elif "70b" in name_lower or "72b" in name_lower:
        params_bn = 70
    elif "34b" in name_lower or "32b" in name_lower:
        params_bn = 34
    elif "13b" in name_lower or "14b" in name_lower:
        params_bn = 13
    elif "7b" in name_lower or "8b" in name_lower:
        params_bn = 7
    elif "3b" in name_lower:
        params_bn = 3
    elif "1b" in name_lower or "1.5b" in name_lower:
        params_bn = 1.5
    else:
        params_bn = 7  # Default fallback
        logs.append("WARNING: Could not determine model size, using 7B as default")
    
    spec = estimate_architecture(params_bn)
    logs.append(f"Using estimated architecture for {params_bn}B parameters")
    
    return spec, f"Estimated {params_bn}B model", logs

# =================================================================================================
# VRAM CALCULATION ENGINE
# =================================================================================================
#
# This section contains the core memory estimation algorithms for transformer models.
# All calculations are based on empirically-validated formulas derived from:
#   - Production deployments of LLMs
#   - Memory profiling of training workloads  
#   - Vendor specifications and benchmarks
#   - Academic research on transformer efficiency
#
# Accuracy: ±10-15% for model weights, ±15-20% for dynamic allocations

def calculate_model_weights(spec: ModelSpec, precision: str) -> float:
    """
    Calculate memory footprint of model parameters.
    
    Computes the storage requirement for all model parameters based on the
    specified precision format. This is the base memory requirement before
    considering any runtime allocations.
    
    Formula: 
        memory_gb = (num_parameters × bytes_per_parameter) / (1024³)
    
    Args:
        spec (ModelSpec): Model architecture specification containing parameter count
        precision (str): Precision format (e.g., 'fp16', 'nf4', 'int8')
                        Must be a valid key in PRECISION_MAP
        
    Returns:
        float: Memory requirement in gigabytes (GB)
        
    Example:
        >>> spec = ModelSpec(params=7_000_000_000, ...)  # 7B parameters
        >>> calculate_model_weights(spec, 'fp16')
        13.0  # 7B × 2 bytes / 1024³ ≈ 13 GB
    """
    bytes_per_param = PRECISION_MAP.get(precision, 2.0)  # Default to fp16 if unknown
    return (spec.params * bytes_per_param) / (1024**3)

def calculate_kv_cache(
    spec: ModelSpec,
    batch_size: int,
    seq_len: int,
    precision: str
) -> float:
    """
    Calculate Key-Value cache memory requirement for transformer inference.
    
    The KV cache stores computed key and value vectors from attention layers
    to avoid recomputation during autoregressive generation. This is the primary
    dynamic memory component during inference.
    
    Formula:
        kv_memory = 2 × L × B × S × H_kv × D × P
        
        Where:
            2 = Keys + Values (separate tensors)
            L = Number of layers
            B = Batch size
            S = Sequence length
            H_kv = Number of key/value heads (for GQA/MQA architectures)
            D = Dimension per head
            P = Bytes per element (precision)
    
    Important Notes:
        - For standard Multi-Head Attention: H_kv = H (total heads)
        - For Grouped Query Attention (GQA): H_kv < H
          Example: Llama 3 uses H=32, H_kv=8 (4:1 ratio for efficiency)
        - KV cache typically kept at higher precision (fp16) even when model
          is quantized, as aggressive quantization degrades generation quality
    
    Args:
        spec (ModelSpec): Model architecture specification
        batch_size (int): Number of sequences processed in parallel
        seq_len (int): Maximum sequence length to cache
        precision (str): Precision format for KV cache storage
        
    Returns:
        float: Memory requirement in gigabytes (GB)
        
    Example:
        >>> spec = ModelSpec(layers=32, kv_heads=8, head_dim=128, ...)
        >>> calculate_kv_cache(spec, batch_size=32, seq_len=2048, precision='fp16')
        4.0  # Approximately 4 GB for this configuration
    """
    bytes_per_elem = PRECISION_MAP.get(precision, 2.0)
    
    # KV cache is typically not quantized as aggressively as model weights
    # to maintain generation quality. Force fp16 for low-bit formats.
    if precision in ['nf4', '4bit', 'int4']:
        bytes_per_elem = 2.0  # Override to fp16 for quality preservation
    
    kv_memory_bytes = (
        2                    # Separate K and V tensors
        * spec.layers        # One cache per transformer layer
        * batch_size         # Parallel sequences
        * seq_len            # Tokens per sequence
        * spec.kv_heads      # Key/value heads (may differ from query heads in GQA)
        * spec.head_dim      # Dimension of each attention head
        * bytes_per_elem     # Precision-dependent storage size
    )
    
    return kv_memory_bytes / (1024**3)  # Convert bytes to GB

def calculate_activations(
    spec: ModelSpec,
    batch_size: int,
    seq_len: int,
    precision: str,
    use_checkpointing: bool = True  # Most frameworks use some form of checkpointing
) -> float:
    """
    Calculate activation memory for training (more accurate estimate).
    
    Args:
        spec: Model specification
        batch_size: Batch size
        seq_len: Sequence length
        precision: Precision format
        use_checkpointing: Whether gradient checkpointing is used
        
    Returns:
        Memory in GB
    """
    bytes_per_elem = PRECISION_MAP.get(precision, 2.0)
    hidden_size = spec.heads * spec.head_dim
    
    # Activation memory depends on gradient checkpointing strategy:
    # - No checkpointing: Store all intermediate activations (~34x)
    # - Selective checkpointing: Recompute some activations (~12x)
    # Most modern frameworks use some form of checkpointing by default
    multiplier = 12 if use_checkpointing else 34
    
    activation_bytes = (
        batch_size
        * seq_len
        * hidden_size
        * spec.layers
        * multiplier
        * bytes_per_elem
    )
    
    return activation_bytes / (1024**3)

def calculate_optimizer_states(
    model_weights_gb: float,
    ft_method: str,
    rank: int,
    spec: ModelSpec
) -> Tuple[float, str]:
    """
    Calculate optimizer state memory.
    
    Args:
        model_weights_gb: Model weights in GB
        ft_method: Fine-tuning method
        rank: LoRA rank (used for LoRA/QLoRA)
        spec: Model specification (for calculating adapter size)
        
    Returns:
        Tuple of (memory in GB, description)
    """
    if ft_method == "Full Fine-Tuning":
        # Adam: momentum + variance, both stored at fp32
        # Model weights are fp16 (2 bytes), optimizer states are fp32 (4 bytes each)
        # Total: 2 states × 4 bytes = 8 bytes per param vs 2 bytes for model
        # = 4x model weight size
        optimizer_gb = model_weights_gb * 4
        return optimizer_gb, "Optimizer (Adam - Full FT)"
    else:
        # LoRA/QLoRA: only optimizer states for adapter weights
        # Adapter parameters per layer: 2 matrices (A and B) of size (rank × hidden_dim)
        hidden_dim = spec.heads * spec.head_dim
        adapter_params = 2 * rank * hidden_dim * spec.layers
        
        # Adapter params stored at fp16
        adapter_params_gb = (adapter_params * 2) / (1024**3)
        
        # Adam optimizer states at fp32: 2x params at fp32 = 4x params at fp16
        optimizer_gb = adapter_params_gb * 4
        
        return optimizer_gb, f"Optimizer (LoRA r={rank})"

def calculate_gradients(
    model_weights_gb: float,
    ft_method: str,
    rank: int,
    spec: ModelSpec
) -> Tuple[float, str]:
    """
    Calculate gradient memory.
    
    Args:
        model_weights_gb: Model weights in GB
        ft_method: Fine-tuning method
        rank: LoRA rank
        spec: Model specification (for calculating adapter size)
        
    Returns:
        Tuple of (memory in GB, description)
    """
    if ft_method == "Full Fine-Tuning":
        # Full fine-tuning: gradients for all parameters
        # Gradients typically stored at fp32 for numerical stability
        gradient_gb = model_weights_gb * 2  # fp32 vs fp16
        return gradient_gb, "Gradients (Full FT)"
    else:
        # LoRA/QLoRA: only gradients for adapter weights
        # Calculate actual adapter size based on model architecture
        hidden_dim = spec.heads * spec.head_dim
        adapter_params = 2 * rank * hidden_dim * spec.layers
        
        # Gradients stored at fp16 (same precision as adapter params)
        adapter_gradient_gb = (adapter_params * 2) / (1024**3)
        
        return adapter_gradient_gb, f"Gradients (LoRA r={rank})"

def calculate_vram(
    spec: ModelSpec,
    precision: str,
    batch_size: int,
    seq_len: int,
    task: str,
    framework: str,
    ft_method: Optional[str] = None,
    rank: Optional[int] = None
) -> Tuple[float, float, float, str]:
    """
    Main VRAM calculation function.
    
    Args:
        spec: Model specification
        precision: Precision format
        batch_size: Batch size
        seq_len: Sequence length
        task: 'Inference' or 'Training'
        framework: Framework being used
        ft_method: Fine-tuning method (for training)
        rank: LoRA rank (for LoRA/QLoRA)
        
    Returns:
        Tuple of (total_vram_gb, weights_gb, variable_gb, variable_label, actual_precision)
    """
    # KEY DIFFERENCE: QLoRA vs LoRA vs Full FT
    # QLoRA: Base model stays quantized (e.g., 4-bit), adapters in fp16/bf16
    # LoRA: Base model in fp16/bf16, adapters in fp16/bf16
    # Full FT: Base model MUST be fp16/bf16 (all params trainable)
    
    # Track the actual precision being used (may differ from user selection)
    actual_precision = precision
    
    if task == "Training" and ft_method in ["LoRA", "Full Fine-Tuning"]:
        # LoRA and Full FT require full precision base model
        # Even if user selected quantization, these methods need bf16/fp16
        if precision in ['nf4', '4bit', 'int4', 'int8', 'awq', 'gptq']:
            # Override to bf16
            base_precision = 'bf16'
            actual_precision = 'bf16'  # Track the override
            weights_gb = calculate_model_weights(spec, base_precision)
        else:
            weights_gb = calculate_model_weights(spec, precision)
    else:
        # QLoRA or Inference: use selected precision
        # QLoRA can use quantized base because it's frozen
        weights_gb = calculate_model_weights(spec, precision)
    
    # Calculate KV cache
    # For training, KV cache is always at compute precision (bf16/fp16), not quantized
    kv_precision = "bf16" if task == "Training" else precision
    kv_cache_gb = calculate_kv_cache(spec, batch_size, seq_len, kv_precision)
    
    if task == "Inference":
        # Inference: weights + KV cache + framework overhead
        overhead_gb = FRAMEWORK_OVERHEAD.get(framework, 2.0)
        variable_gb = kv_cache_gb + overhead_gb
        variable_label = f"KV Cache + {framework.upper()} Overhead"
        
    else:  # Training
        # Training: weights + KV + activations + optimizer + gradients
        # IMPORTANT: Activations are always stored at compute precision (bf16/fp16)
        # even for QLoRA, because the forward pass computes at full precision
        activations_gb = calculate_activations(spec, batch_size, seq_len, "bf16")
        optimizer_gb, _ = calculate_optimizer_states(weights_gb, ft_method or "Full Fine-Tuning", rank or 64, spec)
        gradients_gb, _ = calculate_gradients(weights_gb, ft_method or "Full Fine-Tuning", rank or 64, spec)
        
        # Add LoRA adapter weights (small additional memory)
        if ft_method in ["LoRA", "QLoRA"]:
            # LoRA adapters: A and B matrices per layer
            # Each matrix is (rank × hidden_dim), stored at fp16
            hidden_dim = spec.heads * spec.head_dim
            adapter_params = 2 * rank * hidden_dim * spec.layers
            adapter_gb = (adapter_params * 2) / (1024**3)  # fp16
            
            variable_gb = kv_cache_gb + activations_gb + optimizer_gb + gradients_gb + adapter_gb
            variable_label = "KV + Activations + Optimizer + Gradients + LoRA Adapters"
        else:
            variable_gb = kv_cache_gb + activations_gb + optimizer_gb + gradients_gb
            variable_label = "KV + Activations + Optimizer + Gradients"
    
    total_vram_gb = weights_gb + variable_gb
    
    return total_vram_gb, weights_gb, variable_gb, variable_label, actual_precision

# =================================================================================================
# Hardware Recommendation Engine
# =================================================================================================

def generate_analysis_csv(
    # Input parameters
    model_name: str,
    task: str,
    quant: str,
    framework: str,
    batch_size: int,
    seq_len: int,
    ft_method: Optional[str],
    rank: Optional[int],
    sample_count: int,
    input_tokens: int,
    output_tokens: int,
    pricing_tier: str,
    manufacturers: List[str],
    # Calculated values
    spec: 'ModelSpec',
    weights_gb: float,
    variable_gb: float,
    total_vram: float,
    actual_precision: str,
    # GPU recommendations
    gpu_recommendations: List[Dict[str, Any]],
    source: str,
) -> str:
    """
    Generate a detailed CSV containing all analysis parameters, formulas, and recommendations.
    
    Returns:
        CSV content as a string
    """
    import csv
    import io
    from datetime import datetime
    
    output = io.StringIO()
    writer = csv.writer(output)
    
    # ==========================================
    # SECTION 1: HEADER
    # ==========================================
    writer.writerow(["IndiaAI GPU Infrastructure Recommender - Analysis Report"])
    writer.writerow(["Generated", datetime.now().strftime("%Y-%m-%d %H:%M:%S")])
    writer.writerow([])
    
    # ==========================================
    # SECTION 2: INPUT PARAMETERS
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["INPUT PARAMETERS"])
    writer.writerow(["*" * 50])
    writer.writerow(["Parameter", "Value", "Description"])
    writer.writerow(["Model Name", model_name, "User-specified model identifier"])
    writer.writerow(["Task", task, "Inference or Training"])
    writer.writerow(["Quantization", quant, "Precision format for model weights"])
    writer.writerow(["Framework", framework, "Inference/training framework"])
    writer.writerow(["Batch Size", batch_size, "Number of samples processed together"])
    writer.writerow(["Sequence Length", seq_len, "Maximum context length (tokens)"])
    if task == "Training":
        writer.writerow(["Fine-tuning Method", ft_method or "N/A", "Training strategy (QLoRA/LoRA/Full FT)"])
        writer.writerow(["LoRA Rank", rank if ft_method in ["LoRA", "QLoRA"] else "N/A", "Adapter rank for LoRA/QLoRA"])
    writer.writerow(["Sample Count", sample_count, "Total samples to process"])
    writer.writerow(["Input Tokens", input_tokens, "Average input tokens per sample"])
    writer.writerow(["Output Tokens", output_tokens, "Average output tokens per sample"])
    writer.writerow(["Pricing Tier", pricing_tier, "Selected pricing model"])
    writer.writerow(["GPU Manufacturers", ", ".join(manufacturers), "Filtered GPU vendors"])
    writer.writerow([])
    
    # ==========================================
    # SECTION 3: MODEL ARCHITECTURE (Derived)
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["MODEL ARCHITECTURE (Derived)"])
    writer.writerow(["*" * 50])
    writer.writerow(["Parameter", "Value", "Formula / Source"])
    writer.writerow(["Source", source, "How model info was obtained"])
    writer.writerow(["Parameters (Billions)", f"{spec.params_bn:.2f}", "From model config or estimated from name"])
    writer.writerow(["Parameters (Exact)", f"{spec.params:,}", "params_bn × 1,000,000,000"])
    writer.writerow(["Layers", spec.layers, "Number of transformer layers"])
    writer.writerow(["Attention Heads", spec.heads, "Number of query attention heads"])
    writer.writerow(["KV Heads", spec.kv_heads, "Number of key/value heads (GQA)"])
    writer.writerow(["Head Dimension", spec.head_dim, "Dimension per attention head"])
    writer.writerow(["Hidden Size", spec.heads * spec.head_dim, "heads × head_dim"])
    writer.writerow(["Max Context Length", spec.context, "Maximum supported sequence length"])
    writer.writerow([])
    
    # ==========================================
    # SECTION 4: PRECISION PARAMETERS
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["PRECISION PARAMETERS"])
    writer.writerow(["*" * 50])
    writer.writerow(["Parameter", "Value", "Formula / Explanation"])
    
    bytes_per_param = PRECISION_MAP.get(quant, 2.0)
    writer.writerow(["Selected Precision", quant, "User-selected quantization format"])
    writer.writerow(["Bytes per Parameter", bytes_per_param, f"From PRECISION_MAP['{quant}']"])
    writer.writerow(["Actual Precision Used", actual_precision, "May differ for LoRA/Full FT (requires bf16)"])
    
    if actual_precision != quant:
        actual_bytes = PRECISION_MAP.get(actual_precision, 2.0)
        writer.writerow(["Actual Bytes per Param", actual_bytes, f"Overridden to {actual_precision} for training"])
    writer.writerow([])
    
    # ==========================================
    # SECTION 5: VRAM CALCULATION BREAKDOWN
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["VRAM CALCULATION BREAKDOWN"])
    writer.writerow(["*" * 50])
    writer.writerow(["Component", "Value (GB)", "Formula"])
    
    # Model weights calculation
    actual_bytes_per_param = PRECISION_MAP.get(actual_precision, 2.0)
    weights_formula = f"{spec.params_bn:.2f}B params × {actual_bytes_per_param} bytes / 1024³"
    writer.writerow(["Model Weights", f"{weights_gb:.2f}", weights_formula])
    
    # KV Cache calculation
    # For training, KV cache uses bf16 (compute precision), not quantized
    if task == "Training":
        kv_bytes_per_elem = 2.0  # Always bf16 for training
        kv_cache_gb = calculate_kv_cache(spec, batch_size, seq_len, "bf16")
        kv_formula = f"2 × {spec.layers} layers × {batch_size} batch × {seq_len} seq × {spec.kv_heads} kv_heads × {spec.head_dim} head_dim × {kv_bytes_per_elem} bytes (bf16) / 1024³"
    else:
        kv_bytes_per_elem = 2.0 if quant in ['nf4', '4bit', 'int4'] else PRECISION_MAP.get(quant, 2.0)
        kv_cache_gb = calculate_kv_cache(spec, batch_size, seq_len, quant)
        kv_formula = f"2 × {spec.layers} layers × {batch_size} batch × {seq_len} seq × {spec.kv_heads} kv_heads × {spec.head_dim} head_dim × {kv_bytes_per_elem} bytes / 1024³"
    writer.writerow(["KV Cache", f"{kv_cache_gb:.2f}", kv_formula])
    
    if task == "Training":
        # Activations - always at bf16 compute precision, even for QLoRA
        # The forward pass computes at full precision, so activations are stored at bf16
        hidden_size = spec.heads * spec.head_dim
        activations_gb = calculate_activations(spec, batch_size, seq_len, "bf16")
        act_formula = f"{batch_size} × {seq_len} × {hidden_size} hidden × {spec.layers} layers × 12 (checkpointing) × 2 bytes (bf16) / 1024³"
        writer.writerow(["Activations", f"{activations_gb:.2f}", act_formula])
        
        # Optimizer states
        optimizer_gb, opt_desc = calculate_optimizer_states(weights_gb, ft_method or "Full Fine-Tuning", rank or 64, spec)
        if ft_method == "Full Fine-Tuning":
            opt_formula = f"{weights_gb:.2f} GB weights × 4 (Adam: 2 states × fp32)"
        else:
            adapter_params = 2 * (rank or 64) * hidden_size * spec.layers
            opt_formula = f"LoRA adapters ({adapter_params:,} params) × 4 (Adam states)"
        writer.writerow(["Optimizer States", f"{optimizer_gb:.2f}", opt_formula])
        
        # Gradients
        gradients_gb, grad_desc = calculate_gradients(weights_gb, ft_method or "Full Fine-Tuning", rank or 64, spec)
        if ft_method == "Full Fine-Tuning":
            grad_formula = f"{weights_gb:.2f} GB weights × 2 (fp32 gradients)"
        else:
            grad_formula = f"LoRA adapter gradients (fp16)"
        writer.writerow(["Gradients", f"{gradients_gb:.2f}", grad_formula])
        
        # LoRA adapters
        if ft_method in ["LoRA", "QLoRA"]:
            adapter_params = 2 * (rank or 64) * hidden_size * spec.layers
            adapter_gb = (adapter_params * 2) / (1024**3)
            adapter_formula = f"2 × {rank} rank × {hidden_size} hidden × {spec.layers} layers × 2 bytes / 1024³"
            writer.writerow(["LoRA Adapters", f"{adapter_gb:.4f}", adapter_formula])
    else:
        # Framework overhead for inference
        overhead_gb = FRAMEWORK_OVERHEAD.get(framework, 2.0)
        writer.writerow(["Framework Overhead", f"{overhead_gb:.2f}", f"FRAMEWORK_OVERHEAD['{framework}']"])
    
    writer.writerow([])
    raw_total = weights_gb + variable_gb
    buffer_amount = raw_total * 0.1
    total_with_buffer = raw_total * 1.1
    writer.writerow(["TOTAL VRAM (calculated)", f"{raw_total:.2f}", "Sum of all components"])
    writer.writerow(["Safety Buffer (10%)", f"{buffer_amount:.2f}", "Total × 0.10 (used for GPU selection)"])
    writer.writerow(["TOTAL VRAM + BUFFER", f"{total_with_buffer:.2f}", "Total × 1.10 (GPUs must have >= this VRAM)"])
    writer.writerow([])
    
    # ==========================================
    # SECTION 6: GPU RECOMMENDATIONS (Top 10)
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["GPU RECOMMENDATIONS (Top 10 by Cost)"])
    writer.writerow(["*" * 50])
    writer.writerow([
        "Rank", "GPU Configuration", "Total VRAM (GB)", "VRAM/GPU (GB)", "GPU Count",
        "VRAM Utilization (%)", "Throughput (tok/s)", "Price (₹/hr)", 
        "Cost Efficiency (tok/₹)", "TFLOPS", "Bandwidth (GB/s)"
    ])
    
    for i, gpu in enumerate(gpu_recommendations, 1):
        writer.writerow([
            i,
            gpu["name"],
            gpu["vram"],
            gpu["vram_per_gpu"],
            gpu["gpu_count"],
            f"{gpu['vram_util']:.1f}",
            f"{gpu['throughput']:.0f}",
            f"{gpu['price']:.2f}",
            # f"{gpu['cost_efficiency']:.2f}",
            f"{gpu['tflops']:.0f}",
            f"{gpu['bandwidth']:.0f}",
        ])
    
    writer.writerow([])
    
    # ==========================================
    # SECTION 7: FORMULAS REFERENCE
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["FORMULAS REFERENCE"])
    writer.writerow(["*" * 50])
    writer.writerow(["Calculation", "Formula"])
    writer.writerow(["Model Weights (GB)", "num_parameters × bytes_per_param / 1024³"])
    writer.writerow(["KV Cache (GB)", "2 × layers × batch × seq_len × kv_heads × head_dim × bytes_per_elem / 1024³"])
    writer.writerow(["Activations (GB)", "batch × seq_len × hidden_size × layers × multiplier × bytes / 1024³"])
    writer.writerow(["Optimizer States (GB)", "trainable_params × 8 bytes (Adam: momentum + variance at fp32)"])
    writer.writerow(["Gradients (GB)", "trainable_params × bytes_per_grad"])
    writer.writerow(["LoRA Adapter Size", "2 × rank × hidden_dim × layers × 2 bytes"])
    writer.writerow(["VRAM Utilization (%)", "(required_vram / gpu_vram) × 100"])
    # writer.writerow(["Cost Efficiency", "throughput (tok/s) / price (₹/hr)"])
    writer.writerow([])
    
    # ==========================================
    # SECTION 8: PRECISION MAP REFERENCE
    # ==========================================
    writer.writerow(["*" * 50])
    writer.writerow(["PRECISION MAP REFERENCE"])
    writer.writerow(["*" * 50])
    writer.writerow(["Format", "Bytes per Parameter", "Notes"])
    for fmt, bytes_val in PRECISION_MAP.items():
        notes = {
            # "fp32": "Full precision - maximum accuracy",
            # "float32": "Full precision - maximum accuracy",
            "bf16": "Brain Float16 - preferred for training",
            "fp16": "Half precision - standard for inference",
            "nf4": "NormalFloat4 - QLoRA format",
            "4bit": "4-bit quantization",
            "int4": "Integer 4-bit",
            "int8": "Integer 8-bit - good accuracy/size tradeoff",
            "awq": "Activation-aware Weight Quantization (inference-only)",
            "gptq": "GPTQ quantization (inference-only)",
        }.get(fmt, "")
        writer.writerow([fmt, bytes_val, notes])
    writer.writerow([])
    
    writer.writerow(["*" * 50])
    writer.writerow(["END OF REPORT"])
    writer.writerow(["*" * 50])
    
    return output.getvalue()

def recommend_hardware(
    required_vram: float,
    task: str,
    spec: ModelSpec,
    weights_gb: float,
    batch_size: int,
    pricing_tier: str,
    precision: str = "fp16",
    framework: str = "vllm",
    sample_count: int = 0,
    input_tokens: int = 0,
    output_tokens: int = 0,
    ft_method: Optional[str] = None,
    rank: Optional[int] = None,
    manufacturers: Optional[list[str]] = None,
    seq_len: int = 2048,
) -> Tuple[Optional[str], str, str, Dict[str, Any]]:
    """
    Recommend GPU configurations based on VRAM requirements.
    
    Args:
        required_vram: Required VRAM in GB
        task: Task type
        spec: Model specification
        weights_gb: Model weights in GB
        batch_size: Batch size
        pricing_tier: Pricing tier selection
        precision: Quantization/precision format
        framework: Inference framework
        sample_count: Number of samples (for time estimation)
        input_tokens: Input tokens per sample
        output_tokens: Output tokens per sample
        ft_method: Fine-tuning method (for training tasks)
        rank: LoRA rank (for LoRA/QLoRA training)
        manufacturers: List of GPU manufacturers to filter by
        seq_len: Sequence length (affects throughput calculation)
        
    Returns:
        Tuple of (error_message, budget_rec, runner_up_rec, chart_data)
    """
    # Manufacturer filtering
    selected = manufacturers or []
    if not selected:
        selected = ["Nvidia"]
    filtered_gpus = [
        gpu for gpu in GPU_DATABASE
        if get_manufacturer(gpu.name) in selected
    ]

    # VRAM filtering with 10% headroom
    valid_configs = [
        gpu for gpu in filtered_gpus
        if gpu.vram >= required_vram * 1.1
    ]

    if not valid_configs:
        maxvram = max(g.vram for g in filtered_gpus) if filtered_gpus else 0
        errormsg = (
            f'<div class="error-box">'
            f"<h4>ERROR No suitable GPU configuration found</h4>"
            f"<p>Required VRAM <strong>{required_vram:.1f} GB</strong></p>"
            f"<p>The largest available configuration in the selected manufacturers "
            f"has {maxvram} GB VRAM.</p>"
            f"<p><em>Suggestions</em></p>"
            f"<ul>"
            f"<li>Reduce batch size (current {batch_size})</li>"
            f"<li>Use more aggressive quantization</li>"
            f"<li>Consider model sharding across multiple nodes</li>"
            f"<li>Try including more GPU manufacturers</li>"
            f"</ul>"
            f"</div>"
        )
        return errormsg, "", "", {}
    
    # Sort by cost
    valid_by_cost = sorted(valid_configs, key=lambda x: x.get_price(pricing_tier))
    
    # Build chart data for all valid configurations (limit to top 10 by cost for readability)
    chart_configs = valid_by_cost[:10]
    chart_data = {
        "names": [],
        "throughput": [],
        "cost": [],
        "vram_util": [],
        "cost_efficiency": [],  # tokens per rupee
        "gpu_details": [],  # Full details for CSV export
        "short_names": [],  # Shortened names for mobile-friendly chart display
    }
    
    def get_short_name(full_name: str) -> str:
        """Create a shorter display name for charts."""
        # Remove vendor prefix and simplify
        name = full_name.replace("Nvidia ", "").replace("AMD ", "").replace("Intel ", "")
        # Shorten common patterns
        name = name.replace(" SXM ", " ").replace(" NVL ", " ").replace(" PCIe ", " ")
        return name
    
    for config in chart_configs:
        price = config.get_price(pricing_tier)
        vram_util = (required_vram / config.vram) * 100
        _, total_tps, _ = calculate_throughput(
            config, spec, task, batch_size, precision, framework, ft_method, rank, seq_len
        )
        
        chart_data["names"].append(config.name)
        chart_data["short_names"].append(get_short_name(config.name))
        chart_data["throughput"].append(total_tps)
        chart_data["cost"].append(price)
        chart_data["vram_util"].append(vram_util)
        # Cost efficiency: tokens per rupee per hour
        cost_eff = total_tps / price if price > 0 else 0
        chart_data["cost_efficiency"].append(cost_eff)
        
        # Full GPU details for CSV export
        chart_data["gpu_details"].append({
            "name": config.name,
            "vram": config.vram,
            "vram_per_gpu": config.vram_per_gpu,
            "gpu_count": config.count,
            "vram_util": vram_util,
            "throughput": total_tps,
            "price": price,
            "cost_efficiency": cost_eff,
            "tflops": config.tflops,
            "bandwidth": config.bandwidth,
        })
    
    # Generate recommendation cards
    def make_card(title: str, emoji: str, config: Optional[GPUConfig]) -> str:
        if config is None:
            return f"### {emoji} {title}\n*No configuration available*"
        
        price = config.get_price(pricing_tier)
        vram_util = (required_vram / config.vram) * 100
        
        tps_per_gpu, total_tps, throughput_desc = calculate_throughput(
            config, spec, task, batch_size, precision, framework, ft_method, rank, seq_len
        )
        
        time_estimate = ""
        if sample_count > 0 and (input_tokens > 0 or output_tokens > 0):
            total_tokens = sample_count * (input_tokens + output_tokens)
            time_str = format_time_estimate(total_tokens, total_tps)
            time_estimate = f"\n- **Time Estimate:** {time_str} for {sample_count:,} samples"
        
        return f"""
### {emoji} {title}
**{config.name}**
- VRAM: {config.vram} GB ({config.vram_per_gpu:.0f} GB/GPU)
- VRAM Utilization: {vram_util:.1f}%
- Performance: {config.tflops:,.0f} TFLOPS
- Bandwidth: {config.bandwidth:,.0f} GB/s
- **Throughput: {total_tps:,.0f} tokens/sec**
  - Per GPU: {tps_per_gpu:,.0f} tok/s
  - {throughput_desc}{time_estimate}
- **Price: ₹{price:,.2f}/hour** from [IndiaAI price list](https://staging2.pmgatishakti.gov.in/IndiaAICompute/pricelist)
- Daily: ₹{price*24:,.2f} | Monthly: ₹{price*730:,.2f}
        """
    
    budget_rec = make_card("Best Budget", "🥇", valid_by_cost[0] if valid_by_cost else None)
    runner_up_rec = make_card("Budget Runner-up", "🥈", valid_by_cost[1] if len(valid_by_cost) > 1 else None)
    
    return None, budget_rec, runner_up_rec, chart_data

# =================================================================================================
# Get Manufacturer Function
# =================================================================================================

def get_manufacturer(gpu_name: str) -> str:
    """Infer GPU manufacturer from config.name."""
    name = gpu_name.lower()
    if "nvidia" in name:
        return "Nvidia"
    if "amd" in name:
        return "AMD"
    if "intel" in name or "gaudi" in name:
        return "Intel"
    return "Other"

# =================================================================================================
# Main Processing Function
# =================================================================================================

def process_request(
    model_name: str,
    seq_len: float,
    quant: str,
    task: str,
    fw: str,
    ft_method: str,
    rank: float,
    batch: float,
    sample_count: float,
    input_tokens: float,
    output_tokens: float,
    dataset_tier: str,
    manufacturers: list[str],
) -> Tuple[str, str, str, Any, Optional[str]]:
    """
    Main request processing function.
    
    Orchestrates model resolution, VRAM calculation, and hardware recommendation.
    
    Returns:
        Tuple of (report, budget_rec, runner_up_rec, combined_chart, csv_filepath)
    """
    try:
        # Validate inputs
        validation_warnings = validate_inputs(seq_len, batch, rank, sample_count, input_tokens, output_tokens)
        
        # Convert to integers
        seq_len = int(seq_len)
        batch = int(batch)
        rank = int(rank) if rank is not None else None
        sample_count = int(sample_count)
        input_tokens = int(input_tokens)
        output_tokens = int(output_tokens)
        
        # Resolve model
        spec, source, logs = resolve_model(model_name, HF_TOKEN)
        
        # Check if quantization is compatible with task
        if task == "Training" and quant in INFERENCE_ONLY_QUANT:
            quant = 'nf4'
            logs.append(f"WARNING: Auto-switched to 'nf4' ({quant} is inference-only)")
        
        # Set ft_method to None for inference
        if task == "Inference":
            ft_method = None
        
        # Add explanation for training precision
        if task == "Training" and ft_method:
            if ft_method == "QLoRA":
                logs.append(f"INFO: QLoRA - Base model stays quantized ({quant}), adapters in fp16")
            elif ft_method == "LoRA":
                if quant in ['nf4', '4bit', 'int4', 'int8']:
                    logs.append(f"INFO: LoRA - Base model upgraded to bf16 (LoRA requires full precision)")
                else:
                    logs.append(f"INFO: LoRA - Base model in {quant}, adapters in fp16")
            elif ft_method == "Full Fine-Tuning":
                if quant in ['nf4', '4bit', 'int4', 'int8', 'awq', 'gptq']:
                    logs.append(f"INFO: Full FT - Base model upgraded to bf16 (training requires full precision)")
                else:
                    logs.append(f"INFO: Full Fine-Tuning - Training all parameters at {quant}")
        
        # Calculate VRAM
        total_vram, weights_gb, variable_gb, variable_label, actual_precision = calculate_vram(
            spec, quant, batch, seq_len, task, fw, ft_method, rank
        )
        
        # Check token length warning
        tokens_per_sample = input_tokens + output_tokens
        if tokens_per_sample > seq_len:
            logs.append(f"WARNING: Sample length ({tokens_per_sample}) exceeds context length ({seq_len})")
        
        # Generate hardware recommendations
        error, budget_rec, runner_up_rec, chart_data = recommend_hardware(
            total_vram, task, spec, weights_gb, batch, dataset_tier, actual_precision, fw,
            sample_count, input_tokens, output_tokens, ft_method, rank, manufacturers=manufacturers,
            seq_len=seq_len,
        )
        
        if error:
            return error, "", "", None, None
        
        # Generate report
        report = f"""
### 📊 Analysis Report

**Model Information:**
- Source: {source}
- Parameters: **{spec.params_bn:.2f}B**
- Layers: {spec.layers} | Heads: {spec.heads} | KV Heads: {spec.kv_heads}
- Context Length: {spec.context:,} tokens

**VRAM Breakdown:**
- Model Weights: **{weights_gb:.1f} GB** ({quant})
- {variable_label}: **{variable_gb:.1f} GB**
- **Total Required: {total_vram:.1f} GB**

**Configuration:**
- Task: {task}
- Framework: {fw}
- Batch Size: {batch}
- Sequence Length: {seq_len:,}
{f"- Fine-tuning: {ft_method} (rank={rank})" if ft_method in ["LoRA", "QLoRA"] else f"- Fine-tuning: {ft_method}" if ft_method else ""}

---
{"<br>".join(f"*{log}*" for log in logs)}
{"<br>".join(f'<div class="warning-box">{w}</div>' for w in validation_warnings) if validation_warnings else ""}

*ℹ️ NOTE: VRAM estimates include 10% safety buffer. Actual usage may vary ±15-20% based on framework optimizations, model architecture details, and runtime conditions. Throughput estimates are based on empirical benchmarks and may vary in production.*
        """
        
        # Create grouped multi-bar chart using Plotly
        combined_chart = None
        
        if chart_data and chart_data.get("names"):
            import plotly.graph_objects as go
            
            # Normalize values to percentages for comparison (0-100 scale)
            max_throughput = max(chart_data["throughput"]) if chart_data["throughput"] else 1
            max_cost = max(chart_data["cost"]) if chart_data["cost"] else 1
            
            throughput_normalized = [(t / max_throughput) * 100 for t in chart_data["throughput"]]
            cost_normalized = [(c / max_cost) * 100 for c in chart_data["cost"]]
            
            combined_chart = go.Figure(data=[
                go.Bar(
                    name='Throughput',
                    x=chart_data["short_names"],  # Use shorter names for mobile
                    y=throughput_normalized,
                    marker_color='#22c55e',
                    hovertemplate='%{customdata[0]}<br>Throughput: %{customdata[1]:,.0f} tok/s<extra></extra>',
                    customdata=list(zip(chart_data["names"], chart_data["throughput"]))  # Full name in hover
                ),
                go.Bar(
                    name='Cost',
                    x=chart_data["short_names"],
                    y=cost_normalized,
                    marker_color='#3b82f6',
                    hovertemplate='%{customdata[0]}<br>Cost: ₹%{customdata[1]:,.2f}/hr<extra></extra>',
                    customdata=list(zip(chart_data["names"], chart_data["cost"]))
                ),
                go.Bar(
                    name='VRAM Util',
                    x=chart_data["short_names"],
                    y=chart_data["vram_util"],
                    marker_color='#a855f7',
                    hovertemplate='%{customdata}<br>VRAM Util: %{y:.1f}%<extra></extra>',
                    customdata=chart_data["names"]
                )
            ])
            
            combined_chart.update_layout(
                title_text="GPU Comparison (Top 10 by Cost)",
                xaxis_title="",  # Remove redundant title - GPU names are self-explanatory
                yaxis_title="Normalized Value (% of max)",
                barmode='group',
                height=520,  # Slightly taller for legend spacing
                autosize=True,  # Enable responsive sizing
                xaxis_tickangle=-45,
                margin=dict(b=160, t=50, l=50, r=20),  # More bottom margin for legend
                legend=dict(
                    orientation="h",  # Horizontal legend
                    yanchor="top",
                    y=-0.38,  # Move further below to avoid overlap
                    xanchor="center",
                    x=0.5,
                    bgcolor="rgba(255,255,255,0.9)",
                    bordercolor="rgba(0,0,0,0.1)",
                    borderwidth=1,
                    font=dict(size=12),
                    itemsizing='constant',
                    traceorder='normal',
                ),
                font=dict(size=10),
                title_font=dict(size=13),
                # Make x-axis labels more readable
                xaxis=dict(
                    tickfont=dict(size=9),
                ),
                yaxis=dict(
                    tickfont=dict(size=9),
                    title_font=dict(size=11),
                ),
            )
        
        # Generate CSV content and save to file
        csv_filepath = None
        if chart_data and chart_data.get("gpu_details"):
            csv_content = generate_analysis_csv(
                model_name=model_name,
                task=task,
                quant=quant,
                framework=fw,
                batch_size=batch,
                seq_len=seq_len,
                ft_method=ft_method,
                rank=rank,
                sample_count=sample_count,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                pricing_tier=dataset_tier,
                manufacturers=manufacturers,
                spec=spec,
                weights_gb=weights_gb,
                variable_gb=variable_gb,
                total_vram=total_vram,
                actual_precision=actual_precision,
                gpu_recommendations=chart_data["gpu_details"],
                source=source,
            )
            
            # Save CSV to a temporary file
            import tempfile
            import os
            
            # Create a filename based on model and timestamp
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            safe_model_name = model_name.replace("/", "_").replace(" ", "_")[:30]
            filename = f"gpu_analysis_{safe_model_name}_{timestamp}.csv"
            
            # Save to temp directory
            csv_filepath = os.path.join(tempfile.gettempdir(), filename)
            with open(csv_filepath, 'w', newline='', encoding='utf-8') as f:
                f.write(csv_content)
        
        return report, budget_rec, runner_up_rec, combined_chart, csv_filepath
        
    except Exception as e:
        error_msg = f"""
<div class="error-box">
<h4>ERROR: Error Processing Request</h4>
<p>{str(e)}</p>
</div>
        """
        return error_msg, "", "", None, None

# =================================================================================================
# UI Event Handlers
# =================================================================================================

def update_ui_on_task(task_val: str):
    """Update UI elements when task changes."""
    if task_val == "Training":
        return [
            gr.update(choices=["nf4", "int4", "bf16", "fp16", "int8"], value="nf4"),
            gr.update(choices=["huggingface"], value="huggingface"),
            gr.update(visible=True)
        ]
    else:
        return [
            gr.update(choices=["nf4", "int4", "bf16", "fp16", "int8", "awq", "gptq"], value="nf4"),
            gr.update(choices=["vllm"], value="vllm"),
            gr.update(visible=False)
        ]

def update_rank_visibility(ft_val: str):
    """Update rank input visibility based on fine-tuning method."""
    if ft_val in ["LoRA", "QLoRA"]:
        return gr.update(visible=True)
    return gr.update(visible=False)

# =================================================================================================
# Gradio Interface
# =================================================================================================

def create_interface() -> gr.Blocks:
    """Create and configure the Gradio interface."""
    
    # Try to create with theme, fallback to basic if not supported
    try:
        demo = gr.Blocks(title="IndiaAI GPU Infrastructure Recommender")
    except Exception as e:
        print(f"Note: Using basic Gradio configuration: {e}")
        demo = gr.Blocks()
    
    with demo:
        # Inject custom CSS
        gr.HTML(CUSTOM_CSS)
        
        # Header
        gr.Markdown("""
        # IndiaAI GPU Infrastructure Recommender
        
        Calculate VRAM requirements and get optimal GPU recommendations for your AI workloads using [IndiaAI's price list](https://staging2.pmgatishakti.gov.in/IndiaAICompute/pricelist).
        """)
        
        with gr.Row():
            # Left column: Inputs
            with gr.Column(scale=0.30):
                gr.Markdown("### 🔧 Configuration")
                
                model_name = gr.Dropdown(
                    choices=MODEL_CHOICES,
                    label="Model Name or Size",
                    value="mistralai/Mistral-7B-Instruct-v0.3",
                    allow_custom_value=True,
                    info="Select a model or enter custom (e.g., '7B', 'meta-llama/...')"
                )
                
                with gr.Row():
                    task = gr.Radio(
                        ["Inference", "Training"],
                        label="Task",
                        value="Inference",
                        info="Select your use case"
                    )
                    fw = gr.Dropdown(
                        ["vllm"], #, "huggingface"],
                        label="Framework",
                        value="vllm",
                        info="Framework for the task"
                    )
                    quant = gr.Dropdown(
                        ["nf4", "int4", "bf16", "fp16", "int8", "awq", "gptq"],
                        label="Quantization",
                        value="nf4",
                        info="Precision format"
                    )
                
                with gr.Row(visible=False) as training_row:
                    ft_method = gr.Dropdown(
                        ["QLoRA", "LoRA", "Full Fine-Tuning"],
                        label="Fine-tuning Method",
                        value="QLoRA",
                        info="Training strategy"
                    )
                    rank = gr.Number(
                        label="LoRA Rank (r)",
                        value=16,
                        visible=True,
                        info="Adapter rank for LoRA/QLoRA"
                    )
                
                gr.Markdown("### 📐 Workload Parameters")
                
                with gr.Row():
                    seq_len = gr.Number(
                        label="Max Context Length",
                        value=1024,
                        info="Maximum sequence length (buffer)"
                    )
                    batch = gr.Number(
                        label="Batch Size",
                        value=16,
                        info="Number of samples processed together"
                    )
                
                with gr.Row():
                    sample_count = gr.Number(
                        label="Samples",
                        value=10000,
                        scale=1,
                        info="Number of samples"
                    )
                    input_tokens = gr.Number(
                        label="Input Tokens",
                        value=300,
                        scale=1,
                        info="Avg input length"
                    )
                    output_tokens = gr.Number(
                        label="Output Tokens",
                        value=100,
                        scale=1,
                        info="Avg output length"
                    )
                
                gr.Markdown("### 💰 Cost Estimation")
                
                with gr.Row():
                    dataset_tier = gr.Dropdown(
                        choices=["On Demand", "1 Month Reserved", "6 Month Reserved", "12 Month Reserved"],
                        label="Pricing Tier",
                        value="On Demand",
                        scale=2,
                        info="Select pricing model",
                    )
                    manufacturers = gr.CheckboxGroup(
                        choices=["Nvidia", "AMD", "Intel"],
                        label="GPU Manufacturers",
                        value=["Nvidia", "AMD", "Intel"],  # default: all
                        info="Filter recommendations by GPU manufacturer",
                    )
                
                btn = gr.Button("Calculate Requirements", variant="primary")
            
            # Right column: Results
            with gr.Column(scale=1):
                gr.HTML("<h3 style='text-align: center; margin-top: 0px;'>Recommendations</h3>")
                
                with gr.Row():
                    with gr.Column(elem_classes=["budget-box"]):
                        rec_out_1 = gr.Markdown()
                    with gr.Column(elem_classes=["runner-box"]):
                        rec_out_2 = gr.Markdown()
                
                # Download button for CSV export
                with gr.Row():
                    csv_download = gr.File(
                        label="📥 Download Analysis (CSV)",
                        visible=True,
                        file_count="single",
                        type="filepath",
                        interactive=False,
                    )
                
                # Bar chart comparison section
                with gr.Accordion("📊 GPU Comparison (Top 10 by Cost)", open=False):
                    combined_plot = gr.Plot(label="GPU Comparison")

                with gr.Accordion("📈 Details", open=False):
                    report_out = gr.Markdown()
        
        # Footer
        gr.Markdown("""
        ---
        💡 **Tips:**
        - Start with smaller batch sizes for testing
        - QLoRA is most memory-efficient for training
        - Consider reserved instances for long-term workloads
        - VRAM estimates include 10% safety buffer
        - Higher "Cost Efficiency" means more tokens processed per rupee spent
        """)
        
        # Event handlers
        task.change(
            fn=update_ui_on_task,
            inputs=task,
            outputs=[quant, fw, training_row]
        )
        
        ft_method.change(
            fn=update_rank_visibility,
            inputs=ft_method,
            outputs=rank
        )
        
        btn.click(
            fn=process_request,
            inputs=[
                model_name, seq_len, quant, task, fw,
                ft_method, rank, batch, sample_count,
                input_tokens, output_tokens, dataset_tier, manufacturers,
            ],
            outputs=[report_out, rec_out_1, rec_out_2, combined_plot, csv_download]
        )
    
    return demo

# =================================================================================================
# Main Entry Point
# =================================================================================================

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=False)