components/sidebar.py

import streamlit as st
from mock_data.sample_data import get_base_models, get_target_hardware, get_hardware_specs

def render_sidebar():
    """Render the sidebar with model selection and optimization controls."""

    with st.sidebar:
        st.header("Configuration")

        # Model Selection section
        st.subheader("🤖 Base LLM Model")
        st.caption("Select the foundation model to optimize for deployment")

        # Base model dropdown
        base_models = get_base_models()
        selected_model = st.selectbox(
            "Model",
            base_models,
            index=0,
            key="base_model",
            help="Recent open-source LLMs optimizable for single-GPU deployment"
        )

        st.markdown("---")

        # Optimization section
        st.subheader("Settings")

        st.markdown("**🎯 Stopping Criteria**")
        st.caption("Optimization stops when time limit is reached")

        # Wall clock time input
        max_hours = st.number_input(
            "Max Time (hours)",
            min_value=0.5,
            max_value=168.0,  # 1 week
            value=24.0,
            step=0.5,
            help="Maximum wall clock time for optimization",
            key="max_hours"
        )

        st.markdown("---")

        # Target Hardware selector
        st.markdown("**⚙️ Target Hardware**")
        st.caption("Select deployment hardware platform")

        # Get hardware options
        hardware_dict = get_target_hardware()

        # Flatten hardware options for selectbox with category prefixes
        hardware_options = []
        for category, devices in hardware_dict.items():
            for device in devices:
                hardware_options.append(f"{category}: {device}")

        selected_hardware_full = st.selectbox(
            "Hardware Platform",
            options=hardware_options,
            index=0,  # Default to H100
            help="Target hardware for LLM deployment optimization",
            key="target_hardware"
        )

        # Extract just the device name (remove category prefix)
        selected_hardware = selected_hardware_full.split(": ", 1)[1]

        # Get and display hardware specs
        hw_specs = get_hardware_specs(selected_hardware)
        hw_category = hw_specs.get('category', 'datacenter')

        # Display hardware specifications with compact styling
        if 'tflops' in hw_specs:
            compute_str = f"{hw_specs['tflops']} TF"
        elif 'tops' in hw_specs:
            compute_str = f"{hw_specs['tops']} TOPS"
        else:
            compute_str = "N/A"

        st.markdown(f"""
        <div style='display: flex; justify-content: space-between; padding: 0.5rem 0; font-size: 0.85rem;'>
            <div style='text-align: center; flex: 1;'>
                <div style='color: rgba(255,255,255,0.6); font-size: 0.75rem;'>VRAM</div>
                <div style='color: rgba(255,255,255,0.95); font-weight: 600; font-size: 0.9rem;'>{hw_specs['vram']} GB</div>
            </div>
            <div style='text-align: center; flex: 1;'>
                <div style='color: rgba(255,255,255,0.6); font-size: 0.75rem;'>Compute</div>
                <div style='color: rgba(255,255,255,0.95); font-weight: 600; font-size: 0.9rem;'>{compute_str}</div>
            </div>
            <div style='text-align: center; flex: 1;'>
                <div style='color: rgba(255,255,255,0.6); font-size: 0.75rem;'>Arch</div>
                <div style='color: rgba(255,255,255,0.95); font-weight: 600; font-size: 0.9rem;'>{hw_specs['compute']}</div>
            </div>
        </div>
        """, unsafe_allow_html=True)

        # Number of hardware units for parallel exploration
        st.markdown("**📊 Parallel Search Cluster**")
        st.caption("Each unit explores models in parallel")
        num_units = st.number_input(
            "Number of Units for Optimization",
            min_value=1,
            max_value=1000,
            value=10,
            step=1,
            help="Number of hardware units available for parallel model exploration",
            key="num_units"
        )

        # For backward compatibility
        num_gpus = num_units

        # Calculate exploration metrics based on time, units, and hardware capability
        # Hardware-aware throughput: datacenter GPUs are faster than edge devices
        if hw_category == 'datacenter':
            models_per_hour_per_unit = 2.5  # Fast datacenter GPUs
        elif hw_category == 'consumer':
            models_per_hour_per_unit = 2.0  # Consumer GPUs
        elif hw_category == 'edge':
            models_per_hour_per_unit = 1.0  # Slower edge devices
        else:
            models_per_hour_per_unit = 2.0  # Default

        models_per_hour = num_units * models_per_hour_per_unit

        # Calculate expected models based on time limit and GPU count
        actual_models = int(max_hours * models_per_hour)

        # Set num_models for backward compatibility
        num_models = actual_models

        # Convert to readable time format
        def format_time(hours):
            if hours < 1:
                return f"{int(hours * 60)} min"
            elif hours < 24:
                h = int(hours)
                m = int((hours - h) * 60)
                return f"{h}h {m}m" if m > 0 else f"{h}h"
            else:
                d = int(hours / 24)
                rh = int(hours % 24)
                return f"{d}d {rh}h"

        actual_time_str = format_time(max_hours)

        # Show exploration metrics
        st.markdown("---")
        st.markdown("**⚡ Exploration Estimate**")

        col1, col2 = st.columns(2)
        with col1:
            st.metric(
                label="Expected Models",
                value=f"{actual_models:,}",
                help="Models that will be explored in the time limit"
            )
        with col2:
            st.metric(
                label="Optimization Time",
                value=actual_time_str,
                help="Time optimization will run"
            )

        st.caption(f"💡 Throughput: {int(models_per_hour)} models/hour | Target: {selected_hardware}")

        st.markdown("---")

        # Optimization objectives (dynamic defaults based on hardware category)
        st.subheader("Objectives")

        # Show hardware-specific optimization focus
        if hw_category == 'edge':
            st.caption("🎯 Edge Focus: Size, Latency, Power Efficiency")
        elif hw_category == 'consumer':
            st.caption("🎯 Consumer Focus: VRAM, Throughput, Accuracy")
        elif hw_category == 'datacenter':
            st.caption("🎯 Datacenter Focus: Throughput, Batch Size, Accuracy")
        else:
            st.caption("Select what to optimize")

        # Set default objectives based on hardware category
        # Edge Hardware: Model Size (CRITICAL), Latency (HIGH), Power (HIGH), Accuracy (MEDIUM)
        # Consumer GPUs: VRAM (CRITICAL), Throughput (HIGH), Accuracy (HIGH), Cost (MEDIUM)
        # Datacenter: Throughput (CRITICAL), Batch Size (HIGH), Accuracy (HIGH)

        if hw_category == 'edge':
            default_size = True       # CRITICAL for edge
            default_latency = True    # HIGH for edge
            default_energy = True     # HIGH for edge
            default_memory = True     # Important for edge
            default_throughput = False
            default_cost = False
            default_batch = False
        elif hw_category == 'consumer':
            default_memory = True     # CRITICAL (VRAM) for consumer
            default_throughput = True # HIGH for consumer
            default_size = True       # HIGH for consumer
            default_cost = True       # MEDIUM for consumer
            default_latency = False
            default_energy = False
            default_batch = False
        elif hw_category == 'datacenter':
            default_throughput = True # CRITICAL for datacenter
            default_batch = True      # HIGH for datacenter
            default_size = True       # Always useful
            default_cost = True       # Important for cloud
            default_memory = False
            default_latency = False
            default_energy = False
        else:
            # Default balanced approach
            default_size = True
            default_throughput = True
            default_memory = True
            default_cost = False
            default_latency = False
            default_energy = False
            default_batch = False

        # Accuracy is always enabled (primary objective)
        accuracy_obj = st.checkbox(
            "🎯 Accuracy",
            value=True,
            disabled=True,
            help="Primary objective - always optimized for model quality",
            key="obj_accuracy"
        )

        # Model size - CRITICAL for edge
        size_label = "📦 Model Size" + (" 🔴" if hw_category == 'edge' else "")
        size_obj = st.checkbox(
            size_label,
            value=default_size,
            help="Minimize model size (MB/GB) - CRITICAL for edge deployment",
            key="obj_size"
        )

        # VRAM / Memory footprint - CRITICAL for consumer
        memory_label = "💾 VRAM Usage" + (" 🔴" if hw_category == 'consumer' else "")
        memory_obj = st.checkbox(
            memory_label,
            value=default_memory,
            help="Minimize VRAM/memory usage - CRITICAL for consumer GPUs with limited memory",
            key="obj_memory"
        )

        # Throughput - CRITICAL for datacenter, HIGH for consumer
        throughput_label = "⚡ Throughput" + (" 🔴" if hw_category == 'datacenter' else "")
        throughput_obj = st.checkbox(
            throughput_label,
            value=default_throughput,
            help="Maximize tokens/sec - CRITICAL for datacenter, HIGH for consumer",
            key="obj_throughput"
        )

        # Latency - HIGH for edge
        latency_label = "⏱️ Latency" + (" 🟠" if hw_category == 'edge' else "")
        latency_obj = st.checkbox(
            latency_label,
            value=default_latency,
            help="Minimize ms/token - HIGH priority for edge real-time applications",
            key="obj_latency"
        )

        # Energy efficiency - HIGH for edge
        energy_label = "🔋 Power Consumption" + (" 🟠" if hw_category == 'edge' else "")
        energy_obj = st.checkbox(
            energy_label,
            value=default_energy,
            help="Minimize power (watts) - HIGH priority for edge and battery-powered devices",
            key="obj_energy"
        )

        # Inference cost - MEDIUM for consumer/datacenter
        cost_obj = st.checkbox(
            "💰 Cost per Inference",
            value=default_cost,
            help="Minimize cost per 1M requests - important for cloud deployment",
            key="obj_cost"
        )

        # Batch size capacity - HIGH for datacenter
        batch_label = "📊 Batch Size" + (" 🟠" if hw_category == 'datacenter' else "")
        batch_obj = st.checkbox(
            batch_label,
            value=default_batch,
            help="Maximize concurrent batch processing - HIGH priority for datacenter multi-tenant serving",
            key="obj_batch"
        )

        st.markdown("---")

        # Optimization techniques
        st.subheader("Optimization Techniques")
        st.caption("Select techniques to apply during optimization")

        # Quantization - always useful for LLMs
        quantization_enabled = st.checkbox(
            "🔢 Quantization",
            value=True,
            help="Reduce precision (INT8, INT4) to decrease model size and increase speed",
            key="tech_quantization"
        )
        if quantization_enabled:
            with st.expander("📋 Quantization Hyperparameters", expanded=False):
                st.markdown("""
                **Precision Options:**
                - **FP16**: 16-bit float (2x smaller, 1.5-2x faster)
                - **INT8**: 8-bit integer (4x smaller, 2-4x faster)
                - **INT4**: 4-bit integer (8x smaller, 3-6x faster)
                - **NF4/GPTQ**: 4-bit optimized (8x smaller, near-FP16 quality)

                **Key Parameters:**
                - `bits`: 4, 8, 16 (quantization bit-width)
                - `calibration_samples`: 128-1024 samples
                - `group_size`: 32, 64, 128 (for GPTQ/AWQ)
                - `per_channel`: True/False (channel-wise vs layer-wise)
                - `symmetric`: True/False (symmetric vs asymmetric range)
                - `scheme`: PTQ (post-training) vs QAT (quantization-aware training)
                """)

        # Pruning - remove less important weights
        pruning_enabled = st.checkbox(
            "✂️ Pruning",
            value=True,
            help="Remove redundant weights to reduce model size",
            key="tech_pruning"
        )
        if pruning_enabled:
            with st.expander("📋 Pruning Hyperparameters", expanded=False):
                st.markdown("""
                **Sparsity Ratio:** 10%-90% (typical: 30-60%)
                - Low (10-30%): <1% accuracy loss, 1.2-1.5x speedup
                - Medium (30-60%): 1-3% accuracy loss, 1.5-2x speedup
                - High (60-90%): >3% accuracy loss, 2-5x speedup

                **Pruning Methods:**
                - **Magnitude**: Remove weights with smallest absolute values
                - **Structured**: Remove entire channels/filters (better hardware efficiency)
                - **Unstructured**: Remove individual weights (higher compression)
                - **Movement**: Track weight importance during training

                **Schedule:**
                - `initial_sparsity`: 0% (start dense)
                - `final_sparsity`: 30-60% (target sparsity)
                - `pruning_frequency`: Every 100-1000 steps
                - `recovery_epochs`: 3-10 epochs for fine-tuning
                """)

        # Knowledge Distillation - train smaller model from larger one
        distillation_enabled = st.checkbox(
            "🧪 Knowledge Distillation",
            value=False,
            help="Train a smaller student model using a larger teacher model",
            key="tech_distillation"
        )
        if distillation_enabled:
            with st.expander("📋 Distillation Hyperparameters", expanded=False):
                st.markdown("""
                **Temperature (T):** 1.0-4.0
                - T=1.0: Hard targets (standard cross-entropy)
                - T=2.0-3.0: Soft targets (typical for distillation)
                - T=4.0+: Very soft (extreme compression scenarios)

                **Alpha (α):** 0.1-0.9 (loss weighting)
                - Loss = α × distill_loss + (1-α) × task_loss
                - α=0.5-0.7: Typical, balanced approach
                - α=0.8-0.9: Prioritize teacher knowledge
                - α=0.1-0.3: Prioritize task performance

                **Student/Teacher Size Ratio:** 0.1x - 0.5x
                - 0.1x: 10% of teacher size (aggressive compression)
                - 0.25x: 25% of teacher size (typical)
                - 0.5x: 50% of teacher size (conservative)

                **Training Epochs:** 5-20 epochs (typically 2-3x longer than standard training)

                **Learning Rate:** 1e-5 to 5e-5 (lower than training from scratch)
                """)

        # LoRA/QLoRA - parameter-efficient fine-tuning
        lora_enabled = st.checkbox(
            "🎛️ LoRA/QLoRA",
            value=True,
            help="Low-rank adaptation for efficient fine-tuning and deployment",
            key="tech_lora"
        )
        if lora_enabled:
            with st.expander("📋 LoRA Hyperparameters", expanded=False):
                st.markdown("""
                **Rank (r):** 4, 8, 16, 32, 64, 128
                - r=4-8: 0.1-0.3M params, minimal overhead, good for simple tasks
                - r=16-32: 0.5-2M params, balanced performance/efficiency
                - r=64-128: 2-8M params, higher capacity for complex tasks

                **Alpha (α):** Scaling factor
                - Typical: α = r (equal scaling) or α = 2×r
                - Controls magnitude of LoRA updates
                - Higher α = stronger adaptation

                **Learning Rate:** 1e-4 to 5e-4 (10-50x higher than full fine-tuning)

                **Dropout:** 0.0-0.1 (LoRA layer dropout for regularization)

                **Target Modules:** q_proj, v_proj, k_proj, o_proj, gate_proj, up_proj, down_proj
                - Attention layers (q,k,v,o): Always recommended
                - FFN layers (gate, up, down): Optional for better quality
                """)

        # Flash Attention - optimize attention mechanism
        flash_attention_enabled = st.checkbox(
            "⚡ Flash Attention",
            value=True,
            help="Optimized attention computation for faster inference",
            key="tech_flash_attention"
        )
        if flash_attention_enabled:
            with st.expander("📋 Flash Attention Configuration", expanded=False):
                st.markdown("""
                **Version:**
                - **Flash Attention 2**: 2-4x faster, 2x more memory efficient
                - **Flash Attention 3** (H100 only): Up to 1.5-2x faster than v2

                **Memory Complexity:**
                - Standard attention: O(N²) memory (N = sequence length)
                - Flash Attention: O(N) memory (IO-aware tiling)

                **Key Features:**
                - **Tiling**: Blocks fit in SRAM (faster than HBM)
                - **Recomputation**: Trade compute for memory (enabling longer sequences)
                - **Causal masking**: Efficient for autoregressive generation
                - **Sequence length**: Supports up to 128K+ tokens

                **Hardware Requirements:**
                - FlashAttn v2: NVIDIA Ampere+ (A100, H100, RTX 30xx/40xx series)
                - FlashAttn v3: NVIDIA Hopper only (H100, H200)
                - Requires CUDA 11.6+ and compute capability 8.0+
                """)

        # Kernel Fusion - fuse operations
        kernel_fusion_enabled = st.checkbox(
            "🔗 Kernel Fusion",
            value=True,
            help="Fuse adjacent operations to reduce memory bandwidth",
            key="tech_kernel_fusion"
        )
        if kernel_fusion_enabled:
            with st.expander("📋 Kernel Fusion Patterns", expanded=False):
                st.markdown("""
                **Common Fusion Patterns:**
                - **GELU Fusion**: GELU + Linear (reduces memory traffic)
                - **LayerNorm Fusion**: LayerNorm + Linear/Add (common in transformers)
                - **Attention Fusion**: Q/K/V projection + softmax + output projection
                - **Residual Fusion**: Add + LayerNorm (skip connections)
                - **FFN Fusion**: Gate + Up + Down projection (MLP blocks)

                **Mixed Precision:**
                - **FP32**: Full precision (baseline, highest accuracy)
                - **FP16**: Half precision (2x faster, 2x less memory)
                - **BF16**: Brain float16 (better numerical stability than FP16)
                - **TF32**: TensorFloat-32 (Ampere+, good accuracy/speed balance)

                **Optimization Backends:**
                - **TensorRT**: NVIDIA optimized (best for deployment)
                - **torch.compile()**: PyTorch 2.0+ native compiler
                - **ONNX Runtime**: Cross-platform inference
                - **Triton**: Custom kernel generation (flexible)

                **Expected Speedup:** 1.3-2x over naive implementation
                """)


        st.markdown("---")

        # Start Optimization button with dynamic state
        is_running = st.session_state.get('optimization_running', False)
        is_completed = st.session_state.get('optimization_completed', False)

        if is_running:
            # Show disabled button with "Optimizing..." text
            st.button(
                "⚙️ Optimizing...",
                use_container_width=True,
                type="primary",
                key="start_optimization_disabled",
                disabled=True
            )
            # Show progress bar
            progress_step = st.session_state.get('progress_step', 0)
            progress_pct = min(progress_step / 20, 1.0)
            st.progress(progress_pct, text=f"Progress: {int(progress_pct * 100)}%")
            start_button = False
        elif is_completed:
            # Show disabled button with "Completed" text
            st.button(
                "✅ Optimization Complete",
                use_container_width=True,
                type="primary",
                key="start_optimization_completed",
                disabled=True
            )
            start_button = False
        else:
            # Show active "Start Optimization" button
            start_button = st.button(
                "🚀 Start Optimization",
                use_container_width=True,
                type="primary",
                key="start_optimization"
            )

        # Reset button (only show if optimization is running or completed)
        reset_button = False
        if is_running or is_completed:
            reset_button = st.button(
                "🔄 Reset & Start New",
                use_container_width=True,
                type="secondary",
                key="reset_optimization"
            )

        return {
            'selected_model': selected_model,
            'target_hardware': selected_hardware,
            'hardware_specs': hw_specs,
            'hardware_category': hw_category,
            'num_units': num_units,
            'num_models': num_models,
            'max_hours': max_hours,
            'num_gpus': num_gpus,  # For backward compatibility (same as num_units)
            'actual_models': actual_models,
            'actual_time': max_hours,
            'estimated_hours': max_hours,  # For backward compatibility with charts
            'estimated_time': actual_time_str,
            'models_per_hour': int(models_per_hour),
            'limiting_factor': 'time',
            'start_button': start_button,
            'reset_button': reset_button,
            'objectives': {
                'accuracy': accuracy_obj,
                'size': size_obj,
                'cost': cost_obj,
                'throughput': throughput_obj,
                'latency': latency_obj,
                'memory': memory_obj,
                'energy': energy_obj,
                'batch_size': batch_obj
            },
            'techniques': {
                'quantization': quantization_enabled,
                'pruning': pruning_enabled,
                'distillation': distillation_enabled,
                'lora': lora_enabled,
                'flash_attention': flash_attention_enabled,
                'kernel_fusion': kernel_fusion_enabled
            }
        }