"""
GPU Runtime Predictor - Gradio Space
=====================================
Paste your PyTorch/CUDA code, select GPUs from the catalog,
and get predicted runtimes for each GPU.
"""

import gradio as gr
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import json
import re
import pickle
import os
from huggingface_hub import hf_hub_download

# ============================================================================
# LOAD MODEL ARTIFACTS
# ============================================================================
MODEL_REPO = "RajBhope/gpu-runtime-predictor"

def download_artifacts():
    """Download all model artifacts from Hub."""
    files = ['model_gbr.pkl', 'model_rf.pkl', 'model_nn.pt', 'scaler_X.pkl',
             'scaler_params.json', 'gpu_catalog.json', 'nn_config.json', 'metrics.json']
    paths = {}
    for f in files:
        paths[f] = hf_hub_download(repo_id=MODEL_REPO, filename=f)
    return paths

print("Downloading model artifacts...")
artifact_paths = download_artifacts()

# Load models
with open(artifact_paths['model_gbr.pkl'], 'rb') as f:
    model_gbr = pickle.load(f)

with open(artifact_paths['model_rf.pkl'], 'rb') as f:
    model_rf = pickle.load(f)

with open(artifact_paths['scaler_X.pkl'], 'rb') as f:
    scaler_X = pickle.load(f)

with open(artifact_paths['scaler_params.json'], 'r') as f:
    scaler_params = json.load(f)

with open(artifact_paths['gpu_catalog.json'], 'r') as f:
    GPU_CATALOG = json.load(f)

with open(artifact_paths['nn_config.json'], 'r') as f:
    nn_config = json.load(f)

with open(artifact_paths['metrics.json'], 'r') as f:
    metrics = json.load(f)


# Load NN model
class RuntimeMLP(nn.Module):
    def __init__(self, input_dim, hidden_dims=[512, 256, 128], dropout=0.15):
        super().__init__()
        layers = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.LayerNorm(h_dim),
                nn.GELU(),
                nn.Dropout(dropout),
            ])
            prev_dim = h_dim
        layers.append(nn.Linear(prev_dim, 1))
        self.net = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.net(x).squeeze(-1)


model_nn = RuntimeMLP(**nn_config)
model_nn.load_state_dict(torch.load(artifact_paths['model_nn.pt'], map_location='cpu', weights_only=True))
model_nn.eval()

GPU_FEATURE_COLS = [
    'cuda_cores', 'tensor_cores', 'memory_gb', 'memory_bandwidth_gbps',
    'base_clock_mhz', 'boost_clock_mhz', 'sm_count', 'fp32_tflops',
    'fp16_tflops', 'tdp_watts', 'compute_capability', 'l2_cache_mb',
]

print("Models loaded!")


# ============================================================================
# CODE FEATURE EXTRACTION
# ============================================================================

def extract_code_features(code_text):
    """Extract features from source code text."""
    features = {}
    
    lines = code_text.strip().split('\n')
    features['num_lines'] = len(lines)
    features['num_chars'] = len(code_text)
    features['avg_line_length'] = np.mean([len(l) for l in lines]) if lines else 0
    
    tokens = re.findall(r'[a-zA-Z_]\w*|[0-9]+\.?[0-9]*', code_text)
    features['num_tokens'] = len(tokens)
    
    numbers = re.findall(r'\b(\d+\.?\d*)\b', code_text)
    nums = [float(n) for n in numbers if n]
    features['num_numeric_literals'] = len(nums)
    features['max_numeric'] = max(nums) if nums else 0
    features['min_numeric'] = min(nums) if nums else 0
    features['mean_numeric'] = np.mean(nums) if nums else 0
    features['sum_numeric_log'] = np.log1p(sum(nums)) if nums else 0
    
    large_nums = [n for n in nums if n >= 64]
    features['num_large_dims'] = len(large_nums)
    features['product_large_dims_log'] = np.log1p(np.prod(large_nums[:5])) if large_nums else 0
    
    pytorch_ops = {
        'matmul': r'torch\.matmul|torch\.mm|@',
        'conv': r'Conv[12]d|conv[12]d',
        'attention': r'attention|Attention|MultiheadAttention|softmax.*matmul',
        'linear': r'nn\.Linear|linear',
        'batchnorm': r'BatchNorm|batchnorm',
        'layernorm': r'LayerNorm|layernorm',
        'softmax': r'softmax|Softmax',
        'relu': r'relu|ReLU',
        'gelu': r'gelu|GELU',
        'sigmoid': r'sigmoid|Sigmoid',
        'tanh': r'tanh|Tanh',
        'dropout': r'Dropout|dropout',
        'embedding': r'Embedding|embedding',
        'pooling': r'Pool|pool|MaxPool|AvgPool',
        'fft': r'fft|FFT',
        'sort': r'torch\.sort',
        'backward': r'backward|grad',
        'loss': r'Loss|loss|CrossEntropy',
        'cat': r'torch\.cat|concatenate',
        'reshape': r'reshape|view|contiguous',
        'transpose': r'transpose|\.t\(\)|permute',
        'reduce': r'torch\.sum|torch\.mean|torch\.max|torch\.min|reduce',
    }
    
    for op_name, pattern in pytorch_ops.items():
        features[f'has_{op_name}'] = 1 if re.search(pattern, code_text) else 0
    
    features['uses_float16'] = 1 if re.search(r'float16|half|fp16', code_text) else 0
    features['uses_float32'] = 1 if re.search(r'float32|float(?!16)', code_text) else 0
    features['uses_cuda'] = 1 if re.search(r"'cuda'|\.cuda\(\)|device='cuda'", code_text) else 0
    
    features['num_for_loops'] = len(re.findall(r'\bfor\b', code_text))
    features['num_function_defs'] = len(re.findall(r'\bdef\b', code_text))
    features['num_class_defs'] = len(re.findall(r'\bclass\b', code_text))
    features['num_imports'] = len(re.findall(r'\bimport\b', code_text))
    
    features['num_torch_calls'] = len(re.findall(r'torch\.', code_text))
    features['num_nn_calls'] = len(re.findall(r'nn\.', code_text))
    
    dim_patterns = [r'\((\d+),\s*(\d+)\)', r'\((\d+),\s*(\d+),\s*(\d+)\)', r'\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)']
    all_dims = []
    for pattern in dim_patterns:
        for match in re.finditer(pattern, code_text):
            dims = [int(g) for g in match.groups()]
            all_dims.extend(dims)
    
    features['num_dim_specs'] = len(all_dims)
    features['max_dim'] = max(all_dims) if all_dims else 0
    features['total_elements_log'] = 0
    if all_dims:
        tuples = re.findall(r'\([\d,\s]+\)', code_text)
        for t in tuples:
            dims = [int(d) for d in re.findall(r'\d+', t)]
            if len(dims) >= 2:
                prod = 1
                for d in dims:
                    prod *= d
                features['total_elements_log'] = max(features['total_elements_log'], np.log1p(prod))
    
    features['compute_bound_score'] = features.get('has_matmul', 0) + features.get('has_conv', 0) + features.get('has_linear', 0)
    features['memory_bound_score'] = features.get('has_embedding', 0) + features.get('has_cat', 0) + features.get('has_transpose', 0) + features.get('has_relu', 0)
    features['mixed_score'] = features.get('has_attention', 0) + features.get('has_batchnorm', 0) + features.get('has_layernorm', 0)
    
    return features


def estimate_flops_and_memory(code_text):
    """Heuristic estimate of FLOPs and memory bytes from code."""
    numbers = re.findall(r'\b(\d+)\b', code_text)
    nums = [int(n) for n in numbers if int(n) > 0]
    
    # Detect dtype
    dtype_bytes = 2 if re.search(r'float16|half', code_text) else 4
    
    # Try to identify tensor dimensions for FLOPs estimation
    flops = 0
    memory = 0
    
    # Matrix multiplication: look for matmul patterns
    if re.search(r'matmul|torch\.mm|@', code_text):
        dims = [n for n in nums if n >= 8]
        if len(dims) >= 3:
            M, K, N = dims[0], dims[1], dims[2] if len(dims) > 2 else dims[1]
            flops = 2 * M * N * K
            memory = dtype_bytes * (M*K + K*N + M*N)
    
    # Conv2D
    elif re.search(r'Conv[12]d', code_text):
        dims = [n for n in nums if n >= 1]
        if len(dims) >= 5:
            batch, in_ch, out_ch = dims[0], dims[1], dims[2]
            H = W = dims[3] if len(dims) > 3 else 56
            ks = dims[4] if len(dims) > 4 else 3
            flops = 2 * batch * out_ch * H * W * in_ch * ks * ks
            memory = dtype_bytes * (batch*in_ch*H*W + out_ch*in_ch*ks*ks + batch*out_ch*H*W)
    
    # Attention
    elif re.search(r'attention|Attention', code_text):
        dims = [n for n in nums if n >= 4]
        if len(dims) >= 3:
            batch, seq_len, hidden = dims[0], dims[1], dims[2]
            flops = 4 * batch * seq_len * seq_len * hidden
            memory = dtype_bytes * batch * 3 * seq_len * hidden * 2
    
    # Linear
    elif re.search(r'nn\.Linear', code_text):
        dims = [n for n in nums if n >= 8]
        if len(dims) >= 2:
            in_f, out_f = dims[0], dims[1]
            batch = dims[2] if len(dims) > 2 else 1
            flops = 2 * batch * in_f * out_f
            memory = dtype_bytes * (batch * in_f + in_f * out_f + batch * out_f)
    
    # Generic fallback: estimate from tensor sizes
    if flops == 0:
        large_nums = sorted([n for n in nums if n >= 32], reverse=True)[:4]
        if large_nums:
            total_elements = 1
            for n in large_nums:
                total_elements *= n
            flops = total_elements * 2
            memory = dtype_bytes * total_elements * 2
    
    return flops, memory, dtype_bytes


def predict_runtime(code_text, selected_gpus, model_choice="Ensemble"):
    """Predict runtime for code on selected GPUs."""
    if not code_text.strip():
        return "⚠️ Please paste some code.", None
    
    if not selected_gpus:
        return "⚠️ Please select at least one GPU.", None
    
    # Extract code features
    code_feats = extract_code_features(code_text)
    code_feat_names = sorted(code_feats.keys())
    code_feat_vec = [code_feats[k] for k in code_feat_names]
    
    # Estimate FLOPs and memory
    flops, memory_bytes, dtype_bytes = estimate_flops_and_memory(code_text)
    arithmetic_intensity = flops / max(memory_bytes, 1)
    
    results = []
    
    for gpu_key in selected_gpus:
        gpu_spec = GPU_CATALOG.get(gpu_key)
        if gpu_spec is None:
            continue
        
        # GPU features
        gpu_feat_vec = [gpu_spec[col] for col in GPU_FEATURE_COLS]
        
        # Extra features
        extra_feats = [np.log1p(flops), np.log1p(memory_bytes), arithmetic_intensity, dtype_bytes]
        
        # Combine
        all_feats = np.array(code_feat_vec + gpu_feat_vec + extra_feats, dtype=np.float32).reshape(1, -1)
        
        # Normalize
        all_feats_scaled = scaler_X.transform(all_feats)
        all_feats_scaled = np.nan_to_num(all_feats_scaled, nan=0.0, posinf=0.0, neginf=0.0)
        
        # Predict
        if model_choice == "GBR":
            pred_log = model_gbr.predict(all_feats_scaled)[0]
        elif model_choice == "Random Forest":
            pred_log = model_rf.predict(all_feats_scaled)[0]
        elif model_choice == "Neural Net":
            with torch.no_grad():
                pred_log = model_nn(torch.tensor(all_feats_scaled, dtype=torch.float32)).item()
        else:  # Ensemble
            pred_gbr = model_gbr.predict(all_feats_scaled)[0]
            pred_rf = model_rf.predict(all_feats_scaled)[0]
            with torch.no_grad():
                pred_nn = model_nn(torch.tensor(all_feats_scaled, dtype=torch.float32)).item()
            pred_log = 0.5 * pred_gbr + 0.3 * pred_rf + 0.2 * pred_nn
        
        runtime_ms = np.expm1(pred_log)
        runtime_ms = max(runtime_ms, 0.001)
        
        results.append({
            'GPU': gpu_spec['name'],
            'Runtime (ms)': round(runtime_ms, 4),
            'FP32 TFLOPS': gpu_spec['fp32_tflops'],
            'Mem BW (GB/s)': gpu_spec['memory_bandwidth_gbps'],
            'VRAM (GB)': gpu_spec['memory_gb'],
            'Relative Speed': None,
        })
    
    if not results:
        return "⚠️ No valid GPUs selected.", None
    
    # Sort by runtime
    results.sort(key=lambda x: x['Runtime (ms)'])
    
    # Calculate relative speed (fastest = 1.0x)
    fastest = results[0]['Runtime (ms)']
    for r in results:
        r['Relative Speed'] = f"{r['Runtime (ms)'] / fastest:.2f}x"
    
    # Format output
    df_results = pd.DataFrame(results)
    
    # Summary text
    summary = f"### 🏆 Fastest: **{results[0]['GPU']}** ({results[0]['Runtime (ms)']:.4f} ms)\n"
    summary += f"### 🐢 Slowest: **{results[-1]['GPU']}** ({results[-1]['Runtime (ms)']:.4f} ms)\n"
    summary += f"### ⚡ Speedup: **{results[-1]['Runtime (ms)']/results[0]['Runtime (ms)']:.1f}x** (fastest vs slowest)\n\n"
    
    summary += f"**Estimated FLOPs:** {flops:,.0f}\n\n"
    summary += f"**Estimated Memory:** {memory_bytes:,.0f} bytes\n\n"
    summary += f"**Arithmetic Intensity:** {arithmetic_intensity:.2f} FLOP/byte\n\n"
    
    if arithmetic_intensity > 10:
        summary += "🔥 **Compute-bound** workload — faster GPUs with more TFLOPS will help most"
    else:
        summary += "💾 **Memory-bound** workload — GPUs with higher memory bandwidth will help most"
    
    return summary, df_results


# ============================================================================
# EXAMPLE CODES
# ============================================================================

EXAMPLE_CODES = {
    "Matrix Multiplication (2048x2048)": """import torch

def matmul_kernel(A, B):
    # Matrix multiplication: (2048, 2048) x (2048, 2048) -> (2048, 2048)
    C = torch.matmul(A, B)
    return C

A = torch.randn(2048, 2048, dtype=torch.float32, device='cuda')
B = torch.randn(2048, 2048, dtype=torch.float32, device='cuda')
C = matmul_kernel(A, B)
torch.cuda.synchronize()""",

    "Self-Attention (batch=8, seq=1024)": """import torch
import torch.nn.functional as F

def self_attention(Q, K, V, num_heads=16):
    B, S, D = Q.shape
    head_dim = D // num_heads
    
    Q = Q.view(B, S, num_heads, head_dim).transpose(1, 2)
    K = K.view(B, S, num_heads, head_dim).transpose(1, 2)
    V = V.view(B, S, num_heads, head_dim).transpose(1, 2)
    
    attn = torch.matmul(Q, K.transpose(-2, -1)) / (head_dim ** 0.5)
    attn = F.softmax(attn, dim=-1)
    out = torch.matmul(attn, V)
    return out.transpose(1, 2).contiguous().view(B, S, D)

hidden_dim = 1024
Q = torch.randn(8, 1024, hidden_dim, dtype=torch.float32, device='cuda')
K = torch.randn(8, 1024, hidden_dim, dtype=torch.float32, device='cuda')
V = torch.randn(8, 1024, hidden_dim, dtype=torch.float32, device='cuda')
out = self_attention(Q, K, V)
torch.cuda.synchronize()""",

    "Conv2D ResNet Block": """import torch
import torch.nn as nn

def conv2d_forward(x, conv):
    # Conv2D: batch=16, in_channels=256, out_channels=512
    # Input: (16, 256, 56, 56), Kernel: 3x3
    return conv(x)

conv = nn.Conv2d(256, 512, kernel_size=3, padding=1).to('cuda')
x = torch.randn(16, 256, 56, 56, dtype=torch.float32, device='cuda')
out = conv2d_forward(x, conv)
torch.cuda.synchronize()""",

    "Transformer Block": """import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn = nn.MultiheadAttention(768, 12, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(768, 3072),
            nn.GELU(),
            nn.Linear(3072, 768)
        )
        self.ln1 = nn.LayerNorm(768)
        self.ln2 = nn.LayerNorm(768)
    
    def forward(self, x):
        attn_out, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x))
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x

block = TransformerBlock().to('cuda')
x = torch.randn(8, 512, 768, dtype=torch.float32, device='cuda')
out = block(x)
torch.cuda.synchronize()""",

    "Elementwise GELU (100M elements)": """import torch

def elementwise_op(x):
    # Elementwise gelu on tensor of size 100000000
    return torch.nn.functional.gelu(x)

x = torch.randn(100000000, dtype=torch.float32, device='cuda')
out = elementwise_op(x)
torch.cuda.synchronize()""",

    "LLM Linear Layer (fp16, vocab=50257)": """import torch
import torch.nn as nn

def linear_forward(x, linear):
    # Linear layer: (32, 4096) -> (32, 50257)
    return linear(x)

linear = nn.Linear(4096, 50257).to('cuda')
x = torch.randn(32, 4096, dtype=torch.float16, device='cuda')
out = linear_forward(x, linear)
torch.cuda.synchronize()""",
}


# ============================================================================
# GRADIO UI
# ============================================================================

gpu_choices = list(GPU_CATALOG.keys())
gpu_display_names = {k: v['name'] for k, v in GPU_CATALOG.items()}


def load_example(example_name):
    return EXAMPLE_CODES.get(example_name, "")


with gr.Blocks(
    title="GPU Runtime Predictor",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown("""
    # ⚡ GPU Runtime Predictor
    
    Predict how fast your PyTorch/CUDA code will run on different GPU hardware.
    Paste your code, select GPUs from the catalog, and get instant runtime estimates.
    
    **Model**: Ensemble of GBR + Random Forest + Neural Network | **R² = 0.993** | **12 GPUs** | **15 workload types**
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            example_dropdown = gr.Dropdown(
                choices=list(EXAMPLE_CODES.keys()),
                label="📝 Load Example Code",
                value=None,
                interactive=True,
            )
            
            code_input = gr.Code(
                label="Your PyTorch/CUDA Code",
                language="python",
                lines=20,
                value=EXAMPLE_CODES["Matrix Multiplication (2048x2048)"],
            )
        
        with gr.Column(scale=2):
            gpu_selector = gr.CheckboxGroup(
                choices=[(gpu_display_names[k], k) for k in gpu_choices],
                value=list(GPU_CATALOG.keys()),
                label="🖥️ Select GPUs to Compare",
            )
            
            model_selector = gr.Radio(
                choices=["Ensemble", "GBR", "Random Forest", "Neural Net"],
                value="Ensemble",
                label="🤖 Prediction Model",
            )
            
            predict_btn = gr.Button("⚡ Predict Runtime", variant="primary", size="lg")
    
    gr.Markdown("---")
    
    with gr.Row():
        with gr.Column():
            summary_output = gr.Markdown(label="Summary")
        
    with gr.Row():
        results_table = gr.DataFrame(
            label="📊 Runtime Predictions (sorted fastest → slowest)",
            interactive=False,
        )
    
    gr.Markdown("""
    ---
    ### ℹ️ How It Works
    
    1. **Code Analysis**: Extracts 48 features from your code (tensor dimensions, operation types, complexity indicators)
    2. **GPU Encoding**: Uses 12 hardware specs for each GPU (CUDA cores, memory bandwidth, TFLOPS, etc.)
    3. **ML Prediction**: Ensemble predicts `log(runtime_ms)` → converted back to milliseconds
    
    **Powered by**: [Training Dataset](https://huggingface.co/datasets/RajBhope/gpu-runtime-prediction-dataset) | [Model](https://huggingface.co/RajBhope/gpu-runtime-predictor)
    
    *Runtimes are estimates based on a roofline performance model. Actual runtimes may vary based on driver version, CUDA toolkit, memory state, and other factors.*
    """)
    
    # Event handlers
    example_dropdown.change(
        fn=load_example,
        inputs=[example_dropdown],
        outputs=[code_input],
    )
    
    predict_btn.click(
        fn=predict_runtime,
        inputs=[code_input, gpu_selector, model_selector],
        outputs=[summary_output, results_table],
    )

demo.launch()