from pathlib import Path
import csv
import json
import math


ROOT = Path(
    "/home/schwarc/diplomovka/"
    "aya_300k_lr_validation"
)

RESULT_DIR = ROOT / "results"
RESULT_DIR.mkdir(
    parents=True,
    exist_ok=True
)


RUNS = [
    ("0.00018", ROOT / "lr_0p00018"),
    ("0.00020", ROOT / "lr_0p00020"),
    ("0.00023", ROOT / "lr_0p00023"),
]


def load_best_state(output):

    files = list(
        output.rglob(
            "trainer_state.json"
        )
    )

    if not files:
        return None

    candidates = []

    for p in files:
        try:
            data = json.loads(
                p.read_text(
                    encoding="utf-8"
                )
            )

            candidates.append(
                (
                    int(
                        data.get(
                            "global_step",
                            0
                        )
                    ),
                    data,
                    p,
                )
            )

        except Exception:
            pass

    if not candidates:
        return None

    return max(
        candidates,
        key=lambda x: x[0]
    )


rows = []
eval_curves = {}
grad_curves = {}


for lr, output in RUNS:

    result = load_best_state(
        output
    )

    if result is None:
        continue

    global_step, state, state_file = result

    history = state.get(
        "log_history",
        []
    )

    evals = []
    grads = []

    for item in history:

        step = int(
            item.get("step", 0)
        )

        if "eval_loss" in item:
            try:
                value = float(
                    item["eval_loss"]
                )
                if math.isfinite(value):
                    evals.append(
                        (step, value)
                    )
            except Exception:
                pass

        if "grad_norm" in item:
            try:
                value = float(
                    item["grad_norm"]
                )
                if math.isfinite(value):
                    grads.append(
                        (step, value)
                    )
            except Exception:
                pass


    if evals:
        final_eval_step, final_eval = (
            evals[-1]
        )

        min_eval_step, min_eval = min(
            evals,
            key=lambda x: x[1]
        )
    else:
        final_eval_step = ""
        final_eval = ""
        min_eval_step = ""
        min_eval = ""


    grad_values = [
        value
        for _, value in grads
    ]

    if grad_values:

        max_grad = max(
            grad_values
        )

        ordered = sorted(
            grad_values
        )

        p95_index = int(
            0.95
            * (len(ordered) - 1)
        )

        p95_grad = (
            ordered[p95_index]
        )

    else:
        max_grad = ""
        p95_grad = ""


    rows.append(
        {
            "learning_rate": lr,
            "global_step": global_step,
            "final_eval_step":
                final_eval_step,
            "final_eval_loss":
                final_eval,
            "min_eval_step":
                min_eval_step,
            "min_eval_loss":
                min_eval,
            "max_grad_norm":
                max_grad,
            "p95_grad_norm":
                p95_grad,
            "trainer_state":
                str(state_file),
        }
    )

    eval_curves[lr] = evals
    grad_curves[lr] = grads


csv_path = (
    RESULT_DIR
    / "aya_300k_lr_summary.csv"
)

with csv_path.open(
    "w",
    encoding="utf-8",
    newline=""
) as f:

    fields = [
        "learning_rate",
        "global_step",
        "final_eval_step",
        "final_eval_loss",
        "min_eval_step",
        "min_eval_loss",
        "max_grad_norm",
        "p95_grad_norm",
        "trainer_state",
    ]

    writer = csv.DictWriter(
        f,
        fieldnames=fields
    )

    writer.writeheader()
    writer.writerows(rows)


print()
print("===== 300K LR SUMMARY =====")

for row in rows:
    print(
        f"LR={row['learning_rate']} | "
        f"step={row['global_step']} | "
        f"final_eval={row['final_eval_loss']} | "
        f"min_eval={row['min_eval_loss']} "
        f"@ {row['min_eval_step']} | "
        f"max_grad={row['max_grad_norm']} | "
        f"p95_grad={row['p95_grad_norm']}"
    )

print()
print("CSV:", csv_path)


try:
    import matplotlib.pyplot as plt

    if eval_curves:

        plt.figure(
            figsize=(12, 7)
        )

        for lr, points in (
            eval_curves.items()
        ):
            if not points:
                continue

            plt.plot(
                [x for x, _ in points],
                [y for _, y in points],
                marker="o",
                label=f"LR {lr}",
            )

        plt.xlabel(
            "optimizer step"
        )
        plt.ylabel(
            "eval loss"
        )
        plt.title(
            "Aya 300k - eval loss"
        )
        plt.legend()
        plt.grid(alpha=0.25)
        plt.tight_layout()

        plt.savefig(
            RESULT_DIR
            / "eval_loss_over_steps.png",
            dpi=200,
            bbox_inches="tight"
        )

        plt.close()


    if grad_curves:

        plt.figure(
            figsize=(12, 7)
        )

        for lr, points in (
            grad_curves.items()
        ):
            if not points:
                continue

            plt.plot(
                [x for x, _ in points],
                [y for _, y in points],
                label=f"LR {lr}",
            )

        plt.xlabel(
            "optimizer step"
        )
        plt.ylabel(
            "grad norm"
        )
        plt.title(
            "Aya 300k - gradient norm"
        )
        plt.yscale("log")
        plt.legend()
        plt.grid(alpha=0.25)
        plt.tight_layout()

        plt.savefig(
            RESULT_DIR
            / "grad_norm_over_steps.png",
            dpi=200,
            bbox_inches="tight"
        )

        plt.close()


except Exception as exc:
    print(
        "Grafy sa nepodarilo vytvorit:",
        exc
    )