import pandas as pd
import numpy as np
import json
from pathlib import Path


# GPU specifications data

# Notes:
# FP16_TFLOPS is either FP16 perf or BF16 if available
# High BW Interconnect refers to the ability to use NVLink or SXM interconnects to connect multiple GPUs together, this is either 0 or 1
# Specs moved to external JSON file (gpu_data.json)

def load_gpu_data(json_path: str | None = None):
    base_dir = Path(__file__).parent
    path = Path(json_path) if json_path else (base_dir / "gpu_data.json")
    with path.open("r") as f:
        return json.load(f)


def build_df(gpu_dict):
    """Convert nested GPU dictionary to DataFrame"""
    data = []
    for name, specs in gpu_dict.items():
        row = {"name": name}
        row.update(specs)
        data.append(row)
    
    df = pd.DataFrame(data)
    return df


def gpu_score(
    df,
    memory_weight=0.7,
    compute_weight=0.3,
    bandwidth_bonus_weight=0.3,
    interconnect_weight=0.3,
    min_floor=0.05,
):
    """
    GPU score calculation using:
      - Memory capacity (0-1), moderately boosted by memory bandwidth
      - FP16 TFLOPs (0-1) as a separate, tunable weight
      - Optional high-bandwidth interconnect bonus (default off)
    
    Args:
        df: DataFrame with MEMORY_GB, MEMORY_BW_GBPS, FP16_TFLOPS, HIGH_BW_INTERCONNECT_EXISTS
        memory_weight: Weight for memory component (0-1)
        compute_weight: Weight for FP16 compute component (0-1)
        bandwidth_bonus_weight: Scales bandwidth effect on memory (e.g., 0.3 => up to +30% boost)
        interconnect_weight: Optional bonus for high-BW interconnect (0 disables)
        min_floor: Minimum normalized value (>0 ensures no exact zeros); result scaled to [min_floor, 1].
    
    Notes:
        - To make bandwidth influence more/less, adjust bandwidth_bonus_weight.
        - To favor compute vs memory, adjust compute_weight vs memory_weight.
        - Final combined score is normalized to 0-1 across GPUs.
    """
    # Normalize memory capacity
    mem = df["MEMORY_GB"].astype(float)
    mem_min, mem_max = mem.min(), mem.max()
    mem_score = pd.Series(1.0, index=df.index) if mem_max == mem_min else (mem - mem_min) / (mem_max - mem_min)
    
    # Normalize memory bandwidth
    bw = df["MEMORY_BW_GBPS"].astype(float)
    bw_min, bw_max = bw.min(), bw.max()
    bw_score = pd.Series(1.0, index=df.index) if bw_max == bw_min else (bw - bw_min) / (bw_max - bw_min)
    
    # Apply a moderate multiplicative bonus to memory based on bandwidth
    # Example: with bandwidth_bonus_weight=0.3, highest-bandwidth memory gets up to +30% boost
    bandwidth_weighted_memory = mem_score * (1.0 + bandwidth_bonus_weight * bw_score)
    
    # Normalize FP16 TFLOPs
    fp16 = df["FP16_TFLOPS"].astype(float)
    fp16_min, fp16_max = fp16.min(), fp16.max()
    compute_score = pd.Series(1.0, index=df.index) if fp16_max == fp16_min else (fp16 - fp16_min) / (fp16_max - fp16_min)
    
    # Optional interconnect bonus
    interconnect_bonus = df["HIGH_BW_INTERCONNECT_EXISTS"].astype(float) * interconnect_weight
    
    # Combine components
    combined = (memory_weight * bandwidth_weighted_memory) + (compute_weight * compute_score) + interconnect_bonus
    
    # Normalize to 0-1 for comparability
    cmin, cmax = combined.min(), combined.max()
    if cmax == cmin:
        return pd.Series(1.0, index=df.index)
    combined01 = (combined - cmin) / (cmax - cmin)
    return (combined01 * (1.0 - min_floor)) + min_floor


def main():
    """Run GPU score calculation and display results in a table"""
    # Build dataframe
    gpu_data = load_gpu_data()
    df = build_df(gpu_data)

    # Default weights: equal memory/compute; moderate bandwidth bonus; no interconnect bonus
    df["score"] = gpu_score(
        df,
        memory_weight=0.6,
        compute_weight=0.4,
        bandwidth_bonus_weight=0.4,
        interconnect_weight=0.1,
    )
    
    # Create results table with GPU names and scores
    results = df[["name", "score"]].copy()
    
    # Sort by score (descending) for better readability
    results = results.sort_values("score", ascending=False)
    
    # Format scores to 3 decimal places for cleaner display
    results["score"] = results["score"].round(3)
    
    # Print table
    print("\nGPU Ranking Results (0-1 scale, higher is better)\n")
    print("=" * 80)
    print(f"{'GPU Name':<30} {'Score':<12}")
    print("=" * 80)
    
    for _, row in results.iterrows():
        print(f"{row['name']:<30} {row['score']:<12.3f}")
    
    print("=" * 80)
    
    # Also output JSON (list of {name, score})
    json_payload = results.to_dict(orient="records")
    print("\nJSON Results:\n")
    print(json.dumps(json_payload, indent=2))

if __name__ == "__main__":
    main()