import json
import os

import numpy as np
import pytest

from conftest import solar_disc
from suvi import corpus, dataset, paths, segments as segments_lib

CAD = paths.CADENCE
BASE = 1715385600          # 2024-05-11 00:00 UTC


def make_segments(lengths, gap=400):
    """Contiguous runs of the given slot lengths, spaced apart."""
    out, t = [], BASE
    for length in lengths:
        out.append((t, t + length * CAD))
        t += (length + gap) * CAD
    return out


def test_slices_are_capped_so_no_segment_dominates():
    segments = make_segments([300, 2000, 350])
    chosen, used = corpus.choose_slices(segments, cap=512, budget_bytes=1e12,
                                        satellites=(16, 18), seed=0,
                                        val_count=0, test_count=0)
    lengths = sorted((b - a) // CAD for a, b in chosen["train"])
    assert lengths == [300, 350, 512]          # the 2000-slot run was truncated
    assert used == pytest.approx(sum(lengths) * corpus.SLOT_BYTES * 2)


def test_slices_stay_inside_their_segment():
    segments = make_segments([2000])
    for seed in range(8):
        chosen, _ = corpus.choose_slices(segments, cap=512, budget_bytes=1e12,
                                         satellites=(16,), seed=seed,
                                         val_count=0, test_count=0)
        (start, end), = chosen["train"]
        assert BASE <= start and end <= BASE + 2000 * CAD
        assert (end - start) // CAD == 512


def test_budget_is_respected_and_prefers_more_segments():
    segments = make_segments([400] * 20)
    budget = 10 * 400 * corpus.SLOT_BYTES * 2          # room for exactly ten
    chosen, used = corpus.choose_slices(segments, cap=512, budget_bytes=budget,
                                        satellites=(16, 18), seed=1,
                                        val_count=0, test_count=0)
    assert len(chosen["train"]) == 10
    assert used <= budget


def test_evaluation_sets_are_admitted_first_and_are_diverse():
    """val/test must not be squeezed out by a full disk, and must span many
    segments -- three long stretches would make validation measure three days."""
    segments = make_segments([400] * 30)
    budget = 12 * 400 * corpus.SLOT_BYTES * 2
    chosen, _ = corpus.choose_slices(segments, cap=512, budget_bytes=budget,
                                     satellites=(16, 18), seed=2,
                                     val_count=5, test_count=5)
    assert len(chosen["val"]) == 5 and len(chosen["test"]) == 5
    assert len(chosen["train"]) == 2
    everything = chosen["train"] + chosen["val"] + chosen["test"]
    assert len({s for s, _ in everything}) == len(everything)   # no reuse


def test_evaluation_segments_are_stable_across_budgets():
    """Repacking at a different budget or cap must not silently change what the
    model is being measured on."""
    segments = make_segments([400] * 30)
    big, _ = corpus.choose_slices(segments, cap=512, budget_bytes=1e12,
                                  satellites=(16,), seed=3, val_count=4, test_count=4)
    small, _ = corpus.choose_slices(segments, cap=512,
                                    budget_bytes=10 * 400 * corpus.SLOT_BYTES,
                                    satellites=(16,), seed=3, val_count=4, test_count=4)
    assert big["val"] == small["val"] and big["test"] == small["test"]


def test_wanted_slots_groups_every_chosen_slice_by_day_and_satellite():
    """The slices say which Sun to ship; the builder needs it keyed the way shards
    are, one ascending list per (day, satellite)."""
    chosen = {"train": [(BASE, BASE + 5 * CAD)],
              "val": [(BASE + 100 * CAD, BASE + 102 * CAD)],   # same day: 360/day
              "test": []}
    grouped = corpus.wanted_slots(chosen, (16, 18))
    assert set(grouped) == {("2024-05-11", 16), ("2024-05-11", 18)}
    for times in grouped.values():
        assert times == sorted(times)
        assert len(times) == 7                      # 5 train + 2 val on one day
        assert all((t - BASE) % CAD == 0 for t in times)


def test_wanted_slots_splits_a_slice_that_crosses_midnight():
    """Shards are day-keyed, so a slice spanning a day boundary must land in two."""
    start = BASE + 358 * CAD                        # 23:52 UTC
    grouped = corpus.wanted_slots({"train": [(start, start + 4 * CAD)]}, (16,))
    assert sorted(grouped) == [("2024-05-11", 16), ("2024-05-12", 16)]
    assert len(grouped[("2024-05-11", 16)]) == 2
    assert len(grouped[("2024-05-12", 16)]) == 2


def test_no_day_hosts_two_splits():
    """Shards are day-keyed, so a shared day would place val frames inside the
    training sampler's context reach -- mild leakage, easy to prevent here."""
    # Adjacent segments separated by only a few slots share calendar days.
    segments = [(BASE + i * 150 * CAD, BASE + (i * 150 + 130) * CAD)
                for i in range(24)]
    chosen, _ = corpus.choose_slices(segments, cap=512, budget_bytes=1e12,
                                     satellites=(16, 18), seed=4,
                                     val_count=6, test_count=6)
    owner = {}
    for role, slices in chosen.items():
        for pair in slices:
            for day in segments_lib.days_touched([pair]):
                assert owner.setdefault(day, role) == role, f"{day} shared"
    assert chosen["val"] and chosen["test"] and chosen["train"]


def test_default_cap_is_half_a_solar_rotation():
    """The scale at which a contiguous segment stops being one view of one
    hemisphere.  Capping below this discards new Sun, not redundancy."""
    slots = segments_lib.decorrelation_slots()
    days = slots * CAD / 86400
    assert 12.5 < days < 14.0, f"{days:.1f} d is not half a synodic rotation"
    # Nothing in the present corpus reaches it, so nothing gets truncated.
    segments = make_segments([2570])
    chosen, _ = corpus.choose_slices(segments, cap=slots, budget_bytes=1e12,
                                     satellites=(16, 18), seed=0,
                                     val_count=0, test_count=0)
    (a, b), = chosen["train"]
    assert (b - a) // CAD == 2570