# ---- Data from model (e.g. ExecuTorch/ PyTorch):
baseline_time_ms = 1   # Avg. Time for an Inference (warm-start forward pass).
pte_size_bytes = 9e5
weight_dtype = 'int8'  # Parameter dtype

from lib_edge_eval import fetch_edge_hardware_dataframe, Plotting

# ======= Load the Hardware Table: ========
df_final = fetch_edge_hardware_dataframe()
display(df_final)

# ======= Specify the Baseline from the Hardware Table: ========
baseline_index = 8 # 8 = NVIDIA T4 GPU

Plotting.plot_hardware_comparisons(df_final)
Plotting.plot_hardware_comparisons(df_final, 
                                   metrics=['Peak TOPS (INT8/NPU)', 'Peak TFLOPS (FP32/GPU)'],
                                   titles=['Peak TOPS (INT8/NPU)', 'Peak TFLOPS (FP32/GPU)']
                                   )

from lib_edge_eval import calculate_edge_metrics, Plotting

# ======= Calculate from Baseline: ========
df_comparison = calculate_edge_metrics(
    df=df_final,
    baseline_idx=baseline_index,
    baseline_ms=baseline_time_ms,
    baseline_mem=pte_size_bytes,
    baseline_dtype=weight_dtype if weight_dtype ['int8','fp8'] else 'fp8'
)

display(df_comparison.round(2))
Plotting.plot_comparison_metrics(df_comparison)

import torchvision.models as models
import torch
# model = models.alexnet(); input_shape=(batch_size, 3, 224, 224) # For example.
model = ... # PyTorch model

from deepspeed.profiling.flops_profiler import get_model_profile
from deepspeed.accelerator import get_accelerator

with get_accelerator().device(0):
    batch_size = 1
    flops, macs, params = get_model_profile(model=model, # model
                                    input_shape=(batch_size, 10, 5), # input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument.
                                    args=None, # list of positional arguments to the model.
                                    kwargs=None, # dictionary of keyword arguments to the model.
                                    print_profile=False, # prints the model graph with the measured profile attached to each module
                                    detailed=True, # print the detailed profile
                                    module_depth=-1, # depth into the nested modules, with -1 being the inner most modules
                                    top_modules=1, # the number of top modules to print aggregated profile
                                    warm_up=10, # the number of warm-ups before measuring the time of each module
                                    as_string=True, # print raw numbers (e.g. 1000) or as human-readable strings (e.g. 1k)
                                    output_file=None, # path to the output file. If None, the profiler prints to stdout.
                                    ignore_modules=None) # the list of modules to ignore in the profiling
flops, macs, params

import os
import torch

def export_pytorch_and_get_size(model, filename="model_pytorch.pth"):
    """Exports the PyTorch model to a .pth file and returns its size in bytes."""
    # Save the model state dict
    torch.save(model.state_dict(), filename)

    # Calculate size
    size_bytes = os.path.getsize(filename)
    return size_bytes

# Run the function
pth_filename = "model_pytorch.pth"
pth_size_bytes = export_pytorch_and_get_size(model, pth_filename)

print(f"PyTorch Model (.pth) size: {pth_size_bytes} bytes ({pth_size_bytes/1024:.2f} KB)")

# Get weight dtype
weight_dtype = str(next(model.parameters()).dtype).split('.')[-1]   

from executorch.exir import to_edge_transform_and_lower
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from torch.export import export
import torch
import os

def export_to_executorch(model, example_inputs):
    """Exports a PyTorch model to an ExecuTorch program using XNNPACK."""
    model = model.eval()    
    # Get weight dtype
    weight_dtype = str(next(model.parameters()).dtype).split('.')[-1]   
    # Export to Program
    exported_program = export(model, example_inputs)
    # Lower to Edge and convert to ExecuTorch
    executorch_program = to_edge_transform_and_lower(
        exported_program,
        partitioner=[XnnpackPartitioner()],
    ).to_executorch()
    
    return executorch_program, weight_dtype

# Define inputs
# The SimpleTCN model expects (batch_size, sequence_length, num_features)
example_inputs = (torch.randn(1, 10, 5),)

# Call the function
executorch_program, weight_dtype = export_to_executorch(model, example_inputs)

# The .pte file size represents the binary size of the program including weights
pte_size_bytes = os.path.getsize(pte_path)

print(f"ExecuTorch Program (.pte) size: {pte_size_bytes} bytes ({pte_size_bytes/1024:.2f} KB)")
print(f"PyTorch Model (.pth) size: {pth_size_bytes} bytes ({pth_size_bytes/1024:.2f} KB)")
reduction = (1 - (pte_size_bytes / pth_size_bytes)) * 100
print(f"\nSize Difference: {reduction:.2f}% change in footprint")

print(f"\nInference weights size (approx): {pte_size_bytes} bytes")
print(f"Weight datatype: {weight_dtype}")

import time, torch

def benchmark_executorch_ms(method, input_data, trials=20):
    # Warm-up
    method.execute([input_data])
    start = time.perf_counter()
    for _ in range(trials): method.execute([input_data])
    avg_ms = ((time.perf_counter() - start) / trials) * 1000
    return avg_ms

def benchmark_pytorch_ms(model, input_data, trials=20):
    model.eval()
    # Warm-up
    with torch.no_grad():
        model(input_data)
    with torch.no_grad():
        start = time.perf_counter()
        for _ in range(trials): model(input_data)
    avg_ms = ((time.perf_counter() - start) / trials) * 1000
    return avg_ms

# Run benchmarks
t = torch.randn(1, 10, 5)
num_trials = 150
time_ms_pytorch = benchmark_pytorch_ms(model, t, num_trials)
time_ms_inf = benchmark_executorch_ms(method, t, num_trials)

print(f"ExecuTorch Avg Inference Time ({num_trials} trials): {time_ms_inf:.3f} ms")
print(f"PyTorch Avg Inference Time ({num_trials} trials): {time_ms_pytorch:.3f} ms")
print(f"Speedup: {time_ms_pytorch / time_ms_inf:.2f}x")

import pandas as pd
import sys

def generate_model_performance_report(model, executorch_program,
                                      pth_size_bytes, pte_size_bytes,
                                      time_ms_pytorch, time_ms_inf,
                                      weight_dtype, params, flops):
    """Generates and displays a comprehensive performance comparison report."""

    # 1. Calculate dynamic differences
    size_pct = ((pte_size_bytes - pth_size_bytes) / pth_size_bytes) * 100
    latency_pct = ((time_ms_inf - time_ms_pytorch) / time_ms_pytorch) * 100

    performance_data = {
        "Metric": ["Precision", "Parameters", "FLOPs", "Model Size (KB)", "Inference Time (ms)"],
        "PyTorch (Source)": [
            weight_dtype,
            params,
            flops,
            f"{pth_size_bytes / 1024:.2f}",
            f"{time_ms_pytorch:.3f}"
        ],
        "ExecuTorch (.pte)": [
            weight_dtype,
            params,
            flops,
            f"{pte_size_bytes / 1024:.2f}",
            f"{time_ms_inf:.3f}"
        ],
        "Diff": [
            "--",
            "--",
            "--",
            f"{size_pct:+.2f}%",
            f"{latency_pct:+.2f}%"
        ]
    }

    # Display Report
    report_df = pd.DataFrame(performance_data)
    print("=== Model Performance Report ===")
    display(report_df)

    # 2. Print File Size Comparison (Memory Footprint)
    print(f"\n--- In-Memory (RAM) Footprint Change ---")
    print(f"PyTorch Model: {pth_size_bytes / 1024:.2f} KB")
    print(f"ExecuTorch Program: {pte_size_bytes / 1024:.2f} KB")
    mem_diff = ((pte_size_bytes - pth_size_bytes) / pth_size_bytes) * 100
    print(f"Difference: {mem_diff:+.2f}%")

    # 3. Throughput Analysis
    numeric_flops = float(flops.split()[0]) * 1e6
    et_gflops = (numeric_flops / (time_ms_inf / 1000)) / 1e9
    pt_gflops = (numeric_flops / (time_ms_pytorch / 1000)) / 1e9

    print(f"\n--- Throughput Change (on baseline hardware) ---")
    print(f"PyTorch: {pt_gflops:.4f} GFLOPS")
    print(f"ExecuTorch: {et_gflops:.4f} GFLOPS")
    print(f"Speedup Factor: {time_ms_pytorch / time_ms_inf:.2f}x faster")

# Execute the function
generate_model_performance_report(
    model,
    executorch_program,
    pth_size_bytes,
    pte_size_bytes,
    time_ms_pytorch,
    time_ms_inf,
    weight_dtype,
    params,
    flops
)