"""The rotating training pool: slices chosen, written, read, retired, and the
producer and trainer that share them.

What must hold: a slice lies inside one pristine segment with the model's span of
margin; a reader never sees a half-built slice and keeps reading one retired under
it; unreadable slots split a slice instead of entering it; the pool stays inside its
budget, oldest out first; and the manifest keeps evaluation and holdout data out.
"""

import datetime
import os

import numpy as np
import pytest

import pool_producer
from conftest import solar_disc, write_fits
from suvi import availability, dataset, db, detectors, paths, pool

BASE = 1715400000 // paths.CADENCE * paths.CADENCE
SPAN = 4
SATELLITES = (16, 18)


def block(value=1000):
    return np.full((6, dataset.SHARD_SIZE, dataset.SHARD_SIZE), value, dtype=np.uint16)


def all_good(satellite, times):
    for t in times:
        yield t, block()


def slots(count, start=0):
    return (BASE + start * paths.CADENCE, BASE + (start + count) * paths.CADENCE)


# --------------------------------------------------------------------- choosing


def test_a_slice_lies_inside_one_segment_with_its_margin():
    rng = np.random.default_rng(0)
    segments = [slots(40), slots(30, start=100)]
    for _ in range(200):
        t0, t1 = pool.choose_slice(segments, SPAN, rng, targets=10)
        assert any(a <= t0 and t1 <= b for a, b in segments)
        assert (t1 - t0) // paths.CADENCE == 10 + 2 * SPAN


def test_a_short_segment_gives_its_whole_length():
    t0, t1 = pool.choose_slice([slots(12)], SPAN, np.random.default_rng(1), targets=50)
    assert (t0, t1) == slots(12)


def test_slices_are_drawn_by_target_count_and_never_from_spanless_segments():
    rng = np.random.default_rng(2)
    short = slots(2 * SPAN, start=500)              # holds no target at this span
    segments = [slots(20), slots(200, start=50), short]
    picks = [pool.choose_slice(segments, SPAN, rng, targets=4) for _ in range(600)]
    inside = [sum(a <= t0 < b for t0, _ in picks) for a, b in segments]
    assert inside[2] == 0
    assert inside[1] > 5 * inside[0]                # 192 targets against 12


def test_choosing_from_nothing_trainable_is_an_error():
    with pytest.raises(ValueError):
        pool.choose_slice([slots(2 * SPAN)], SPAN, np.random.default_rng(0))


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


def test_a_written_slice_reads_back_whole(tmp_path):
    piece = slots(20)
    path = pool.write_slice(str(tmp_path), "0000001_x", piece, SATELLITES, all_good, SPAN)
    shards, segments = pool.open_slice(path)
    assert segments == [piece]
    assert {satellite for _, satellite in shards} == set(SATELLITES)
    for shard in shards.values():
        assert len(shard.times()) == 20
        np.testing.assert_array_equal(shard.raw(shard.times()[0]), block())
        shard.close()


def test_a_slice_across_midnight_has_a_shard_per_day(tmp_path):
    midnight = (BASE // 86400 + 1) * 86400
    piece = (midnight - 5 * paths.CADENCE, midnight + 15 * paths.CADENCE)
    path = pool.write_slice(str(tmp_path), "s", piece, SATELLITES, all_good, SPAN)
    shards, _ = pool.open_slice(path)
    assert len(shards) == 4
    assert sum(len(s.times()) for s in shards.values()) == 2 * 20
    for shard in shards.values():
        shard.close()


def test_an_unreadable_slot_splits_the_slice(tmp_path):
    piece = slots(30)
    hole = BASE + 15 * paths.CADENCE

    def holed(satellite, times):
        for t in times:
            yield t, None if (satellite == 18 and t == hole) else block()

    path = pool.write_slice(str(tmp_path), "s", piece, SATELLITES, holed, SPAN)
    _, segments = pool.open_slice(path)
    assert segments == [(BASE, hole), (hole + paths.CADENCE, piece[1])]


def test_a_slice_with_nothing_left_to_train_on_is_dropped(tmp_path):
    def nothing(satellite, times):
        for t in times:
            yield t, None

    assert pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, nothing, SPAN) is None
    assert os.listdir(tmp_path) == []


def test_a_failed_write_leaves_nothing_behind(tmp_path):
    def broken(satellite, times):
        yield times[0], block()
        raise OSError("share went away")

    with pytest.raises(OSError):
        pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, broken, SPAN)
    assert os.listdir(tmp_path) == []


def test_readers_see_only_complete_slices_oldest_first(tmp_path):
    first = pool.write_slice(str(tmp_path), "b", slots(20), SATELLITES, all_good, SPAN)
    second = pool.write_slice(str(tmp_path), "a", slots(20, 40), SATELLITES, all_good, SPAN)
    os.makedirs(tmp_path / ".c.partial")
    assert [path for path, _ in pool.ready_slices(str(tmp_path))] == [first, second]
    assert pool.ready_slices(str(tmp_path / "absent")) == []


def test_retiring_drops_the_oldest_until_the_pool_fits(tmp_path):
    made = [pool.write_slice(str(tmp_path), f"{i}", slots(20, 40 * i), SATELLITES,
                             all_good, SPAN) for i in range(4)]
    size = pool.ready_slices(str(tmp_path))[0][1]
    gone = pool.retire(str(tmp_path), budget_bytes=2.5 * size)
    assert gone == made[:2]
    assert [p for p, _ in pool.ready_slices(str(tmp_path))] == made[2:]
    assert pool.retire(str(tmp_path), budget_bytes=10 * size) == []


def test_partials_from_a_stopped_producer_are_cleared(tmp_path):
    os.makedirs(tmp_path / ".x.partial")
    kept = pool.write_slice(str(tmp_path), "k", slots(20), SATELLITES, all_good, SPAN)
    pool.clear_partials(str(tmp_path))
    assert os.listdir(tmp_path) == [os.path.basename(kept)]


def test_the_newest_slice_is_never_retired(tmp_path):
    only = pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, all_good, SPAN)
    assert pool.retire(str(tmp_path), budget_bytes=0) == []
    assert [p for p, _ in pool.ready_slices(str(tmp_path))] == [only]


def test_a_slice_retired_mid_visit_stays_readable(tmp_path):
    path = pool.write_slice(str(tmp_path), "s", slots(20), SATELLITES, all_good, SPAN)
    pool.write_slice(str(tmp_path), "t", slots(20, 40), SATELLITES, all_good, SPAN)
    shards, _ = pool.open_slice(path)
    pool.retire(str(tmp_path), budget_bytes=0)
    assert not os.path.exists(path)
    for shard in shards.values():
        np.testing.assert_array_equal(shard.raw(shard.times()[-1]), block())
        shard.close()
    with pytest.raises(OSError):
        pool.open_slice(path)


def test_slice_names_sort_by_counter():
    names = [pool.slice_name(i, slots(20)) for i in (9, 10, 100)]
    assert names == sorted(names)


# -------------------------------------------------------------------- manifest


def test_a_manifest_round_trips_and_a_foreign_version_is_refused(tmp_path):
    manifest = {"version": pool.MANIFEST_VERSION, "segments": [[1, 2]], "frames": {}}
    pool.save_manifest(manifest, str(tmp_path / "m.zst"))
    assert pool.load_manifest(str(tmp_path / "m.zst")) == manifest
    pool.save_manifest(dict(manifest, version=0), str(tmp_path / "old.zst"))
    with pytest.raises(ValueError, match="pool_manifest.py"):
        pool.load_manifest(str(tmp_path / "old.zst"))


def index(tmp_path, count, bad=(), missing=()):
    """`count` slots per satellite from BASE, every band judged good unless named."""
    conn = db.connect(str(tmp_path / "index.sqlite"))
    run = db.create_detector_run(conn, detectors.JUDGMENT_RUN, detectors.judgment_config())
    rows = []
    for satellite in SATELLITES:
        for slot in range(count):
            for wavelength in paths.WAVELENGTHS:
                key = (satellite, slot, wavelength)
                if key in missing:
                    continue
                when = BASE + slot * paths.CADENCE
                name = paths.FrameName(satellite=satellite, wavelength=wavelength,
                                       t_start=when, t_end=when + paths.CADENCE,
                                       version="1-0-2")
                frame_id = db.upsert_frame(conn, name, name.relpath(), size_bytes=1_500_000)
                rows.append((frame_id, "bad" if key in bad else "good", None, None, 0))
    db.record_detections(conn, run, rows)
    conn.commit()
    return conn, run


def pristine(count, cuts=()):
    """Every band of both satellites GOOD, except one band BAD at each slot in `cuts`
    -- which splits the era into segments there."""
    states = np.full((2, len(availability.BANDS), count), availability.GOOD, np.uint8)
    for cut in cuts:
        states[0, 0, cut] = availability.BAD
    return availability.Availability(states=states, t0=BASE, cadence=paths.CADENCE,
                                     satellites=SATELLITES)


def manifest_for(conn, run, model, old_val=(), test=(), holdout=(), **kwargs):
    kwargs.setdefault("val_new", 0)
    return pool.build_manifest(conn, model, run, SPAN, list(old_val), list(test),
                               list(holdout), SATELLITES, **kwargs)


def test_the_manifest_lists_every_frame_of_its_segments(tmp_path):
    conn, run = index(tmp_path, 30)
    manifest = manifest_for(conn, run, pristine(30))
    assert manifest["segments"] == [list(slots(30))]
    assert manifest["val"] == [] and manifest["val_segments"] == []
    assert len(manifest["frames"]["16"]) == len(manifest["frames"]["18"]) == 30
    first = manifest["frames"]["18"][str(BASE)]
    assert len(first) == 6 and all("g18" in path for path in first)


def test_a_slot_without_a_good_frame_of_every_band_splits_its_segment(tmp_path):
    conn, run = index(tmp_path, 30, bad={(16, 12, 171)}, missing={(18, 20, 304)})
    manifest = manifest_for(conn, run, pristine(30))
    # 13..20 is 7 slots, below the 2*span+1 a target needs, so it goes too.
    assert manifest["segments"] == [list(slots(12)),
                                    [BASE + 21 * paths.CADENCE, BASE + 30 * paths.CADENCE]]
    assert str(BASE + 12 * paths.CADENCE) not in manifest["frames"]["16"]


def test_test_segments_and_holdout_days_stay_out_of_the_pool(tmp_path):
    conn, run = index(tmp_path, 30)
    model = pristine(30)
    near = (BASE + 200 * paths.CADENCE, BASE + 300 * paths.CADENCE)   # 12 h later
    assert manifest_for(conn, run, model, test=[near])["segments"] == []
    far = (BASE + 10 * 86400, BASE + 10 * 86400 + 100 * paths.CADENCE)
    assert manifest_for(conn, run, model, test=[far])["segments"] == [list(slots(30))]
    day = datetime.datetime.fromtimestamp(BASE, datetime.timezone.utc).strftime("%Y-%m-%d")
    assert manifest_for(conn, run, model, holdout=[day])["segments"] == []


#: Four segments: slots 0..49, 50..99 and 100..149 (49 each), and 150..200 (50).
CUTS = (49, 99, 149)


def segment(index):
    end = 50 * index + 49 if index < 3 else 200
    return [BASE + 50 * index * paths.CADENCE, BASE + end * paths.CADENCE]


def test_validation_spreads_over_the_era_and_leaves_the_pool(tmp_path, monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=2, val_targets=10)
    assert manifest["val_segments"] == [segment(1), segment(3)]
    assert segment(0) in manifest["segments"] and segment(2) in manifest["segments"]
    for a, b in manifest["segments"]:            # the rest of 1 and 3 stays trainable
        assert not any(a < high and b > low for low, high in manifest["val"])
    assert len(manifest["val"]) == 2
    for (a, b), (low, high) in zip(manifest["val"], manifest["val_segments"]):
        assert low <= a and b <= high and (b - a) // paths.CADENCE == 10 + 2 * SPAN
    for a, b in manifest["val"]:
        assert all(str(t) in manifest["frames"]["16"] for t in range(a, b, paths.CADENCE))


def test_the_old_validation_segments_stay_in_validation(tmp_path, monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), old_val=[segment(0)],
                            val_new=1, val_targets=10)
    assert manifest["val_segments"] == [segment(0), segment(2)]
    assert segment(1) in manifest["segments"] and segment(3) in manifest["segments"]


def test_new_validation_avoids_the_previous_models_training_data(tmp_path, monkeypatch):
    """So the previous model can be scored on the new set without having seen it."""
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=4, val_targets=10,
                            old_train=[segment(2)])
    assert manifest["val_segments"] == [segment(0), segment(1), segment(3)]
    assert segment(2) in manifest["segments"]


def test_subtract_cuts_out_the_guarded_spans_and_drops_crumbs(monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 2 * paths.CADENCE)
    a, b = slots(100)
    cut = (a + 40 * paths.CADENCE, a + 50 * paths.CADENCE)
    assert pool.subtract([(a, b)], [cut], minimum=9) == [
        (a, a + 38 * paths.CADENCE), (a + 52 * paths.CADENCE, b)]
    near_end = (b - 5 * paths.CADENCE, b)
    assert pool.subtract([(a, b)], [near_end], minimum=9) == [(a, b - 7 * paths.CADENCE)]
    assert pool.subtract([(a, a + 10 * paths.CADENCE)], [(a + 4 * paths.CADENCE,
                                                         a + 5 * paths.CADENCE)], 9) == []


def test_a_validation_window_never_comes_from_a_test_segment(tmp_path, monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), test=[segment(1)],
                            val_new=3, val_targets=10)
    assert segment(1) not in manifest["val_segments"]
    assert segment(1) not in manifest["segments"]


def test_test_segments_are_kept_whole_with_their_frames(tmp_path, monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), test=[segment(1)])
    assert manifest["test"] == [segment(1)]
    a, b = segment(1)
    assert all(str(t) in manifest["frames"]["18"] for t in range(a, b, paths.CADENCE))


def test_a_short_validation_segment_gives_its_whole_length(tmp_path, monkeypatch):
    monkeypatch.setattr(pool, "GUARD_SECONDS", 0)
    conn, run = index(tmp_path, 200)
    old = [BASE + 3 * paths.CADENCE, BASE + 15 * paths.CADENCE]       # 12 slots
    manifest = manifest_for(conn, run, pristine(200, CUTS), old_val=[old],
                            val_targets=40)
    assert manifest["val"] == [old]


def test_the_pool_keeps_a_day_clear_of_validation(tmp_path):
    conn, run = index(tmp_path, 200)
    manifest = manifest_for(conn, run, pristine(200, CUTS), val_new=1, val_targets=10)
    assert manifest["val_segments"] and manifest["segments"] == []


# -------------------------------------------------------------------- producer


def archive(tmp_path, count):
    """A FITS archive of `count` slots per satellite plus its pool manifest."""
    root = tmp_path / "Data"
    frames = {str(s): {} for s in SATELLITES}
    image = solar_disc(size=dataset.SHARD_SIZE, radius=200)
    for satellite in SATELLITES:
        for slot in range(count):
            when = BASE + slot * paths.CADENCE
            files = []
            for wavelength in paths.WAVELENGTHS:
                rel = f"g{satellite}/{wavelength}/{when}.fits"
                os.makedirs(root / os.path.dirname(rel), exist_ok=True)
                write_fits(str(root / rel), image)
                files.append(rel)
            frames[str(satellite)][str(when)] = files
    manifest = {"version": pool.MANIFEST_VERSION, "satellites": list(SATELLITES),
                "span": SPAN, "judgment_run": 1, "segments": [list(slots(count))],
                "val": [list(slots(count))], "val_segments": [list(slots(count))],
                "test": [list(slots(count // 2 + SPAN))], "frames": frames}
    pool.save_manifest(manifest, str(tmp_path / "manifest.zst"))
    return root


def test_the_producer_fills_the_pool_and_keeps_it_in_budget(tmp_path):
    root = archive(tmp_path, 12)
    directory = tmp_path / "pool"
    pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
                        str(root), "--pool", str(directory), "--workers", "2",
                        "--slices", "3", "--seed", "0", "--budget-gb", "0"])
    resident = pool.ready_slices(str(directory))
    assert len(resident) == 1                       # a zero budget keeps only the newest
    assert os.path.basename(resident[0][0]).startswith("0000003")
    shards, segments = pool.open_slice(resident[0][0])
    assert segments == [slots(12)]
    frame = next(iter(shards.values())).frames(BASE)
    assert frame.shape == (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE)
    for shard in shards.values():
        shard.close()


def test_an_unreadable_archive_file_becomes_a_hole_not_a_crash(tmp_path):
    root = archive(tmp_path, 12)
    os.remove(root / f"g18/171/{BASE + 6 * paths.CADENCE}.fits")
    directory = tmp_path / "pool"
    pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
                        str(root), "--pool", str(directory), "--workers", "1",
                        "--slices", "1", "--seed", "0"])
    assert pool.ready_slices(str(directory)) == []  # 6 + 5 slots: below 2*span+1 each


def test_fixed_mode_writes_each_validation_window_once(tmp_path):
    root = archive(tmp_path, 12)
    directory = tmp_path / "val"
    argv = ["--manifest", str(tmp_path / "manifest.zst"), "--archive", str(root),
            "--pool", str(directory), "--workers", "1", "--fixed", "val"]
    pool_producer.main(argv)
    assert sorted(os.listdir(directory)) == ["val000"]
    made = os.path.getmtime(directory / "val000" / pool.SLICE_FILE)
    pool_producer.main(argv)                        # a rerun skips what is complete
    assert os.path.getmtime(directory / "val000" / pool.SLICE_FILE) == made
    test_dir = tmp_path / "test"
    pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
                        str(root), "--pool", str(test_dir), "--workers", "1",
                        "--fixed", "test"])
    _, segments = pool.open_slice(str(test_dir / "test000"))
    assert segments == [slots(12 // 2 + SPAN)]


def test_encode_slot_declines_a_band_it_cannot_read(tmp_path):
    assert pool_producer.encode_slot([str(tmp_path / "missing.fits")] * 6) is None


# ------------------------------------------------------------------ the trainer


torch = pytest.importorskip("torch")


def test_the_pool_set_yields_model_ready_samples(tmp_path):
    import train

    root = archive(tmp_path, 14)
    directory = tmp_path / "pool"
    pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
                        str(root), "--pool", str(directory), "--workers", "1",
                        "--slices", "1", "--seed", "0"])
    data = train.PoolSet(str(directory), SATELLITES, offsets=(1, 4), per_visit=3)
    stream = iter(data)
    items = [next(stream) for _ in range(5)]        # crosses a visit boundary
    for item in items:
        assert item["stack"].shape[1:] == (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE)
        assert item["lost"].sum() >= 1
    batch = train.collate(items[:2])
    assert batch["stack"].shape[0] == 2


def test_the_slice_set_is_a_fixed_validation_set(tmp_path):
    import train

    root = archive(tmp_path, 14)
    directory = tmp_path / "val"
    pool_producer.main(["--manifest", str(tmp_path / "manifest.zst"), "--archive",
                        str(root), "--pool", str(directory), "--workers", "1",
                        "--fixed", "val"])
    data = train.SliceSet(str(directory), SATELLITES, offsets=(1, 4))
    assert len(data) == 2 * (14 - 2 * SPAN)
    first, again = data[3], data[3]
    assert torch.equal(first["stack"], again["stack"])      # damage pinned to the index
    assert torch.equal(first["lost"], again["lost"])


def test_the_slice_set_refuses_an_empty_directory(tmp_path):
    import train

    with pytest.raises(SystemExit):
        train.SliceSet(str(tmp_path), SATELLITES)


def test_the_pool_set_gives_up_when_no_producer_appears(tmp_path, monkeypatch):
    import train

    monkeypatch.setattr(train.time, "sleep", lambda seconds: None)
    data = train.PoolSet(str(tmp_path / "empty"), SATELLITES, wait=10.0)
    with pytest.raises(RuntimeError, match="producer"):
        next(iter(data))


def test_gpu_train_runs_the_producer_beside_a_pool_run():
    import gpu_train

    args = gpu_train.argparse.Namespace(
        shards_dir="shards", name="v10", epochs=6, batch=4, lr=2e-4, base=32,
        depth=3, workers=6, val_samples=96, eval_every=500, steps=8923,
        lr_range=None, span=64, warmup=None, relative=False, corr_flow=False,
        attention=True, max_steps=57000, overfit=None, resume=False, pool=True,
        pool_manifest="pool_manifest.json.zst", pool_budget_gb=48.0, pool_workers=16,
        val_dir="val10", evaluate=None)
    assert f"--pool {gpu_train.POOL_DIR}" in gpu_train.command(args)
    assert "--val-slices /work/val10" in gpu_train.command(args)
    producer = gpu_train.producer_command(args)
    assert "--archive /archive/Data" in producer and gpu_train.POOL_DIR in producer
    args.pool = False
    assert "--pool" not in gpu_train.command(args)