import os

import numpy as np
import pytest
import torch

from cellquant.model import CellUNet, N_OUT, load_checkpoint, pad_to_multiple, save_checkpoint

WEIGHTS = os.environ.get("CELLQUANT_WEIGHTS", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "weights", "cellquant.pt"))


def test_pad_to_multiple():
    x = torch.zeros(1, 3, 70, 100)
    xp, (ph, pw) = pad_to_multiple(x, 32)
    assert xp.shape[-2:] == (96, 128) and (ph, pw) == (26, 28)
    x2, (ph2, pw2) = pad_to_multiple(torch.zeros(1, 3, 64, 64), 32)
    assert x2.shape[-2:] == (64, 64) and (ph2, pw2) == (0, 0)


def test_model_forward_and_checkpoint_roundtrip(tmp_path):
    m = CellUNet("resnet34", pretrained=False).eval()
    x = torch.rand(1, 3, 64, 96)
    with torch.no_grad():
        y = m(x)
    assert y.shape == (1, N_OUT, 64, 96)
    p = tmp_path / "m.pt"
    save_checkpoint(m, str(p), {"note": "test"})
    m2, meta = load_checkpoint(str(p))
    assert meta["note"] == "test" and meta["encoder"] == "resnet34"
    with torch.no_grad():
        assert torch.allclose(m2(x), y, atol=1e-5)


@pytest.mark.skipif(not os.path.exists(WEIGHTS), reason="trained weights not present")
def test_predictor_deterministic_and_tiling_consistent():
    from cellquant.infer import Predictor
    rng = np.random.default_rng(0)
    img = (rng.random((300, 420, 3)) * 60).astype(np.uint8)
    p = Predictor(WEIGHTS, device="cpu", tta=1)
    a = p.probs(img)
    b = p.probs(img)
    assert all(np.array_equal(a[k], b[k]) for k in a)
    assert a["cell"].shape == (3, 300, 420)
    # force tiling and compare with the whole-image pass
    pt = Predictor(WEIGHTS, device="cpu", tta=1, tile=256, overlap=64, max_whole=128)
    c = pt.probs(img)
    diff = np.abs(c["cell"] - a["cell"]).mean()
    assert diff < 0.05, diff


@pytest.mark.skipif(not os.path.exists(WEIGHTS), reason="trained weights not present")
def test_tta_outputs_valid_probabilities():
    from cellquant.infer import Predictor
    img = (np.random.default_rng(1).random((128, 160, 3)) * 60).astype(np.uint8)
    for tta in (4, 8):
        pr = Predictor(WEIGHTS, device="cpu", tta=tta).probs(img)
        assert np.allclose(pr["cell"].sum(0), 1, atol=1e-4) and np.allclose(pr["nos"].sum(0), 1, atol=1e-4)


def test_multiscale_tta_keeps_shapes_and_matches_single_scale():
    from cellquant.model import CellUNet
    from cellquant.infer import Predictor
    import numpy as np
    img = (np.random.default_rng(3).random((70, 90, 3)) * 120).astype(np.uint8)
    m = CellUNet("resnet34", pretrained=False).eval()
    base = Predictor(m, device="cpu", tta=1).probs(img)
    same = Predictor(m, device="cpu", tta=1, scales=(1.0,)).probs(img)
    assert all(np.array_equal(base[k], same[k]) for k in base)
    ms = Predictor(m, device="cpu", tta=4, scales=(0.8, 1.0, 1.25)).probs(img)
    assert set(ms) == set(base) and all(ms[k].shape == base[k].shape for k in base)
    assert np.allclose(ms["cell"].sum(0), 1, atol=1e-4) and 0 <= ms["dist"].min() and ms["dist"].max() <= 1
    assert not np.array_equal(ms["cell"], base["cell"])  # the scales changed the average
    # the tiled path takes the scales too
    tiled = Predictor(m, device="cpu", tta=1, scales=(0.8, 1.0), tile=64, overlap=16, max_whole=32).probs(img)
    assert tiled["dist"].shape == (70, 90)
    import pytest
    with pytest.raises(AssertionError):
        Predictor(m, device="cpu", scales=(0.1,))