import matplotlib.pyplot as plt
import numpy as np

performance_data = {
    "Mem": {
        "regular": {"CPU_cycles": 33940, "L1_DTLB_access": 13202, "L1_DTLB_miss": 36, "L2_DTLB_miss": 10},
        "modified": {"CPU_cycles": 2457, "L1_DTLB_access": 151, "L1_DTLB_miss": 0, "L2_DTLB_miss": 0},
    },
    "Test_C": {
        "regular": {"CPU_cycles": 34040, "L1_DTLB_access": 13284, "L1_DTLB_miss": 36, "L2_DTLB_miss": 10},
        "modified": {"CPU_cycles": 2789, "L1_DTLB_access": 396, "L1_DTLB_miss": 0, "L2_DTLB_miss": 0},
    },
    "Glibc": {
        "regular": {"CPU_cycles": 36053, "L1_DTLB_access": 14027, "L1_DTLB_miss": 36, "L2_DTLB_miss": 10},
        "modified": {"CPU_cycles": 4219, "L1_DTLB_access": 783, "L1_DTLB_miss": 0, "L2_DTLB_miss": 0},
    },
    "Richards": {
        "regular": {"CPU_cycles": 37808, "L1_DTLB_access": 14363, "L1_DTLB_miss": 10, "L2_DTLB_miss": 6},
        "modified": {"CPU_cycles": 4433, "L1_DTLB_access": 874, "L1_DTLB_miss": 0, "L2_DTLB_miss": 0},
    },
    "Matrix_mul": {
        "regular": {"CPU_cycles": 42383, "L1_DTLB_access": 17020, "L1_DTLB_miss": 14, "L2_DTLB_miss": 10},
        "modified": {"CPU_cycles": 12149, "L1_DTLB_access": 4411, "L1_DTLB_miss": 0, "L2_DTLB_miss": 0},
    }
}

def compute_percent_diff(regular, modified):
    out = {}
    for k in regular:
        if k in modified and regular[k] != 0:
            out[k] = ((regular[k] - modified[k]) / regular[k]) * 100
        elif k in modified and regular[k] == 0 and modified[k] == 0:
            out[k] = 0.0
        elif k in modified and regular[k] == 0 and modified[k] != 0:
            out[k] = -100.0 
    return out

for t, d in performance_data.items():
    d["percent_diff"] = compute_percent_diff(d["regular"], d["modified"])

tests = list(performance_data.keys())
metrics = ["CPU_cycles", "L1_DTLB_access", "L1_DTLB_miss", "L2_DTLB_miss"]

# 1. Summary plot: Percent Improvement for all metrics grouped by test
x = np.arange(len(tests))
width = 0.2
fig, ax = plt.subplots(figsize=(12, 7))

for i, metric in enumerate(metrics):
    values = [performance_data[t]["percent_diff"].get(metric, 0) for t in tests]
    offset = (i - len(metrics)/2 + 0.5) * width
    ax.bar(x + offset, values, width, label=metric.replace('_', ' ').title())

ax.set_ylabel('Percent Improvement (%)')
ax.set_title('Percent Improvement (Regular -> Modified) by Test and Metric')
ax.set_xticks(x)
ax.set_xticklabels(tests)
ax.legend()
ax.set_ylim(0, 110)
ax.grid(axis='y', linestyle='--', alpha=0.7)
fig.tight_layout()
plt.savefig('percent_improvement_summary.png')

# 2. Individual metric plots: Regular vs Modified grouped by test
y_label_map = {
    "CPU_cycles": "Cycles",
    "L1_DTLB_access": "Access Count",
    "L1_DTLB_miss": "Miss Count",
    "L2_DTLB_miss": "Miss Count"
}

for metric in metrics:
    fig, ax = plt.subplots(figsize=(10, 6))
    reg_vals = [performance_data[t]["regular"][metric] for t in tests]
    mod_vals = [performance_data[t]["modified"][metric] for t in tests]
    
    x_local = np.arange(len(tests))
    w = 0.35
    
    ax.bar(x_local - w/2, reg_vals, w, label='Regular', color='steelblue')
    ax.bar(x_local + w/2, mod_vals, w, label='Modified', color='salmon')
    
    # Setting dynamic y-label
    ax.set_ylabel(y_label_map.get(metric, metric.replace('_', ' ').title()))
    ax.set_title(f'Comparison: {metric.replace("_", " ").title()}')
    ax.set_xticks(x_local)
    ax.set_xticklabels(tests)
    ax.legend()
    ax.grid(axis='y', linestyle='--', alpha=0.7)
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig(f'grouped_comparison_{metric}.png')