import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path

OUT_DIR = Path("/home/schwarc/diplomovka/lmeval_plots")
OUT_DIR.mkdir(parents=True, exist_ok=True)

MODEL = "slovak-nlp/mistral-sk-7b"
ADAPTER = "Jakub1320/mistral-sk-7b-slovak-alpaca-qlora"

rows = [
    {"task": "arc_sk", "samples": 1169, "metric": "acc", "value": 0.3259},
    {"task": "arc_sk", "samples": 1169, "metric": "acc_norm", "value": 0.3473},

    {"task": "hellaswag_sk", "samples": 9485, "metric": "acc", "value": 0.4134},
    {"task": "hellaswag_sk", "samples": 9485, "metric": "acc_norm", "value": 0.5106},

    {"task": "m_mmlu_sk", "samples": 13062, "metric": "acc", "value": 0.2928},

    {"task": "truthfulqa_sk_mc1", "samples": 778, "metric": "acc", "value": 0.2391},
    {"task": "truthfulqa_sk_mc2", "samples": 778, "metric": "acc", "value": 0.3942},

    {"task": "sklegal", "samples": 1334, "metric": "acc", "value": 0.2841},
    {"task": "sklegal", "samples": 1334, "metric": "acc_norm", "value": 0.4693},

    {"task": "skquad", "samples": 9583, "metric": "exact", "value": 1.9722},
    {"task": "skquad", "samples": 9583, "metric": "f1", "value": 27.7994},
    {"task": "skquad", "samples": 9583, "metric": "HasAns_exact", "value": 2.4106},
    {"task": "skquad", "samples": 9583, "metric": "HasAns_f1", "value": 34.1456},
    {"task": "skquad", "samples": 9583, "metric": "NoAns_exact", "value": 0.0561},
    {"task": "skquad", "samples": 9583, "metric": "NoAns_f1", "value": 0.0561},
    {"task": "skquad", "samples": 9583, "metric": "best_exact", "value": 18.6163},
    {"task": "skquad", "samples": 9583, "metric": "best_f1", "value": 28.3012},
]

df = pd.DataFrame(rows)
df["model"] = MODEL
df["adapter"] = ADAPTER

csv_path = OUT_DIR / "lmeval_results_full_table.csv"
df.to_csv(csv_path, index=False)

# Skrátený výber hlavných metrík pre graf

main_rows = [
    {"benchmark": "ARC-SK", "metric": "acc_norm", "score_percent": 34.73},
    {"benchmark": "HellaSwag-SK", "metric": "acc_norm", "score_percent": 51.06},
    {"benchmark": "M-MMLU-SK", "metric": "acc", "score_percent": 29.28},
    {"benchmark": "TruthfulQA-SK MC1", "metric": "acc", "score_percent": 23.91},
    {"benchmark": "TruthfulQA-SK MC2", "metric": "acc", "score_percent": 39.42},
    {"benchmark": "SKLegal", "metric": "acc_norm", "score_percent": 46.93},
    {"benchmark": "SKQuAD", "metric": "f1", "score_percent": 27.80},
]

main_df = pd.DataFrame(main_rows)
main_csv_path = OUT_DIR / "lmeval_main_metrics_table.csv"
main_df.to_csv(main_csv_path, index=False)

# Graf hlavných metrík
labels = main_df["benchmark"] + " (" + main_df["metric"] + ")"
scores = main_df["score_percent"]

plt.figure(figsize=(11, 6))
bars = plt.barh(labels, scores)

plt.xlabel("Skóre (%)")
plt.ylabel("Benchmark")
plt.title("Výsledky lm-evaluation-harness pre PEFT LoRA model")
plt.xlim(0, 60)
plt.gca().invert_yaxis()

for bar, score in zip(bars, scores):
    plt.text(
        bar.get_width() + 0.8,
        bar.get_y() + bar.get_height() / 2,
        f"{score:.2f} %",
        va="center"
    )

plt.figtext(
    0.01,
    0.01,
    f"Model: {MODEL} | Adaptér: {ADAPTER}",
    fontsize=8
)

plt.tight_layout(rect=[0, 0.04, 1, 1])

png_path = OUT_DIR / "lmeval_main_metrics_chart.png"
pdf_path = OUT_DIR / "lmeval_main_metrics_chart.pdf"

plt.savefig(png_path, dpi=300)
plt.savefig(pdf_path)
plt.close()

print("Hotovo.")
print(f"CSV kompletna tabulka: {csv_path}")
print(f"CSV hlavne metriky:     {main_csv_path}")
print(f"Graf PNG:              {png_path}")
print(f"Graf PDF:              {pdf_path}")