import argparse
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

# # Argparse with the parameter plot_type
parser = argparse.ArgumentParser()
parser.add_argument("--year", "-y", type=int, required=True, help="Year of ENEM exam")
args = parser.parse_args()
year = args.year

dic_exam_name = {'CH': 'Humanities', 'LC': 'Languages and Codes', 'CN': 'Natural Sciences', 'MT': 'Mathematics'}

itens_llm_df = pd.read_csv(f"rank_llms_{year}.csv")
itens_human_df = pd.read_csv(f"rank_humans_{year}.csv")

curves_llms = {}
curves_humans = {}
for subject in ['MT', 'CN', 'CH', 'LC']:
    for year in [year]:
        curves_llms[f'{dic_exam_name[subject]}'] = itens_llm_df[itens_llm_df['subject'] == subject]['discrimination_index']
        curves_humans[f'{dic_exam_name[subject]}'] = itens_human_df[itens_human_df['subject'] == subject]['discrimination_index']

plt.rcParams.update({'font.size': 6})
fig, ax = plt.subplots(1, 2, figsize=(1.4*3.5, 1.65), sharex=True, sharey=True, layout="compressed")

sns.kdeplot(
        curves_humans,
        bw_adjust=0.6,
        common_norm=False,  # Normalize each distribution independently
        ax=ax[0],
        hue_order=sorted(curves_humans.keys())
)

sns.kdeplot(
        curves_llms,
        bw_adjust=0.6,
        common_norm=False,  # Normalize each distribution independently
        ax=ax[1],
        hue_order=sorted(curves_llms.keys())
)

# Remove legend of ax[0]
ax[0].get_legend().remove()

sns.move_legend(ax[1], loc='upper left', ncol=2, fontsize=5, title_fontsize=0, title="")

ax[0].set_title(f"Humans ({year})", fontsize=6)
ax[1].set_title(f"LLMs ({year})", fontsize=6)
ax[0].set_xlabel("Discrimination Index")
ax[1].set_xlabel("Discrimination Index")
ax[0].set_ylabel("Density")
ax[1].set_ylabel("Density")
ax[0].set_xticks([-0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
ax[1].set_xticks([-0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
ax[0].set_yticks([0, 2, 4, 6])
ax[1].set_yticks([0, 2, 4, 6])
plt.tight_layout()
_= plt.savefig(f'../plots/discrimination-indexes/di-llms-humans-{year}.pdf', format='pdf',bbox_inches='tight')

