import json
import numpy as np
import matplotlib.pyplot as plt

#font sizes
plt.rcParams.update({
    "font.size": 18,
    "axes.titlesize": 22,
    "axes.labelsize": 20,
    "xtick.labelsize": 16,
    "ytick.labelsize": 16,
    "legend.fontsize": 16,
})

def getResults(filename):

    with open(filename) as f:
        data = json.load(f)
    
    # Prepare lists
    spmv_sizes = []
    spmv_times = []
    
    dense_sizes = []
    dense_times = []
    
    labels = []
    
    for d in data:
        spmv_sizes.append(d["nnz"])
        spmv_times.append(d["spmv_avg_ns"])
    
        dense_sizes.append(d["dense_rows"])
        dense_times.append(d["dense_avg_ns"])
    
        labels.append(d["label"])

    return spmv_sizes, spmv_times, dense_sizes, dense_times, labels


c_spmv_sizes, c_spmv_times, c_dense_sizes, c_dense_times, c_labels = getResults("cSparseResults.json")
go_spmv_sizes, go_spmv_times, go_dense_sizes, go_dense_times, go_labels = getResults("goSparseResults.json")

# make sure data aligns on x axis
for i in range(len(go_spmv_sizes)):
    assert(c_spmv_sizes[i] == go_spmv_sizes[i])
    assert(go_dense_sizes[i] == c_dense_sizes[i])



#spmv plot
x_spmv = c_spmv_sizes
x_spmv = np.float64(x_spmv)/1e6
go_spmv_float = np.float64(go_spmv_times)
c_spmv_float = np.float64(c_spmv_times)
ratio_spmv = go_spmv_float/c_spmv_float
plt.figure("spmv")
plt.plot(x_spmv, ratio_spmv, 'kd-')
plt.title("Sparse Matrix-Vector multiplication ratio Go native / C MKL")
plt.xlabel("SuiteSparse matrix millions of number of non-zeros")
plt.ylabel("Average time ratio Go/C for SpMV-kernel")
xtick_string = [f"{x:.1f}\n{label}" for x, label in zip(x_spmv, c_labels)]
plt.xticks(x_spmv, xtick_string)
plt.grid()

# dense plot
x_dense = c_dense_sizes
go_dense_float = np.float64(go_dense_times)
c_dense_float = np.float64(c_dense_times)
ratio_dense = go_dense_float/c_dense_float

plt.figure("dense")
plt.plot(x_dense, ratio_dense, 'kd-')
plt.title("Dense Matrix-Matrix multiplication ratio Go native / C MKL")
plt.xlabel("Dense matrix size")
plt.ylabel("Average time ratio Go/C for matmul kernel")
xtick_string = [f"{x}^2" for x in x_dense]
plt.xticks(x_dense, xtick_string)
plt.grid()

plt.show()

## ---- SpMV Plot ----
#plt.figure()
#plt.plot(spmv_sizes, spmv_times, marker='o')
#
#for i, label in enumerate(labels):
#    plt.annotate(label, (spmv_sizes[i], spmv_times[i]))
#
#plt.xlabel("Number of Nonzeros (nnz)")
#plt.ylabel("Average Time (ns)")
#plt.title("SpMV: Avg Time vs Size")
#plt.xscale("log")
#plt.yscale("log")
#plt.grid(True)
#
## ---- Dense Plot ----
#plt.figure()
#plt.plot(dense_sizes, dense_times, marker='o')
#
#for i, label in enumerate(labels):
#    plt.annotate(label, (dense_sizes[i], dense_times[i]))
#
#plt.xlabel("Matrix Size (rows × cols)")
#plt.ylabel("Average Time (ns)")
#plt.title("Dense: Avg Time vs Size")
#plt.xscale("log")
#plt.yscale("log")
#plt.grid(True)
#
#plt.show()