import json

import numpy as np
import pytest

torch = pytest.importorskip("torch")

import train  # noqa: E402
from suvi import dataset, metrics, model, paths, samples  # noqa: E402


# ------------------------------------------------------------------ display mapping


def test_hard_display_map_matches_the_metric_the_bench_scores():
    """With eps=0 the torch map must *be* metrics.to_display, or the reported PSNR
    stops being comparable to the published filler table."""
    rng = np.random.default_rng(0)
    radiance = rng.uniform(-1.0, 60.0, (1, 6, 16, 16)).astype(np.float32)
    coded = torch.from_numpy(samples.encode_for_model(radiance))

    got = train.display_map(coded, eps=0.0).numpy()
    for index, wavelength in enumerate(paths.WAVELENGTHS):
        expected = metrics.to_display(radiance[0, index], wavelength)
        np.testing.assert_allclose(got[0, index], expected, atol=2e-3)


def test_display_map_still_sends_zero_to_zero_when_softened():
    """The softening must not lift the black level: (0+eps)^g - eps^g == 0."""
    coded = samples.encode_for_model(np.zeros((1, 6, 4, 4), np.float32))
    out = train.display_map(torch.from_numpy(coded))
    assert float(out.abs().max()) < 1e-6      # fp32 rounding only; 1/255 is 4e-3


def test_display_map_saturates_without_nan():
    for value in (-1.0, 1.0):
        coded = torch.full((1, 6, 4, 4), value, requires_grad=True)
        out = train.display_map(coded)
        out.sum().backward()
        assert torch.isfinite(out).all() and torch.isfinite(coded.grad).all()


# --------------------------------------------------------------- loss conditioning


def test_the_loss_gradient_is_bounded_everywhere():
    """The property whose absence sank three training runs.

    The old display-space loss had |dL/dpred| of exactly 0 across most of the frame
    (below the vmin clamp) and a 25x spike with a hard discontinuity at the
    threshold; clip_grad_norm then renormalised every update onto that annulus.
    Here the whole coded range, including the sub-vmin corona and the threshold
    itself, must see a per-pixel gradient that is neither zero over wide regions nor
    orders of magnitude apart.
    """
    # Sweep the entire coded range, so the threshold annulus of every band is in,
    # with a *constant* residual: the per-pixel gradient then measures the loss
    # landscape's Jacobian, not the accident of which pixels happened to be right.
    target = torch.linspace(-1.0, 1.0, 64 * 64).reshape(1, 1, 64, 64).repeat(1, 6, 1, 1)
    prediction = (target + 0.05).requires_grad_(True)

    train.criterion(prediction, target, torch.ones(1, 6)).backward()
    per_pixel = prediction.grad.abs() * prediction.numel()
    assert torch.isfinite(prediction.grad).all()
    assert float(per_pixel.max()) < 100.0, "a gradient spike survived the redesign"
    # The asinh term supervises everything: no dead zones anywhere.
    assert float(per_pixel.min()) > 0.1, "part of the frame is unsupervised"


def every_band(x):
    return torch.ones(x.shape[0], x.shape[1])


def test_charbonnier_is_zero_for_a_perfect_match():
    x = torch.rand(4, 6, 8, 8)
    assert train.charbonnier(x, x, every_band(x)) == pytest.approx(train.CHARBONNIER_EPS,
                                                                   abs=1e-6)


def test_the_loss_counts_only_the_lost_bands():
    """The target's other bands are inputs; an error there is not the model's."""
    truth = torch.zeros(2, 6, 8, 8)
    prediction = truth.clone()
    prediction[:, 2] = 1.0                       # wrong only in 171A
    lost = torch.zeros(2, 6)
    lost[:, 4] = 1.0                             # the sample lost 284A
    assert float(train.criterion(prediction, truth, lost)) < 0.01
    lost[0, 2] = 1.0
    assert float(train.criterion(prediction, truth, lost)) > 0.1


def test_charbonnier_is_less_outlier_driven_than_squared_error():
    """A saturated pixel must not be able to dominate the whole frame's loss."""
    truth = torch.zeros(1, 1, 8, 8)
    spike = truth.clone()
    spike[0, 0, 0, 0] = 50.0
    spread = truth + 0.1

    charbonnier_ratio = float(train.charbonnier(spike, truth, every_band(truth))
                              / train.charbonnier(spread, truth, every_band(truth)))
    squared_ratio = float(((spike - truth) ** 2).mean() / ((spread - truth) ** 2).mean())
    assert charbonnier_ratio < 10
    assert squared_ratio > 1000


def test_gradient_loss_punishes_blur():
    """The term that makes a sharp answer beat a hedged smooth one."""
    truth = torch.zeros(1, 1, 32, 32)
    truth[:, :, 8:24, 8:24] = 1.0
    blurred = torch.nn.functional.avg_pool2d(truth, 5, stride=1, padding=2)
    bands = every_band(truth)
    assert train.gradient_loss(blurred, truth, bands) > train.gradient_loss(truth, truth,
                                                                            bands)


def test_display_psnr_matches_the_bench_definition():
    """Per band, then averaged -- the pooled-MSE PSNR reads ~4 dB lower (Jensen) and
    would not be comparable to the published baselines."""
    rng = np.random.default_rng(1)
    radiance = rng.uniform(0.0, 30.0, (6, 16, 16)).astype(np.float32)
    noisy = radiance + rng.normal(0, 0.5, radiance.shape).astype(np.float32)

    expected = []
    for index, wavelength in enumerate(paths.WAVELENGTHS):
        a = metrics.to_display(radiance[index], wavelength)
        b = metrics.to_display(noisy[index], wavelength)
        expected.append(10.0 * np.log10(1.0 / max(((a - b) ** 2).mean(), 1e-12)))

    got = train.display_psnr(
        torch.from_numpy(samples.encode_for_model(noisy))[None],
        torch.from_numpy(samples.encode_for_model(radiance))[None],
        torch.ones(1, 6),
    )
    assert got == pytest.approx(float(np.mean(expected)), abs=0.1)
    lost = torch.zeros(1, 6)
    lost[0, 1] = 1.0
    only = train.display_psnr(torch.from_numpy(samples.encode_for_model(noisy))[None],
                              torch.from_numpy(samples.encode_for_model(radiance))[None],
                              lost)
    assert only == pytest.approx(float(expected[1]), abs=0.1)


def test_display_psnr_of_identical_inputs_is_finite():
    x = torch.rand(1, 6, 8, 8) * 0.5
    assert np.isfinite(train.display_psnr(x, x, torch.ones(1, 6)))


# ------------------------------------------------------------------------- collate


def item(depth, bands=6, size=16):
    return {
        "stack": torch.randn(depth, bands, size, size),
        "condition": torch.stack(
            [model.frame_conditioning("available", True, 240.0) for _ in range(depth)]
        ),
        "dts": torch.full((depth,), 240.0),
        "gains": torch.ones(depth, bands),
        "offsets": torch.zeros(depth, bands),
        "b0": torch.tensor(0.05),
        "radius": torch.tensor(0.3),
        "target": torch.randn(bands, size, size),
        "present": torch.ones(depth, bands),
        "lost": torch.ones(bands),
    }


def test_collate_pads_ragged_stacks():
    """Stack depth varies with how many anchors were found near a shard's edge."""
    batch = train.collate([item(5), item(7), item(6)])
    assert batch["stack"].shape[:2] == (3, 7)
    assert batch["condition"].shape == (3, 7, model.COND_DIM)
    assert batch["gains"].shape == (3, 7, 6)
    assert batch["target"].shape == (3, 6, 16, 16)
    assert batch["present"].shape == (3, 7, 6) and batch["lost"].shape == (3, 6)
    assert batch["b0"].shape == (3,)


def test_collate_marks_padding_as_missing_with_identity_transfer():
    """Padding must take zero weight and pass alignment untouched, or the batch
    shape would change the answer."""
    batch = train.collate([item(3), item(6)])
    assert torch.all(batch["condition"][0, 3:, 1] == 1.0)   # 'missing' one-hot
    assert torch.all(batch["stack"][0, 3:] == 0.0)
    assert torch.all(batch["gains"][0, 3:] == 1.0)
    assert torch.all(batch["offsets"][0, 3:] == 0.0)
    assert torch.all(batch["present"][0, 3:] == 0.0)


def test_padding_does_not_change_the_prediction():
    """The property that makes ragged batching safe at all."""
    torch.manual_seed(0)
    net = model.build(base=8, depth=2).eval()
    short = item(3, size=32)
    alone = train.collate([short])
    padded = train.collate([short, item(6, size=32)])
    with torch.no_grad():
        a = net(alone["stack"], alone["condition"], alone["present"])
        b = net(padded["stack"][:1], padded["condition"][:1], padded["present"][:1])
    assert torch.allclose(a, b, atol=1e-5)


# ---------------------------------------------------------------------- data feed


class MemoryShard:
    def __init__(self, satellite, times, size=64):
        from conftest import solar_disc

        self.satellite = satellite
        self._frames = {
            t: np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + b * 0.2)
                         for b in range(6)]).astype(np.float32)
            for t in times
        }

    def times(self):
        return sorted(self._frames)

    def frames(self, time):
        got = self._frames.get(time)
        return None if got is None else got.copy()


def shard_set(monkeypatch, tmp_path, days, deterministic=False, seed=0):
    base = 1715400000 // paths.CADENCE * paths.CADENCE
    times = [base + i * paths.CADENCE for i in range(40)]
    monkeypatch.setattr(dataset, "Shard", lambda path: MemoryShard(16, times))
    real_exists = train.os.path.exists
    # Pretend only the shard files exist; a blanket True breaks cv2's lazy loader.
    monkeypatch.setattr(
        train.os.path, "exists",
        lambda p: True if str(p).startswith(str(tmp_path)) else real_exists(p),
    )
    return train.ShardSet(str(tmp_path), days, satellites=(16, 18),
                          deterministic=deterministic, seed=seed)


def test_shard_set_yields_model_ready_tensors(monkeypatch, tmp_path):
    data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
    sample = data[0]
    assert sample["stack"].shape[1:] == (6, 64, 64)
    assert sample["condition"].shape[1] == model.COND_DIM
    assert sample["gains"].shape == (sample["stack"].shape[0], 6)
    assert torch.isfinite(sample["stack"]).all()
    assert torch.isfinite(sample["target"]).all()


def test_validation_damage_is_reproducible(monkeypatch, tmp_path):
    """A val number that moves because the damage moved measures nothing."""
    first = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1)
    second = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1)
    for index in (0, 5, 11):
        assert torch.allclose(first[index]["stack"], second[index]["stack"])


def test_training_damage_varies_between_epochs(monkeypatch, tmp_path):
    """The reason shards store clean frames: unlimited corruption variety."""
    data = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=False)
    assert not torch.allclose(data[3]["stack"], data[3]["stack"])


def test_learning_rate_is_warmup_times_plateau_scale():
    """The stable phase holds indefinitely at the plateau scale: LR only moves
    when the plateau logic halves the scale, never on a pre-committed clock."""
    base = 1e-4
    assert train.learning_rate(0, base) == pytest.approx(base / train.WARMUP_STEPS)
    assert train.learning_rate(train.WARMUP_STEPS, base) == pytest.approx(base)
    assert train.learning_rate(500_000, base) == pytest.approx(base)
    assert train.learning_rate(500_000, base, scale=0.25) == pytest.approx(base / 4)
    # Six halvings from 1.0 reach the floor: the run cannot freeze in fewer
    # than six consecutive full plateau windows.
    scale, halvings = 1.0, 0
    while scale > train.MIN_LR_SCALE:
        scale *= train.DECAY_FACTOR
        halvings += 1
    assert halvings == 6


def test_sweep_rate_ramps_exponentially_between_its_endpoints():
    assert train.sweep_rate(0, 1e-6, 3e-3, 600) == pytest.approx(1e-6)
    assert train.sweep_rate(599, 1e-6, 3e-3, 600) == pytest.approx(3e-3)
    quarter = train.sweep_rate(599 / 2, 1e-6, 3e-3, 600)
    assert quarter == pytest.approx((1e-6 * 3e-3) ** 0.5, rel=1e-6)


def test_getitem_survives_a_long_unreadable_stretch():
    """A vetted day can hold hours of contiguous sentinel slots; the fallback walk
    must be iterative -- 1,200 consecutive absent targets once blew the stack at
    step 7,600 of a training run."""
    from test_samples import HoleyShard

    base = 1715400000 // paths.CADENCE * paths.CADENCE
    times = [base + i * paths.CADENCE for i in range(1300)]
    shard = HoleyShard(16, times, 24, unreadable=times[:1200])
    made = train.ShardSet.__new__(train.ShardSet)
    made.deterministic = False
    made.seed = 0
    made.sampler = samples.Sampler({("d", 16): shard}, satellites=(16,),
                                   loss_probability=0.0)
    made.items = made.sampler.targets()
    got = made[0]                       # begins at the very start of the dead run
    assert got["target"].shape[0] == 6


def test_shard_set_refuses_to_run_on_nothing(tmp_path):
    with pytest.raises(SystemExit, match="no shards"):
        train.ShardSet(str(tmp_path), ["2024-01-01"])


# ------------------------------------------------------------------ training steps


def loader_for(monkeypatch, tmp_path):
    data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
    return torch.utils.data.DataLoader(data, batch_size=2, num_workers=0,
                                       collate_fn=train.collate)


def test_a_training_step_is_finite_end_to_end(monkeypatch, tmp_path):
    """Data feed, alignment, forward, loss, backward: all wired together."""
    loader = loader_for(monkeypatch, tmp_path)
    net = model.build(base=8, depth=2)
    optimiser = torch.optim.AdamW(net.parameters(), lr=1e-4)
    device = torch.device("cpu")

    for index, batch in enumerate(loader):
        if index >= 2:
            break
        aligned, condition, present, target, lost = train.aligned_batch(batch, device)
        prediction = net(aligned, condition, present)
        loss = train.criterion(prediction, target, lost)
        assert torch.isfinite(loss)
        optimiser.zero_grad()
        loss.backward()
        grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0))
        assert np.isfinite(grad)
        optimiser.step()


def test_evaluate_leaves_weights_alone(monkeypatch, tmp_path):
    """Validation must not train, or the held-out split stops being held out."""
    loader = loader_for(monkeypatch, tmp_path)
    net = model.build(base=8, depth=2)
    before = [p.detach().clone() for p in net.parameters()]
    stats = train.evaluate(net, loader, torch.device("cpu"))
    assert np.isfinite(stats["loss"]) and np.isfinite(stats["psnr"])
    for old, new in zip(before, net.parameters()):
        assert torch.equal(old, new)


def test_validation_subset_spans_the_whole_set(monkeypatch, tmp_path):
    """A shortened val pass must not become 'the first N targets': `targets()` is
    sorted by (satellite, time), so a prefix is the earliest days of one satellite."""
    data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
    wanted = 8
    stride = len(data) / wanted
    subset = torch.utils.data.Subset(data, [int(i * stride) for i in range(wanted)])
    indices = subset.indices
    assert len(set(indices)) == wanted
    assert max(indices) > len(data) * 0.8, "subset does not reach the end of the set"
    assert min(indices) < len(data) * 0.2, "subset does not start near the beginning"
    satellites = {data.items[i][0] for i in indices}
    assert len(satellites) > 1, "subset covers only one satellite"


# --------------------------------------------------------------- manifest shape


def test_training_refuses_a_manifest_without_segments(tmp_path):
    """Day-split manifests predate the pristine-segment corpus and no longer
    describe what the sampler needs.  Silently training on one would use a
    corpus whose splits were never verified segment-wise."""
    manifest = tmp_path / "manifest.json"
    manifest.write_text(json.dumps({
        "satellites": [16, 18], "digest": "old",
        "splits": {"train": ["2024-01-01"], "val": [], "test": []},
    }))
    with pytest.raises(SystemExit, match="segments"):
        train.main(["--shards", str(tmp_path), "--manifest", str(manifest),
                    "--out", str(tmp_path / "run")])


# ---------------------------------------------------------------- availability


def test_training_refuses_to_run_without_the_availability_model(tmp_path):
    """Silently reverting to synthetic-only outages would undo the whole point of
    profiling the archive."""
    with pytest.raises(SystemExit, match="availability"):
        train.load_availability(str(tmp_path))


def test_availability_model_round_trips(tmp_path):
    from suvi import availability

    states = np.zeros((2, 6, 4), np.uint8)
    states[0, 1] = [0, 1, 2, 0]
    states[1, 4] = [2, 2, 0, 1]
    model_in = availability.Availability(states=states, t0=1000, cadence=240,
                                         satellites=(16, 18))
    model_in.save(str(tmp_path / "availability.npz"))
    model_out = train.load_availability(str(tmp_path))
    np.testing.assert_array_equal(model_out.states, states)
    assert model_out.satellites == (16, 18)
    assert model_out.t0 == 1000 and model_out.cadence == 240