"""COMFYUI_ROOT=/path/to/ComfyUI python -m pytest code/tests/test_comfy_adapter.py"""
import importlib.util
import os
from pathlib import Path
import sys

import pytest
import torch
from safetensors.torch import save_file

sys.path.insert(0, os.environ['COMFYUI_ROOT'])
from comfy.cli_args import args
args.cpu = True
import comfy.ops
from comfy.ldm.moss_image.model import MOSSImage
from comfy.model_patcher import ModelPatcher

spec = importlib.util.spec_from_file_location('rapid_nodes', Path(__file__).parents[2] / 'comfyui/rapid-moss/nodes.py')
rapid = importlib.util.module_from_spec(spec)
spec.loader.exec_module(rapid)


@pytest.fixture
def adapter(tmp_path, monkeypatch):
    torch.manual_seed(123)
    dim, ffn, rank = 128, 192, 3
    lora = {}
    for name in rapid.PROJ:
        inp, out = (ffn, dim) if name == 'ffn_down' else (dim, 2 * ffn if name == 'ffn_up_gate' else dim)
        prefix = 'diffusion_model.layers.0.proj.' + name
        lora[prefix + '.lora_down.weight'] = torch.randn(rank, inp) * 0.03
        lora[prefix + '.lora_up.weight'] = torch.randn(out, rank) * 0.03
        lora[prefix + '.alpha'] = torch.tensor(float(rank))
    time = {
        'time_embedding.0.weight': torch.randn(dim, 256) * .01,
        'time_embedding.0.bias': torch.randn(dim) * .01,
        'time_embedding.2.weight': torch.randn(dim, dim) * .01,
        'time_embedding.2.bias': torch.randn(dim) * .01,
        'shared_block_modulation.bias': torch.randn(6 * dim) * .01,
        'video_head_modulation.bias': torch.randn(2 * dim) * .01,
    }
    save_file(lora, tmp_path / 'lora.safetensors')
    save_file(time, tmp_path / 'time_embedding.safetensors')
    save_file({'latent_mix.Q': torch.eye(128), 'step_codes': torch.tensor([1000, 968, 926, 871])},
              tmp_path / 'dataspace.safetensors')
    monkeypatch.setattr(rapid, '_adapter_dirs', lambda: {'test': str(tmp_path)})
    return lora


def make_model(fused):
    holder = torch.nn.Module()
    holder.diffusion_model = MOSSImage(dim=128, ffn_dim=192, num_heads=1, num_layers=1,
                                     fused_qkv=fused, dtype=torch.float32, operations=comfy.ops.manual_cast)
    for p in holder.parameters():
        torch.nn.init.normal_(p, std=.02)
    return ModelPatcher(holder, torch.device('cpu'), torch.device('cpu'))


@pytest.mark.parametrize('fused', [False, True])
def test_lora_apply_and_restore(adapter, fused):
    model = make_model(fused)
    layer = model.model.diffusion_model.layers[0]
    base_weights = {k: m.weight.detach().clone() for k, m in layer.proj.items()}
    patched, _, sigmas = rapid.RapidMOSSAdapter().apply(model, 'test', 4, .7)
    assert sigmas.tolist() == pytest.approx([1, .968, .926, .871, 0])
    assert not model.object_patches
    patched.patch_model(load_weights=False)
    for name, module in layer.proj.items():
        x = torch.randn(1, 5, module.weight.shape[1])
        expected = torch.nn.functional.linear(x, base_weights[name])
        names = ('q', 'k', 'v') if name == 'qkv' else (name,)
        deltas = []
        for proj in names:
            prefix = 'diffusion_model.layers.0.proj.' + proj
            deltas.append(.7 * (x @ adapter[prefix + '.lora_down.weight'].T) @ adapter[prefix + '.lora_up.weight'].T)
        expected += torch.cat(deltas, dim=-1)
        torch.testing.assert_close(module(x), expected)
        torch.testing.assert_close(module.weight, base_weights[name], rtol=0, atol=0)
    patched.unpatch_model()
    for name, module in layer.proj.items():
        x = torch.randn(1, 2, module.weight.shape[1])
        torch.testing.assert_close(module(x), torch.nn.functional.linear(x, base_weights[name]))


def test_sampler_matches_dataspace_recurrence():
    torch.manual_seed(4)
    z = torch.randn(1, 3, 2, 2)
    q = torch.randn(3, 3) * .1
    codes = torch.tensor([1000, 968, 926, 871])
    sigmas = torch.cat((codes.float() / 1000, torch.zeros(1)))
    calls = []
    def denoise(x, sigma):
        calls.append(sigma.clone())
        return x - sigma[:, None, None, None] * (x * .3 + sigma[:, None, None, None])
    expected = torch.zeros_like(z)
    for code in codes:
        x = z + torch.einsum('oc,bchw->bohw', q, expected)
        expected = z - (x * .3 + code / 1000)
    actual = rapid.sample_dataspace(denoise, z, sigmas, Q=q, codes=codes)
    torch.testing.assert_close(actual, expected)
    assert len(calls) == 4


def test_reject_too_many_steps(adapter):
    with pytest.raises(ValueError, match='4 step codes'):
        rapid.RapidMOSSAdapter().apply(make_model(True), 'test', 5, 1.)


@pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
@pytest.mark.parametrize('device', ['cpu', pytest.param('cuda', marks=pytest.mark.skipif(
    os.environ.get('RAPID_MOSS_TEST_CUDA') != '1', reason='set RAPID_MOSS_TEST_CUDA=1 for GPU checks'))])
def test_native_forward_fused_matches_split(adapter, dtype, device):
    model = make_model(False)
    fused = make_model(True)
    sd = model.model.state_dict()
    prefix = 'diffusion_model.layers.0.proj.'
    sd[prefix + 'qkv.weight'] = torch.cat([sd.pop(prefix + n + '.weight') for n in ('q', 'k', 'v')])
    fused.model.load_state_dict(sd)
    x = torch.randn(1, 128, 2, 3, device=device, dtype=dtype)
    ctx = torch.randn(1, 4, 4096, device=device, dtype=dtype)
    t = torch.tensor([.968], device=device)
    outputs = []
    for m in (model, fused):
        m.model.to(device=device, dtype=dtype)
        patched, _, _ = rapid.RapidMOSSAdapter().apply(m, 'test', 4, 1.)
        patched.patch_model(load_weights=False)
        with torch.no_grad():
            outputs.append(m.model.diffusion_model(x.clone(), t, ctx))
        patched.unpatch_model()
    assert all(torch.isfinite(out).all() for out in outputs)
    assert all(out.abs().max() < 10 for out in outputs)
    torch.testing.assert_close(*outputs, rtol=.01, atol=.001)