"""Read actual training events and inspect saved checkpoint tensors on CPU."""

import json
from pathlib import Path

import torch
from safetensors import safe_open
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

from portrait import ROOT, save_json, sha256, validate_dataset


def main():
    output = ROOT / "outputs/run01"
    hashes = json.loads((output / "weight-hashes.json").read_text())
    weights = []
    for filename, digest in hashes.items():
        path = output / filename
        assert sha256(path) == digest, filename
        with safe_open(path, framework="pt", device="cpu") as handle:
            metadata = handle.metadata()
            keys = list(handle.keys())
            finite = all(bool(torch.isfinite(handle.get_tensor(key)).all()) for key in keys)
            assert finite, filename
            weights.append({"file": filename, "sha256": digest, "bytes": path.stat().st_size,
                            "steps": int(metadata["ss_steps"]), "rank": int(metadata["ss_network_dim"]),
                            "tensor_count": len(keys), "all_tensors_finite": finite})
    scalar_runs = []
    for path in output.rglob("events.out.tfevents.*"):
        events = EventAccumulator(str(path), size_guidance={"scalars": 0}).Reload()
        scalar_runs.append({"file": str(path.relative_to(ROOT)), "scalars": {
            tag: [{"step": item.step, "value": item.value, "wall_time": item.wall_time}
                  for item in events.Scalars(tag)] for tag in events.Tags()["scalars"]}})
    timeline = [json.loads(line) for line in (ROOT / "logs/events.jsonl").read_text().splitlines()]
    summary = {"dataset_images": len(validate_dataset(ROOT / "dataset")), "weights": weights,
               "training_events": [row for row in timeline if row["stage"].startswith("training_")],
               "tensorboard": scalar_runs}
    save_json(ROOT / "logs/training-summary.json", summary)
    print(json.dumps({"weights": weights, "scalar_tags": [list(r["scalars"]) for r in scalar_runs]}))


if __name__ == "__main__":
    main()
