import numpy as np
import pandas as pd

def estimate_inference_cost(model_size_gb, quantization="FP32"):
    """
    Estimate inference cost per 1M requests based on model size and quantization.

    Args:
        model_size_gb: Model size in GB
        quantization: Quantization type (FP32, FP16, INT8, INT4, Mixed)

    Returns:
        dict with cost metrics
    """
    # Step 1: Determine GPU tier based on model size
    if model_size_gb <= 4:
        gpu_name = "T4"
        gpu_cost_per_hour = 0.35
        throughput_base = 200
    elif model_size_gb <= 10:
        gpu_name = "V100"
        gpu_cost_per_hour = 2.50
        throughput_base = 120
    elif model_size_gb <= 20:
        gpu_name = "A100"
        gpu_cost_per_hour = 4.10
        throughput_base = 80
    else:
        num_gpus = int(np.ceil(model_size_gb / 20))
        gpu_name = f"A100x{num_gpus}"
        gpu_cost_per_hour = 4.10 * num_gpus
        throughput_base = 80

    # Step 2: Adjust throughput for quantization (quantized models are faster)
    quantization_speedup = {
        "FP32": 1.0,
        "FP16": 1.8,
        "INT8": 2.5,
        "INT4": 3.5,
        "Mixed": 2.0
    }
    throughput_qps = throughput_base * quantization_speedup.get(quantization, 1.0)

    # Step 3: Calculate cost per 1M requests
    cost_per_second = gpu_cost_per_hour / 3600
    cost_per_inference = cost_per_second / throughput_qps
    cost_per_1M = cost_per_inference * 1_000_000

    return {
        "cost_per_1M": cost_per_1M,
        "gpu_tier": gpu_name,
        "throughput_qps": throughput_qps,
        "gpu_cost_per_hour": gpu_cost_per_hour
    }

def estimate_latency(model_size_gb, quantization="FP32"):
    """
    Estimate inference latency in milliseconds.
    Latency inversely proportional to throughput.
    """
    cost_info = estimate_inference_cost(model_size_gb, quantization)
    # Latency approximation: 1000ms / QPS (for single request)
    latency_ms = 1000.0 / cost_info["throughput_qps"]
    return latency_ms

def estimate_memory_footprint(model_size_gb, quantization="FP32"):
    """
    Estimate peak memory footprint during inference.
    Memory = model size + activation overhead
    """
    # Activation overhead typically 20-40% of model size
    overhead_factor = 1.3
    memory_gb = model_size_gb * overhead_factor
    return memory_gb

def estimate_energy_consumption(model_size_gb, quantization="FP32"):
    """
    Estimate energy consumption in Watts.
    Based on GPU tier and utilization.
    """
    cost_info = estimate_inference_cost(model_size_gb, quantization)

    # GPU power consumption (TDP)
    gpu_power = {
        "T4": 70,      # 70W
        "V100": 250,   # 250W
        "A100": 400    # 400W
    }

    # Extract base GPU type
    gpu_type = cost_info["gpu_tier"]
    if "x" in gpu_type:
        # Multiple GPUs
        base_type = gpu_type.split("x")[0]
        num_gpus = int(gpu_type.split("x")[1])
        power = gpu_power.get(base_type, 400) * num_gpus
    else:
        power = gpu_power.get(gpu_type, 400)

    # Assume 70% utilization during inference
    power_watts = power * 0.7
    return power_watts

def generate_pareto_front_data():
    """Generate mock Pareto front data with all objectives."""
    np.random.seed(42)

    # Generate many model configurations (100+ points)
    num_models = 150

    # Create models with various sizes and accuracies
    model_sizes = []
    accuracies = []

    # Generate diverse model configurations
    for i in range(num_models):
        size = np.random.uniform(2, 35)

        # Base accuracy trend: larger models tend to be more accurate
        # But with significant variance to create interesting Pareto front
        base_acc = 70 + (size / 35) * 20  # 70-90% range
        noise = np.random.normal(0, 3)  # More variance
        acc = np.clip(base_acc + noise, 68, 92)

        model_sizes.append(size)
        accuracies.append(acc)

    # Calculate other objectives for each model
    # Use mixed quantization types for diversity
    quantization_types = np.random.choice(['FP32', 'FP16', 'INT8', 'INT4'], num_models)

    costs = []
    throughputs = []
    latencies = []
    memories = []
    energies = []

    for i in range(num_models):
        quant = quantization_types[i]
        size = model_sizes[i]

        cost_info = estimate_inference_cost(size, quant)
        costs.append(cost_info['cost_per_1M'])
        throughputs.append(cost_info['throughput_qps'])

        latencies.append(estimate_latency(size, quant))
        memories.append(estimate_memory_footprint(size, quant))
        energies.append(estimate_energy_consumption(size, quant))

    df = pd.DataFrame({
        'accuracy': accuracies,
        'size': model_sizes,
        'cost': costs,
        'throughput': throughputs,
        'latency': latencies,
        'memory': memories,
        'energy': energies
    })

    # Calculate Pareto front for accuracy vs size (default)
    # A point is on the Pareto front if no other point dominates it
    # (dominates = higher accuracy AND smaller size)
    df['is_pareto_accuracy_size'] = False

    for i in range(len(df)):
        is_dominated = False
        for j in range(len(df)):
            if i != j:
                # Check if point j dominates point i
                if (df.iloc[j]['accuracy'] >= df.iloc[i]['accuracy'] and
                    df.iloc[j]['size'] <= df.iloc[i]['size'] and
                    (df.iloc[j]['accuracy'] > df.iloc[i]['accuracy'] or
                     df.iloc[j]['size'] < df.iloc[i]['size'])):
                    is_dominated = True
                    break
        if not is_dominated:
            df.at[i, 'is_pareto_accuracy_size'] = True

    return df

def calculate_pareto_front(df, obj1, obj2, obj1_maximize=True, obj2_maximize=False):
    """
    Calculate Pareto front for any two objectives.

    Args:
        df: DataFrame with objective columns
        obj1: First objective column name
        obj2: Second objective column name
        obj1_maximize: True if obj1 should be maximized, False if minimized
        obj2_maximize: True if obj2 should be maximized, False if minimized

    Returns:
        DataFrame with is_pareto column added
    """
    df = df.copy()
    df['is_pareto'] = False

    for i in range(len(df)):
        is_dominated = False
        for j in range(len(df)):
            if i != j:
                # Check if point j dominates point i
                obj1_better = (df.iloc[j][obj1] >= df.iloc[i][obj1]) if obj1_maximize else (df.iloc[j][obj1] <= df.iloc[i][obj1])
                obj2_better = (df.iloc[j][obj2] >= df.iloc[i][obj2]) if obj2_maximize else (df.iloc[j][obj2] <= df.iloc[i][obj2])
                obj1_strictly_better = (df.iloc[j][obj1] > df.iloc[i][obj1]) if obj1_maximize else (df.iloc[j][obj1] < df.iloc[i][obj1])
                obj2_strictly_better = (df.iloc[j][obj2] > df.iloc[i][obj2]) if obj2_maximize else (df.iloc[j][obj2] < df.iloc[i][obj2])

                if obj1_better and obj2_better and (obj1_strictly_better or obj2_strictly_better):
                    is_dominated = True
                    break
        if not is_dominated:
            df.at[i, 'is_pareto'] = True

    return df

def generate_optimization_progress(budget_hours=2):
    """Generate mock optimization progress data over time."""
    np.random.seed(42)

    # Generate time points
    num_points = 20
    time_points = np.linspace(0, budget_hours, num_points)

    # Accuracy starts lower for LLMs and gradually improves
    base_accuracy = 75
    accuracy_trend = base_accuracy + np.log1p(time_points) * 4
    accuracy = [min(acc + np.random.uniform(-0.3, 0.3), 89) for acc in accuracy_trend]

    # Model size starts large and decreases (GB for LLMs)
    initial_size = 35
    size_reduction = initial_size * (1 - 0.70 * (time_points / budget_hours))
    model_size = [max(s + np.random.uniform(-0.5, 0.5), 2) for s in size_reduction]

    return pd.DataFrame({
        'search_time': time_points,
        'accuracy': accuracy,
        'model_size_gb': model_size
    })

def generate_all_objectives_progress(budget_hours=2):
    """Generate progress data for all optimization objectives over time."""
    np.random.seed(42)

    num_points = 20
    time_points = np.linspace(0, budget_hours, num_points)

    # Accuracy: starts at 72%, improves to ~85%
    base_accuracy = 72
    accuracy_trend = base_accuracy + np.log1p(time_points) * 5.5
    accuracy = [min(acc + np.random.uniform(-0.3, 0.3), 85.2) for acc in accuracy_trend]

    # Model size: starts at 32.5 GB, decreases to ~6.2 GB
    initial_size = 32.5
    size_reduction = initial_size * (1 - 0.81 * (time_points / budget_hours))
    size = [max(s + np.random.uniform(-0.5, 0.5), 6.2) for s in size_reduction]

    # Calculate other metrics based on size and quantization progression
    # Quantization improves gradually from FP32 -> FP16 -> INT8 -> INT4 over time
    # Use a mix to create smoother transitions
    quantizations = (
        ['FP32', 'FP32', 'Mixed', 'Mixed', 'FP16'] +
        ['FP16', 'FP16', 'Mixed', 'INT8', 'INT8'] +
        ['INT8', 'INT8', 'Mixed', 'INT4', 'INT4'] +
        ['INT4', 'INT4', 'INT4', 'INT4', 'INT4']
    )

    cost = []
    throughput = []
    latency = []
    memory = []
    energy = []

    for i in range(num_points):
        quant = quantizations[i]
        model_size = size[i]

        cost_info = estimate_inference_cost(model_size, quant)
        cost.append(cost_info['cost_per_1M'])
        throughput.append(cost_info['throughput_qps'])

        # Add slight variation to latency for smoother visualization
        base_latency = estimate_latency(model_size, quant)
        latency_variation = np.random.uniform(-0.1, 0.1)
        latency.append(max(0.1, base_latency * (1 + latency_variation)))

        memory.append(estimate_memory_footprint(model_size, quant))
        energy.append(estimate_energy_consumption(model_size, quant))

    return pd.DataFrame({
        'time': time_points,
        'accuracy': accuracy,
        'size': size,
        'cost': cost,
        'throughput': throughput,
        'latency': latency,
        'memory': memory,
        'energy': energy
    })

def generate_discovered_models():
    """Generate mock discovered models with their specifications (6 Pareto-optimal configs)."""
    models = [
        {
            'name': 'Optimized-7B-Q8',
            'params': '7.2B',
            'accuracy': 85.2,
            'size_gb': 10.2,
            'quantization': 'INT8'
        },
        {
            'name': 'Optimized-7B-Q4',
            'params': '7.2B',
            'accuracy': 83.8,
            'size_gb': 6.2,
            'quantization': 'INT4'
        },
        {
            'name': 'Optimized-3B-Q8',
            'params': '3.5B',
            'accuracy': 79.2,
            'size_gb': 4.8,
            'quantization': 'INT8'
        },
        {
            'name': 'Optimized-7B-FP16',
            'params': '7.2B',
            'accuracy': 86.5,
            'size_gb': 14.4,
            'quantization': 'FP16'
        },
        {
            'name': 'Optimized-3B-Q4',
            'params': '3.5B',
            'accuracy': 77.8,
            'size_gb': 2.4,
            'quantization': 'INT4'
        },
        {
            'name': 'Optimized-7B-Mixed',
            'params': '7.2B',
            'accuracy': 84.6,
            'size_gb': 8.5,
            'quantization': 'Mixed'
        }
    ]

    return models

def get_base_models():
    """Get list of available base LLM models (recent, single-GPU optimizable)."""
    return [
        "Llama 3.2 3B Instruct",
        "Llama 3.1 8B Instruct",
        "Mistral 7B v0.3",
        "Phi-3.5 Mini (3.8B)",
        "Qwen2.5 3B Instruct",
        "Qwen2.5 7B Instruct",
        "Gemma 2 2B",
        "Gemma 2 9B",
        "Yi 1.5 9B",
        "StableLM 2 1.6B"
    ]

def get_target_hardware():
    """Get list of target hardware platforms for LLM deployment."""
    return {
        "NVIDIA Datacenter": [
            "H100 (80GB)",
            "A100 (80GB)",
            "A100 (40GB)",
            "L40S (48GB)",
            "A10 (24GB)"
        ],
        "NVIDIA Workstation/Consumer": [
            "RTX 4090 (24GB)",
            "RTX 4080 (16GB)",
            "RTX 3090 (24GB)",
            "RTX 3080 Ti (12GB)"
        ],
        "AMD Datacenter": [
            "MI300X (192GB)",
            "MI250X (128GB)",
            "MI210 (64GB)"
        ],
        "AMD Consumer": [
            "RX 7900 XTX (24GB)",
            "RX 7900 XT (20GB)"
        ],
        "Edge Hardware": [
            "Jetson AGX Orin (64GB)",
            "Jetson Orin NX (16GB)",
            "Jetson Orin Nano (8GB)",
            "Apple M2 Ultra (192GB)"
        ]
    }

def get_hardware_specs(hardware_name):
    """Get specifications for a given hardware platform."""
    # Parse hardware name to extract key info
    hardware_specs = {
        # NVIDIA Datacenter
        "H100 (80GB)": {"vram": 80, "compute": "Hopper", "tflops": 1979, "category": "datacenter"},
        "A100 (80GB)": {"vram": 80, "compute": "Ampere", "tflops": 312, "category": "datacenter"},
        "A100 (40GB)": {"vram": 40, "compute": "Ampere", "tflops": 312, "category": "datacenter"},
        "L40S (48GB)": {"vram": 48, "compute": "Ada Lovelace", "tflops": 362, "category": "datacenter"},
        "A10 (24GB)": {"vram": 24, "compute": "Ampere", "tflops": 125, "category": "datacenter"},

        # NVIDIA Workstation/Consumer
        "RTX 4090 (24GB)": {"vram": 24, "compute": "Ada Lovelace", "tflops": 83, "category": "consumer"},
        "RTX 4080 (16GB)": {"vram": 16, "compute": "Ada Lovelace", "tflops": 49, "category": "consumer"},
        "RTX 3090 (24GB)": {"vram": 24, "compute": "Ampere", "tflops": 36, "category": "consumer"},
        "RTX 3080 Ti (12GB)": {"vram": 12, "compute": "Ampere", "tflops": 34, "category": "consumer"},

        # AMD
        "MI300X (192GB)": {"vram": 192, "compute": "CDNA 3", "tflops": 1307, "category": "datacenter"},
        "MI250X (128GB)": {"vram": 128, "compute": "CDNA 2", "tflops": 383, "category": "datacenter"},
        "MI210 (64GB)": {"vram": 64, "compute": "CDNA 2", "tflops": 181, "category": "datacenter"},
        "RX 7900 XTX (24GB)": {"vram": 24, "compute": "RDNA 3", "tflops": 61, "category": "consumer"},
        "RX 7900 XT (20GB)": {"vram": 20, "compute": "RDNA 3", "tflops": 51, "category": "consumer"},

        # Edge
        "Jetson AGX Orin (64GB)": {"vram": 64, "compute": "Ampere", "tops": 275, "category": "edge"},
        "Jetson Orin NX (16GB)": {"vram": 16, "compute": "Ampere", "tops": 100, "category": "edge"},
        "Jetson Orin Nano (8GB)": {"vram": 8, "compute": "Ampere", "tops": 40, "category": "edge"},
        "Apple M2 Ultra (192GB)": {"vram": 192, "compute": "ARM", "tflops": 28, "category": "edge"}
    }

    return hardware_specs.get(hardware_name, {"vram": 24, "compute": "Unknown", "category": "unknown"})