import copy
import torch
from nexora.adapters import LoRALinear, inject_lora, merge_lora
from nexora.model import ModelConfig, NexoraLM


def test_lora_freeze_and_merge():
    torch.manual_seed(2)
    base = torch.nn.Linear(8, 12)
    layer = LoRALinear(base, rank=2)
    x = torch.randn(3, 8)
    torch.testing.assert_close(layer(x), base(x))
    opt = torch.optim.SGD([layer.a, layer.b], lr=.1)
    old = base.weight.detach().clone()
    layer(x).square().mean().backward()
    opt.step()
    assert torch.equal(base.weight, old)
    assert layer.b.abs().sum() > 0
    torch.testing.assert_close(layer(x), layer.merged()(x))


def test_inject_only_adapter_trainable():
    m = NexoraLM(ModelConfig(hidden_size=32, layers=1, heads=4, kv_heads=2, intermediate_size=64))
    assert inject_lora(m) == ["blocks.0.attn.q", "blocks.0.attn.v"]
    assert all(name.endswith((".a", ".b")) for name, p in m.named_parameters() if p.requires_grad)
    merged = merge_lora(copy.deepcopy(m))
    assert not any(isinstance(x, LoRALinear) for x in merged.modules())