import pandas as pd
from matplotlib import pyplot as plt, colors
from matplotlib.font_manager import FontProperties
import numpy as np

subjects = ["Maths", "Reading", "Science"]

data = pd.read_csv("./practice1.csv", names=["country", *subjects])

country_codes = pd.read_csv("./country_codes.csv", names=["country", "code"])

data = data.merge(country_codes, how="outer", on="country")

data = data.dropna(subset=subjects)

data["Average"] = data.loc[:, subjects].mean(axis=1)

data = data.sort_values(by="Average", ascending=False)

data.to_csv("res.csv", index=False, header=False)

if data[data["code"].isna()].size != 0:
    raise ValueError("NaN country code")

fig, ax = plt.subplots(nrows=4)

fig.set_size_inches(16, 11.69)

positions = np.array(list(map(str,range(1, len(data.index)+1))))

x = data["code"] + "\n" + positions

# colors=("#ef476f", "#ffd166", "#06d6a0", "#118ab2", "#073B4C")

colors=("tab:blue", "tab:orange", "tab:green", "tab:red", "tab:grey")

subjects_labels = ("🔢", "📚", "🔬", "")

for i, subj in enumerate((*subjects, "Average")):
    bar = ax[i].bar(x, data[subj], color=colors[i])

    bar_labels = list(map(lambda j: str(int(data[subj][data.index[j]])) if data["code"][data.index[j]] == "UK" else "", range(len(data.index))))

    ax[i].bar_label(bar, labels=bar_labels, weight="bold")

    ax[i].set_title(subjects_labels[i])

    ax[i].autoscale(tight=True)

    ax[i].set_ylim(350)

    mean = int(data[subj].mean())

    ax[i].set_yticks(list(set((*ax[i].get_yticks(), mean, 650))))

    for label in ax[i].get_xticklabels():
        if "UK" in label.get_text():
            label.set(fontweight="bold")

    for label in ax[i].get_yticklabels():
        if label.get_text() == str(mean):
            label.set(fontweight="bold")

    ax[i].hlines([mean], 0, 1, transform=ax[i].get_yaxis_transform(), color=colors[-1])

    ax[i].spines['top'].set_visible(False)
    ax[i].spines['right'].set_visible(False)
    ax[i].spines['bottom'].set_visible(False)
    ax[i].spines['left'].set_visible(False)
    # ax[i].grid()

fig.tight_layout()

fig.savefig("test.png")