"""A rotating in-memory training pool, generated from the archive while training runs.

The on-disk corpus holds ~80,000 slots (323 GB) because that is what fits on the
training host; the archive's pristine two-satellite era is ~5.6 times larger.  So
instead of a fixed corpus the host can generate training data as it trains: a
producer (``pool_producer.py``) reads FITS straight from the archive share, writes
*slices* -- a day of one pristine segment, both satellites, with the model's span of
margin -- into a RAM directory, and retires the oldest when the pool is over budget.
``train.py --pool`` samples from whatever slices are resident.  Validation and test
stay on the fixed on-disk corpus, so scores remain comparable between runs.

A slice is an ordinary little corpus: day-shards in the usual format plus
``slice.json`` naming its verified segments.  It is built in ``.<name>.partial`` and
renamed into place complete, so a reader never sees half of one; a reader opens its
shards eagerly (:class:`suvi.dataset.Shard`), so a slice retired mid-visit stays
readable through the handles already open.

Measured on the host (2026-10-13): 16 workers produce ~8.5 slots/s from the share
(~90 MB/s, near gigabit); training at the v9 pace consumes ~5 new slots/s.
"""

import bisect
import datetime
import json
import os
import shutil
import time

import numpy as np

from . import corpus, dataset, paths, segments as segments_module

#: Targets per slice: one day of one satellite, so a slice holds one evolving view
#: of the Sun and the pool turns over many of them.
SLICE_TARGETS = 360
#: Slices kept clear of the evaluation segments, in seconds either side: a pristine
#: run separated from a val segment by a few bad slots shows the same Sun.
GUARD_SECONDS = 86400
#: The validation set (v10 on): one window of VAL_TARGETS targets in each of the
#: previous corpus's val segments and in VAL_NEW more spread evenly over the era --
#: ~5,800 targets per satellite over 64 segments, against v9's 1,714 over 16.
VAL_NEW = 48
VAL_TARGETS = 90
VAL_SEED = 11
#: Manifest format version.
MANIFEST_VERSION = 2
SLICE_FILE = "slice.json"


# ----------------------------------------------------------------- the manifest


def subtract(segments, cuts, minimum):
    """`segments` with every (low, high) of `cuts` removed, widened by
    :data:`GUARD_SECONDS`; pieces shorter than `minimum` slots are dropped."""
    out = []
    for a, b in segments:
        pieces = [(a, b)]
        for low, high in cuts:
            low, high = low - GUARD_SECONDS, high + GUARD_SECONDS
            pieces = [part for start, end in pieces
                      for part in ((start, min(end, low)), (max(start, high), end))
                      if part[1] > part[0]]
        out.extend((start, start + (end - start) // paths.CADENCE * paths.CADENCE)
                   for start, end in pieces
                   if (end - start) // paths.CADENCE >= minimum)
    return out


def build_manifest(conn, model, run_id, span, old_val, test, holdout_days,
                   satellites=(16, 18), val_new=None, val_targets=None, seed=None,
                   old_train=()):
    """The pool's segments and a fixed validation set, with every slot's frame paths.

    `model` is the availability model of judgment run `run_id`.  Validation is one
    window of `val_targets` targets (plus the span either side) in each of the
    `old_val` segments -- the previous corpus's, so its days stay covered -- and in
    `val_new` further segments spread evenly over the era, none of them touching the
    previous corpus's `old_train` (so the previous model can be scored on the new
    set fairly).  The pool keeps :data:`GUARD_SECONDS` clear of every validation
    window and `test` segment, and off the bench's `holdout_days`.  Paths are the run's
    good, non-stub frames, one per (satellite, band, slot); a slot missing one on
    any satellite splits its segment or window, so the producer judges nothing.
    The `test` segments are kept whole, as the evaluation set's own slices.
    """
    val_new = VAL_NEW if val_new is None else val_new
    val_targets = VAL_TARGETS if val_targets is None else val_targets
    seed = VAL_SEED if seed is None else seed
    minimum = 2 * span + 1
    window = val_targets + 2 * span
    held = set(holdout_days)

    def near(segment, others):
        a, b = segment
        return any(a < high + GUARD_SECONDS and b > low - GUARD_SECONDS
                   for low, high in others)

    candidates = [pair for pair in segments_module.from_availability(model, minimum)
                  if not set(segments_module.days_touched([pair])) & held]
    old_val, test = [tuple(p) for p in old_val], [tuple(p) for p in test]
    old_train = [tuple(p) for p in old_train]
    eligible = [pair for pair in candidates if not near(pair, old_val + test + old_train)
                and (pair[1] - pair[0]) // paths.CADENCE >= window]
    count = min(val_new, len(eligible))
    new_val = [eligible[int((i + 0.5) * len(eligible) / count)] for i in range(count)]
    val_segments = sorted(old_val + new_val)
    rng = np.random.default_rng(seed)
    windows = []
    for a, b in val_segments:
        length = min(window, (b - a) // paths.CADENCE)
        first = int(rng.integers(0, (b - a) // paths.CADENCE - length + 1))
        windows.append((a + first * paths.CADENCE,
                        a + (first + length) * paths.CADENCE))
    chosen = subtract(candidates, windows + test, minimum)

    wanted = sorted(chosen + windows + test)
    frames = {str(s): {} for s in satellites}
    starts = [a for a, _ in wanted]
    marks = ",".join("?" * len(satellites))
    # One pass over the era, not a query per segment: the planner drives a range
    # query from the run's detections, so each of ~700 would rescan all 3 M of them.
    for row in conn.execute(
        f"SELECT f.satellite, f.wavelength, f.t_start, f.path FROM detection d "
        f"JOIN frame f ON f.id = d.frame_id "
        f"WHERE d.run_id = ? AND d.verdict = 'good' AND f.satellite IN ({marks}) "
        f"AND (f.size_bytes IS NULL OR f.size_bytes >= ?)",
        (run_id, *satellites, dataset.MIN_FRAME_BYTES),
    ) if wanted else ():
        t = row["t_start"]
        index = bisect.bisect_right(starts, t) - 1
        if index < 0 or t >= wanted[index][1]:
            continue
        slot = frames[str(row["satellite"])].setdefault(
            str(t), [None] * len(paths.WAVELENGTHS))
        slot[paths.WAVELENGTHS.index(row["wavelength"])] = row["path"]

    def complete(t):
        return all(None not in frames[str(s)].get(str(t), [None]) for s in satellites)

    verified = segments_module.verify(chosen, complete, minimum_slots=minimum)
    val = segments_module.verify(windows, complete, minimum_slots=minimum)
    tested = segments_module.verify(test, complete, minimum_slots=minimum)
    keep = {str(t) for a, b in verified + val + tested
            for t in range(a, b, paths.CADENCE)}
    frames = {s: {t: v for t, v in slots.items() if t in keep}
              for s, slots in frames.items()}
    return {"version": MANIFEST_VERSION, "satellites": list(satellites),
            "span": span, "judgment_run": int(run_id),
            "segments": [list(pair) for pair in verified],
            "val": [list(pair) for pair in val],
            "val_segments": [list(pair) for pair in val_segments],
            "test": [list(pair) for pair in tested],
            "frames": frames}


def save_manifest(manifest, path):
    compress, _ = dataset._zstd()
    payload = compress(json.dumps(manifest).encode(), 9)
    temporary = f"{path}.partial"
    with open(temporary, "wb") as handle:
        handle.write(payload)
    os.replace(temporary, path)


def load_manifest(path):
    _, decompress = dataset._zstd()
    with open(path, "rb") as handle:
        manifest = json.loads(decompress(handle.read()))
    if manifest.get("version") != MANIFEST_VERSION:
        raise ValueError(f"{path}: pool manifest version {manifest.get('version')}, "
                         f"expected {MANIFEST_VERSION}; rebuild it with pool_manifest.py")
    return manifest


def target_count(segments, span):
    """Targets the segments hold at `span`, per satellite."""
    return sum(max((b - a) // paths.CADENCE - 2 * span, 0) for a, b in segments)


# ------------------------------------------------------------------- slices


def choose_slice(segments, span, rng, targets=SLICE_TARGETS):
    """A random (t0, t1) of at most `targets` targets plus `span` slots of margin
    either side, inside one segment.  Segments are drawn by target count, so every
    target in the manifest is equally likely to be covered."""
    weights = np.array([max((b - a) // paths.CADENCE - 2 * span, 0)
                        for a, b in segments], dtype=float)
    if weights.sum() <= 0:
        raise ValueError(f"no segment holds a target at span {span}")
    a, b = segments[int(rng.choice(len(segments), p=weights / weights.sum()))]
    length = min(targets + 2 * span, (b - a) // paths.CADENCE)
    first = int(rng.integers(0, (b - a) // paths.CADENCE - length + 1))
    t0 = a + first * paths.CADENCE
    return t0, t0 + length * paths.CADENCE


def write_slice(directory, name, piece, satellites, blocks, span):
    """Write one slice: its day-shards, then ``slice.json``, then the rename.

    `blocks(satellite, times)` yields (t, block or None) in ascending order -- None
    for a slot that could not be read, which splits the slice's segment.  Returns
    the slice's path, or None if nothing in it is left to train on.
    """
    partial = os.path.join(directory, f".{name}.partial")
    final = os.path.join(directory, name)
    os.makedirs(partial)
    unusable = set()
    try:
        for (day, satellite), times in sorted(
                corpus.wanted_slots({"train": [piece]}, satellites).items()):
            with open(os.path.join(partial, dataset.shard_name(day, satellite)),
                      "wb") as handle:
                def records():
                    for t, block in blocks(satellite, times):
                        if block is None:
                            unusable.add(t)
                        yield t, block
                dataset.write_shard_stream(handle, day, satellite,
                                           paths.WAVELENGTHS, records())
        verified = segments_module.verify([piece], lambda t: t not in unusable,
                                          minimum_slots=2 * span + 1)
        if not verified:
            shutil.rmtree(partial)
            return None
        with open(os.path.join(partial, SLICE_FILE), "w") as handle:
            json.dump({"segments": [list(pair) for pair in verified],
                       "satellites": list(satellites), "made": time.time()}, handle)
        os.rename(partial, final)
    except BaseException:
        shutil.rmtree(partial, ignore_errors=True)
        raise
    return final


def ready_slices(directory):
    """Complete slices, oldest first: [(path, bytes)]."""
    found = []
    for name in os.listdir(directory) if os.path.isdir(directory) else ():
        path = os.path.join(directory, name)
        if name.startswith(".") or not os.path.exists(os.path.join(path, SLICE_FILE)):
            continue
        try:
            made = json.load(open(os.path.join(path, SLICE_FILE)))["made"]
            size = sum(entry.stat().st_size for entry in os.scandir(path))
        except (OSError, ValueError, KeyError):
            continue                       # retired between listing and reading
        found.append((made, path, size))
    return [(path, size) for _, path, size in sorted(found)]


def retire(directory, budget_bytes):
    """Delete the oldest slices until the pool fits `budget_bytes`, never the newest:
    an empty pool would starve training.  Returns the deleted paths."""
    slices = ready_slices(directory)
    total = sum(size for _, size in slices)
    gone = []
    for path, size in slices[:-1]:
        if total <= budget_bytes:
            break
        shutil.rmtree(path, ignore_errors=True)
        total -= size
        gone.append(path)
    return gone


def clear_partials(directory):
    """Remove half-built slices a stopped producer left behind."""
    for name in os.listdir(directory) if os.path.isdir(directory) else ():
        if name.startswith(".") and name.endswith(".partial"):
            shutil.rmtree(os.path.join(directory, name), ignore_errors=True)


def open_slice(path):
    """(shards {(day, satellite): Shard}, segments) of one ready slice.  Raises
    OSError if it was retired before every shard could be opened."""
    info = json.load(open(os.path.join(path, SLICE_FILE)))
    shards = {}
    try:
        for name in os.listdir(path):
            if name.endswith(".shard"):
                day, satellite = name[:-len(".shard")].rsplit("_g", 1)
                shards[(day, int(satellite))] = dataset.Shard(os.path.join(path, name))
    except BaseException:
        for shard in shards.values():
            shard.close()
        raise
    return shards, [tuple(pair) for pair in info["segments"]]


def slice_name(counter, piece):
    """Sortable, unique and self-describing: when made, and what it covers."""
    day = datetime.datetime.fromtimestamp(piece[0], datetime.timezone.utc)
    return f"{counter:07d}_{day:%Y%m%dT%H%M}_{(piece[1] - piece[0]) // paths.CADENCE}"