#!/usr/bin/env python
"""Keep a rotating pool of training slices in memory, generated from the archive.

    pool_producer.py --manifest /work/pool_manifest.json.zst --archive /archive/Data \\
                     --pool /dev/shm/suvi_pool --budget-gb 48 --workers 16

Runs on the training host beside ``train.py --pool``: picks a random slice of a
pristine segment (:func:`suvi.pool.choose_slice`), reads its six bands per slot for
both satellites from the archive share, encodes them exactly as the on-disk corpus
does (:func:`suvi.dataset.encode_frames`), writes the slice, then retires the oldest
slices past the budget.  Nothing is judged here: the manifest lists only frames the
stored judgment called good.  A frame that cannot be read is an absent slot, which
splits its slice's segment rather than stopping the pool.

With ``--fixed val`` (or ``test``) it instead writes the manifest's validation
windows (or test segments), once each, to a disk directory (``val000`` ...), skipping
any already complete, and exits: evaluation sets holding nothing but their own slots.

See :mod:`suvi.pool` for the design and the measured rates.
"""

import argparse
import os
import sys
import time
from multiprocessing import Pool

import numpy as np

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

from suvi import dataset, fitsio, paths, pool

#: Slots encoded per worker dispatch.
CHUNK = 4


def encode_slot(files):
    """Worker: six archive paths -> a 6x640x640 block, or None if any band fails."""
    arrays = []
    for path in files:
        try:
            image, _ = fitsio.read_image(path)
        except OSError:
            return None
        if image is None:
            return None
        arrays.append(image)
    try:
        return dataset.encode_frames(arrays)
    except ValueError:
        return None


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--manifest", required=True, help="from pool_manifest.py")
    parser.add_argument("--archive", required=True, help="the share's Data directory")
    parser.add_argument("--pool", required=True, help="a RAM directory, e.g. /dev/shm/...")
    parser.add_argument("--budget-gb", type=float, default=48.0)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--slices", type=int, default=None,
                        help="stop after this many (default: run until stopped)")
    parser.add_argument("--fixed", choices=("val", "test"), default=None,
                        help="write this evaluation set from the manifest to --pool "
                             "and exit")
    parser.add_argument("--seed", type=int, default=None,
                        help="default: from the clock, so restarts draw new slices")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    fitsio.quiet_astropy()
    manifest = pool.load_manifest(args.manifest)
    satellites = tuple(manifest["satellites"])
    span = manifest["span"]
    segments = [tuple(pair) for pair in manifest["segments"]]
    frames = manifest["frames"]
    rng = np.random.default_rng(args.seed if args.seed is not None else time.time_ns())
    os.makedirs(args.pool, exist_ok=True)
    pool.clear_partials(args.pool)
    budget = args.budget_gb * 1e9
    print(f"pool {args.pool}: {len(segments)} segments, "
          f"{pool.target_count(segments, span):,} targets per satellite at span {span}; "
          f"budget {args.budget_gb:.0f} GB, {args.workers} readers", flush=True)

    counter = len(pool.ready_slices(args.pool))
    made = 0
    started = time.time()
    slots_done = 0
    with Pool(args.workers, initializer=fitsio.quiet_astropy,
              maxtasksperchild=500) as workers:
        def blocks(satellite, times):
            files = [[os.path.join(args.archive, rel)
                      for rel in frames[str(satellite)][str(t)]] for t in times]
            yield from zip(times, workers.imap(encode_slot, files, chunksize=CHUNK))

        if args.fixed:
            windows = [tuple(pair) for pair in manifest[args.fixed]]
            for index, piece in enumerate(windows):
                name = f"{args.fixed}{index:03d}"
                if os.path.exists(os.path.join(args.pool, name, pool.SLICE_FILE)):
                    continue
                path = pool.write_slice(args.pool, name, piece, satellites, blocks, span)
                print(f"  {name} ({index + 1}/{len(windows)}): "
                      f"{'written' if path else 'nothing usable, dropped'}", flush=True)
            return 0

        while args.slices is None or made < args.slices:
            piece = pool.choose_slice(segments, span, rng)
            counter += 1
            path = pool.write_slice(args.pool, pool.slice_name(counter, piece), piece,
                                    satellites, blocks, span)
            made += 1
            slots_done += (piece[1] - piece[0]) // paths.CADENCE * len(satellites)
            gone = pool.retire(args.pool, budget)
            resident = pool.ready_slices(args.pool)
            rate = slots_done / max(time.time() - started, 1e-6)
            print(f"  {os.path.basename(path) if path else 'slice dropped'}: "
                  f"{len(resident)} resident ({sum(s for _, s in resident) / 1e9:.1f} GB), "
                  f"{len(gone)} retired, {rate:.1f} slots/s", flush=True)
    return 0


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