#!/usr/bin/env python
"""Write the training pool's manifest: its segments, the validation set, frame paths.

    pool_manifest.py --availability .../corpus_v9/availability.npz \\
                     --corpus-manifest .../corpus_v9/manifest.json \\
                     --out .../corpus_v9/pool_manifest.json.zst

Read-only against the index.  The pool draws its training slices from these
segments (:mod:`suvi.pool`).  Validation is a fixed set of windows -- in the corpus
manifest's val segments and in more spread over the era -- written once to disk by
``pool_producer.py --fixed``; the pool keeps clear of every validation segment, of
the corpus's test segments, and of the bench holdout days.
"""

import argparse
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import dataset_build
from suvi import availability, db, detectors, pool


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--availability", required=True)
    parser.add_argument("--corpus-manifest", required=True,
                        help="the on-disk corpus's manifest.json (its val/test stay out)")
    parser.add_argument("--out", required=True)
    parser.add_argument("--run", type=int, default=None,
                        help="judgment run (default: the one under the current rules)")
    parser.add_argument("--db", default=None)
    args = parser.parse_args(argv)

    conn = db.connect(args.db, readonly=True)
    run_id = args.run or detectors.judgment_run(conn)
    if run_id is None:
        raise SystemExit("no judgment run under the current rules")
    corpus = json.load(open(args.corpus_manifest))
    model = availability.Availability.load(args.availability)
    manifest = pool.build_manifest(conn, model, run_id, corpus["span_min"],
                                   corpus["segments"]["val"], corpus["segments"]["test"],
                                   dataset_build.holdout_days(),
                                   satellites=tuple(corpus["satellites"]),
                                   old_train=corpus["segments"]["train"])
    pool.save_manifest(manifest, args.out)
    span = manifest["span"]
    print(f"{len(manifest['segments'])} segments, "
          f"{pool.target_count(manifest['segments'], span):,} targets per satellite at "
          f"span {span} (the corpus's train split: "
          f"{pool.target_count(corpus['segments']['train'], span):,}); validation "
          f"{len(manifest['val'])} windows in {len(manifest['val_segments'])} segments, "
          f"{pool.target_count(manifest['val'], span):,} targets per satellite (the "
          f"corpus's: {pool.target_count(corpus['segments']['val'], span):,}); judgment run "
          f"{run_id}; wrote {args.out} ({os.path.getsize(args.out) / 1e6:.0f} MB)")
    return 0


if __name__ == "__main__":
    sys.exit(main())