from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import trackio
from experiment import ROOMS, run_agent
from safetensors.torch import save_file

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "curiosity-cabinet"
DATA_DIR = PROJECT_DIR / "data"


def aggregate(records: list[dict], reward: str) -> dict:
    subset = [record for record in records if record["reward"] == reward]
    return {
        "runs": len(subset),
        "mean_middle_action_fraction": {
            room: float(
                np.mean(
                    [
                        record["middle_action_fraction"][room]
                        for record in subset
                    ]
                )
            )
            for room in ROOMS
        },
        "mean_final_action_fraction": {
            room: float(
                np.mean(
                    [record["final_action_fraction"][room] for record in subset]
                )
            )
            for room in ROOMS
        },
        "learnable_world_model_mse": {
            "mean": float(
                np.mean(
                    [record["learnable_world_model_mse"] for record in subset]
                )
            ),
            "median": float(
                np.median(
                    [record["learnable_world_model_mse"] for record in subset]
                )
            ),
        },
    }


def main() -> None:
    runs = 60
    steps = 1_200
    trackio.init(
        project="curiosity-cabinet",
        name="error-versus-learning-progress-v1",
        config={
            "runs_per_reward": runs,
            "steps_per_run": steps,
            "rooms": ROOMS,
            "epsilon": 0.15,
        },
    )
    records = []
    representative = {}
    for run in range(runs):
        for reward in ["prediction_error", "learning_progress"]:
            record = run_agent(reward, seed=9000 + run, steps=steps)
            if run == 0:
                representative[reward] = record
            records.append(
                {
                    key: value
                    for key, value in record.items()
                    if key not in {"actions", "losses", "learnable_state_dict"}
                }
            )
        if (run + 1) % 10 == 0:
            recent = records[-20:]
            trackio.log(
                {
                    "completed_runs": run + 1,
                    "error_noisy_tv_fraction": float(
                        np.mean(
                            [
                                item["middle_action_fraction"]["noisy_tv"]
                                for item in recent
                                if item["reward"] == "prediction_error"
                            ]
                        )
                    ),
                    "progress_noisy_tv_fraction": float(
                        np.mean(
                            [
                                item["middle_action_fraction"]["noisy_tv"]
                                for item in recent
                                if item["reward"] == "learning_progress"
                            ]
                        )
                    ),
                }
            )
    summary = {
        reward: aggregate(records, reward)
        for reward in ["prediction_error", "learning_progress"]
    }
    report = {
        "benchmark": "Curiosity Cabinet noisy-TV stress test",
        "runs_per_reward": runs,
        "steps_per_run": steps,
        "summary": summary,
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    for reward, record in representative.items():
        save_file(
            record["learnable_state_dict"],
            ARTIFACT_DIR / f"{reward}_world_model.safetensors",
        )
        np.savez_compressed(
            ARTIFACT_DIR / f"{reward}_trajectory.npz",
            actions=np.asarray(record["actions"], dtype=np.int8),
            losses=np.asarray(record["losses"], dtype=np.float32),
        )
    rows = []
    for record in records:
        row = {
            "reward": record["reward"],
            "seed": record["seed"],
            "learnable_world_model_mse": record["learnable_world_model_mse"],
        }
        for window in ["middle_action_fraction", "final_action_fraction"]:
            for room in ROOMS:
                row[f"{window}_{room}"] = record[window][room]
        rows.append(row)
    pd.DataFrame(rows).to_parquet(DATA_DIR / "seeded_results.parquet", index=False)
    trackio.log(
        {
            "error_noisy_tv_fraction": summary["prediction_error"][
                "mean_middle_action_fraction"
            ]["noisy_tv"],
            "progress_noisy_tv_fraction": summary["learning_progress"][
                "mean_middle_action_fraction"
            ]["noisy_tv"],
            "error_learnable_mse": summary["prediction_error"][
                "learnable_world_model_mse"
            ]["mean"],
            "progress_learnable_mse": summary["learning_progress"][
                "learnable_world_model_mse"
            ]["mean"],
        }
    )
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()