import pytest
import torch


@pytest.fixture(autouse=True)
def _needs_comfy(comfyui):
    """Imports comfy."""


def _load(tiny, **v):
    from comfy.patcher_extension import WrappersMP
    from modules.sampling.modifiers.nodes import FunUILoadModifiers
    patched, status = FunUILoadModifiers.execute(
        tiny.patcher, {"sharpness": {"enabled": True, **v}}).result
    wrap = [w for ws in patched.wrappers[WrappersMP.APPLY_MODEL].values() for w in ws][-1]
    return patched, wrap, status


def _to(i, steps=4):
    s = torch.linspace(1.0, 0.0, steps + 1)
    return {"sample_sigmas": s, "sigmas": s[i:i + 1]}


def test_off_installs_nothing(tiny_h3):
    from comfy.patcher_extension import WrappersMP
    from modules.sampling.modifiers.nodes import FunUILoadModifiers
    patched, _ = FunUILoadModifiers.execute(tiny_h3.patcher, {"sharpness": {"enabled": False}}).result
    assert not patched.wrappers.get(WrappersMP.APPLY_MODEL)


def test_a_plain_latent_is_pushed_away_from_the_previous_prediction_in_the_last_part_only(tiny_h3):
    _p, wrap, status = _load(tiny_h3, amount=0.5, last_part=0.5)
    assert "sharpness" in status
    x = torch.randn(1, 4, 3, 8, 8)
    outs = [x * f for f in (0.9, 0.5, 0.2, 0.1)]
    ts = torch.linspace(1.0, 0.0, 5)
    got = [wrap(lambda *a, o=outs[i], **kw: o, x, ts[i:i + 1], None, None, None, _to(i)) for i in range(4)]
    assert torch.equal(got[0], outs[0]) and torch.equal(got[1], outs[1])      # outside the last half
    kept = list(got)                                       # the boost is measured against the sharpened answer before
    for i in (2, 3):
        assert torch.allclose(got[i], outs[i] + 0.25 * (outs[i] - kept[i - 1]), atol=1e-5)


def test_a_packed_model_is_edited_on_its_answer_including_the_last_step_and_never_the_sound(tiny_h3):
    from conftest import packed_av, unpacked
    _p, wrap, _ = _load(tiny_h3, amount=1.0, last_part=1.0)
    x0, shapes = packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5))
    answers = [x0 * f for f in (0.9, 0.5, 0.2, 0.1)]
    ts = torch.linspace(1.0, 0.0, 5)
    got = [wrap(lambda *a, o=answers[i], **k: o, x0, ts[i:i + 1], None, None, None, _to(i), latent_shapes=shapes)
           for i in range(4)]
    assert torch.equal(got[0], answers[0])                 # nothing to compare on step 0
    for i in (1, 3):
        assert not torch.allclose(unpacked(got[i], shapes)[0], unpacked(answers[i], shapes)[0])
        assert torch.equal(unpacked(got[i], shapes)[1], unpacked(answers[i], shapes)[1])   # sound untouched


def test_a_hand_over_to_a_bigger_picture_is_never_compared_with_the_smaller_one_before_it(tiny_h3):
    """Two stages: the step index counts the whole schedule, and the first call at the full size has nothing of its
    own size to compare with, so it is not sharpened against the first stage's smaller answer."""
    _p, wrap, _ = _load(tiny_h3, amount=1.0, last_part=1.0)
    ts = torch.linspace(1.0, 0.0, 5)
    small = torch.randn(1, 4, 3, 4, 4)
    big = torch.randn(1, 4, 3, 8, 8)
    first = wrap(lambda *a, **k: small, small, ts[0:1], None, None, None, _to(0))
    second = wrap(lambda *a, **k: small * 0.5, small, ts[1:2], None, None, None, _to(1))
    assert torch.equal(first, small) and not torch.equal(second, small * 0.5)       # same size: compared
    handed = wrap(lambda *a, **k: big, big, ts[2:3], None, None, None, _to(2))
    assert torch.equal(handed, big)                                                  # new size: left as the model said
    after = wrap(lambda *a, **k: big * 0.5, big, ts[3:4], None, None, None, _to(3))
    assert not torch.equal(after, big * 0.5)                                         # and sharpening goes on in stage 2