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).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 test_charbonnier_is_zero_for_a_perfect_match():
    x = torch.rand(4, 6, 8, 8)
    assert train.charbonnier(x, x) == pytest.approx(train.CHARBONNIER_EPS, abs=1e-6)


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)
                              / train.charbonnier(spread, 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)
    assert train.gradient_loss(blurred, truth) > train.gradient_loss(truth, truth)


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],
    )
    assert got == pytest.approx(float(np.mean(expected)), 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))


# ------------------------------------------------------------------------- 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),
    }


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["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)


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"])
        b = net(padded["stack"][:1], padded["condition"][: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_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, target = train.aligned_batch(batch, device)
        prediction = net(aligned, condition)
        loss = train.criterion(prediction, target)
        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"