import csv
import math
import matplotlib.pyplot as plot
import numpy as np

def geomean(data):
    split_data = []
    for x in data:
        split_data.append(data[x])
    return round(math.prod(split_data) ** (1 / len(split_data)), 2)

def make_comparison_list(list_a, list_b, encoder):
    results_720p = {}
    results_1080p = {}
    results_4k = {}

    for x in range(len(list_a)):
        if list_a[x]["Encoder"] == encoder:
            difference = round(float(list_a[x]["Run Time"]) / float(list_b[x]["Run Time"]), 2)
    
            if list_a[x]["Resolution"] == "720p":
                results_720p.update({f"Preset {list_a[x]["Preset"]}": float(difference)})
            elif list_a[x]["Resolution"] == "1080p":
                results_1080p.update({f"Preset {list_a[x]["Preset"]}": float(difference)})
            elif list_a[x]["Resolution"] == "4k":
                results_4k.update({f"Preset {list_a[x]["Preset"]}": float(difference)})

    if encoder == "libx265":
        results_720p = dict(reversed(results_720p.items()))
        results_1080p = dict(reversed(results_1080p.items()))
        results_4k = dict(reversed(results_4k.items()))

    return results_720p, results_1080p, results_4k

def make_graph(results_list, title):
    fig, return_plot = plot.subplots(facecolor="#5b306f")

    return_plot.set_facecolor("#be95d1")

    return_plot.spines['top'].set_color("#ffaed2")
    return_plot.spines['bottom'].set_color("#ffaed2")
    return_plot.spines['left'].set_color("#ffaed2")
    return_plot.spines['right'].set_color("#ffaed2")
    return_plot.tick_params(axis='x', colors="#ffaed2")
    return_plot.tick_params(axis='y', colors="#ffaed2")

    bars = return_plot.barh(np.arange(len(results_list)), results_list.values(), color="#ffaed2")
    return_plot.set_yticks(np.arange(len(results_list)), results_list.keys(), color="#ffaed2")
    return_plot.set_title(title, color="#ffaed2")
    for bar, value in zip(bars, results_list.values()):
        return_plot.text(0.99, bar.get_y() + bar.get_height()/2, f'{value}x', va='center', ha='right', color="#2c1a4a", fontsize=10, fontweight='bold')

    return_plot.set_xlabel("Execution time improvement\n 1.0 value based on the first CPU", color="#ffaed2")

    return_plot.axvline(x=1, color="#2c1a4a", linewidth=1)
    fig.text(0, 0, "https://blog.neet.works", color="#ffaed2")

    plot.savefig(f"{title}.png", bbox_inches='tight', dpi=600)

    return return_plot

csv_fieldnames = ["AVX512", "Resolution", "Encoder", "Preset", "Run Time", "Temperature", "MHz", "Watts"]

results_9950X_avx = []
results_9950X_noavx = []
results_14700F = []

with open('results_9950X.csv', 'r', newline='') as results_file:
    reader = csv.DictReader(results_file, csv_fieldnames)
    for row in reader:
        if row["AVX512"] == "On":
            results_9950X_avx.append(row)
        elif row["AVX512"] == "Off":
            results_9950X_noavx.append(row)

with open('results_14700F.csv', 'r', newline='') as results_file:
    reader = csv.DictReader(results_file, csv_fieldnames)
    for row in reader:
        results_14700F.append(row)

results_14700F.pop(0)

print("Ryzen 9 9950X AVX-512 performance gain. \"Zero\" point of 1.0 based on the speed of the non-AVX-512 performance.")
print("\nHEVC encoded by the x265 encoder.\n")

list_9950X_x265_vs_avx512_720p, list_9950X_x265_vs_avx512_1080p, list_9950X_x265_vs_avx512_4k = make_comparison_list(results_9950X_noavx, results_9950X_avx, "libx265")

plot_9950X_x265_noavx_vs_9950X_avx_720p = make_graph(list_9950X_x265_vs_avx512_720p, "9950X AVX-512 Off vs 9950X AVX-512 On, 720p, x265")
plot_9950X_x265_noavx_vs_9950X_avx_1080p = make_graph(list_9950X_x265_vs_avx512_1080p, "9950X AVX-512 Off vs 9950X AVX-512 On, 1080p, x265")
plot_9950X_x265_noavx_vs_9950X_avx_4k = make_graph(list_9950X_x265_vs_avx512_4k, "9950X AVX-512 Off vs 9950X AVX-512 On, 4k, x265")

print("\nAV1 encoded by the SVT-AV1 encoder. Note: In SVT-AV1, preset 0 is the heaviest, and preset 13 is the lightest.\n")

list_9950X_av1_vs_avx512_720p, list_9950X_av1_vs_avx512_1080p, list_9950X_av1_vs_avx512_4k = make_comparison_list(results_9950X_noavx, results_9950X_avx, "libsvtav1")

plot_9950X_av1_noavx_vs_9950X_avx_720p = make_graph(list_9950X_av1_vs_avx512_720p, "9950X AVX-512 Off vs 9950X AVX-512 On, 720p, SVT-AV1")
plot_9950X_av1_noavx_vs_9950X_avx_1080p = make_graph(list_9950X_av1_vs_avx512_1080p, "9950X AVX-512 Off vs 9950X AVX-512 On, 1080p, SVT-AV1")
plot_9950X_av1_noavx_vs_9950X_avx_4k = make_graph(list_9950X_av1_vs_avx512_4k, "9950X AVX-512 Off vs 9950X AVX-512 On, 4k, SVT-AV1")

print("Ryzen 9 9950X AVX-512 performance gain vs the i7-14700F. \"Zero\" point of 1.0 based on the i7-14700F performance.")
print("\nHEVC encoded by the x265 encoder.\n")

list_9950X_x265_vs_14700F_720p, list_9950X_x265_vs_14700F_1080p, list_9950X_x265_vs_14700F_4k = make_comparison_list(results_14700F, results_9950X_avx, "libx265")

plot_9950X_avx_x265_vs_14700F_720p = make_graph(list_9950X_x265_vs_14700F_720p, "i7-14700F vs 9950X AVX-512 On, 720p, x265")
plot_9950X_avx_x265_vs_14700F_1080p = make_graph(list_9950X_x265_vs_14700F_1080p, "i7-14700F vs 9950X AVX-512 On, 1080p, x265")
plot_9950X_avx_x265_vs_14700F_4k = make_graph(list_9950X_x265_vs_14700F_4k, "i7-14700F vs 9950X AVX-512 On, 4k, x265")


print("\nAV1 encoded by the SVT-AV1 encoder. Note: In SVT-AV1, preset 0 is the heaviest, and preset 13 is the lightest.\n")

list_9950X_av1_vs_14700F_720p, list_9950X_av1_vs_14700F_1080p, list_9950X_av1_vs_14700F_4k = make_comparison_list(results_14700F, results_9950X_avx, "libsvtav1")

plot_9950X_avx_av1_vs_14700F_720p = make_graph(list_9950X_av1_vs_14700F_720p, "i7-14700F vs 9950X AVX-512 On, 720p, SVT-AV1")
plot_9950X_avx_av1_vs_14700F_1080p = make_graph(list_9950X_av1_vs_14700F_1080p, "i7-14700F vs 9950X AVX-512 On, 1080p, SVT-AV1")
plot_9950X_avx_av1_vs_14700F_4k = make_graph(list_9950X_av1_vs_14700F_4k, "i7-14700F vs 9950X AVX-512 On, 4k, SVT-AV1")

geomean_dict = {
    "9950X AVX-512 Off vs 9950X AVX-512 On, 720p, x265": geomean(list_9950X_x265_vs_avx512_720p),
    "9950X AVX-512 Off vs 9950X AVX-512 On, 1080p, x265": geomean(list_9950X_x265_vs_avx512_1080p),
    "9950X AVX-512 Off vs 9950X AVX-512 On, 4k, x265": geomean(list_9950X_x265_vs_avx512_4k),
    "9950X AVX-512 Off vs 9950X AVX-512 On, 720p, SVT-AV1": geomean(list_9950X_av1_vs_avx512_720p),
    "9950X AVX-512 Off vs 9950X AVX-512 On, 1080p, SVT-AV1": geomean(list_9950X_av1_vs_avx512_1080p),
    "9950X AVX-512 Off vs 9950X AVX-512 On, 4k, SVT-AV1": geomean(list_9950X_av1_vs_avx512_4k),
    "i7-14700F vs 9950X AVX-512 On, 720p, x265": geomean(list_9950X_x265_vs_14700F_720p),
    "i7-14700F vs 9950X AVX-512 On, 1080p, x265": geomean(list_9950X_x265_vs_14700F_1080p),
    "i7-14700F vs 9950X AVX-512 On, 4k, x265": geomean(list_9950X_x265_vs_14700F_4k),
    "i7-14700F vs 9950X AVX-512 On, 720p, SVT-AV1": geomean(list_9950X_av1_vs_14700F_720p),
    "i7-14700F vs 9950X AVX-512 On, 1080p, SVT-AV1": geomean(list_9950X_av1_vs_14700F_1080p),
    "i7-14700F vs 9950X AVX-512 On, 4k, SVT-AV1": geomean(list_9950X_av1_vs_14700F_4k)
}

geomean_dict = dict(reversed(geomean_dict.items()))

geomean_graph = make_graph(geomean_dict, "Geomean performance difference at major resolutions, across all presets")

#plot.show()