"""Evaluate a C1 checkpoint only on the frozen nonblind snapshots."""
import argparse
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))


def main():
    from vimeml.posttraining.infer import (TRACKS, read_snapshot, load_conditional_checkpoint,
                                         fp32_evaluation, evaluate_snapshot)
    from vimeml.training.data import write_json
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--tracks", choices=tuple(TRACKS), nargs="+", default=["standard-development"])
    parser.add_argument("--condition", choices=("context_reading", "empty_reading", "empty_context"), default="context_reading")
    args = parser.parse_args()
    if args.output.exists() and any(args.output.iterdir()):
        parser.error("Use a new evaluation output directory")
    rows = {track: read_snapshot(track) for track in args.tracks}
    model, codec, metadata = load_conditional_checkpoint(args.checkpoint,
        ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192/tokenizer.model", args.device)
    report = {"checkpoint": metadata, "condition": args.condition, "tracks": {}}
    with fp32_evaluation(model):
        for track, data in rows.items():
            report["tracks"][track] = evaluate_snapshot(model, codec, data, args.output / track, args.condition)
    write_json(args.output / "summary.json", report)
    print(report)


if __name__ == "__main__":
    main()