"""vela3-fast span adapters for Core ML: build the release SchemaNet (frozen backbone + adapter), merge the LoRA deltas into plain
Linear weights (W + alpha / r * B A), and export the merged encoder with the same fixed-length graph as the decision backbone.

    python span_export.py                     # parity: merged export vs the release SchemaNet hidden states
    python span_export.py convert L [adapter] # -> coreml/vela3_<adapter>_L<L>.mlpackage (default adapter: span_sets)"""
import sys
import torch
from torch import nn
from vela3_export import Vela3EncoderExport, load_decision

sys.path.insert(0, "src")
import vela3_schema as VS  # noqa: E402


def build(adapter="span_sets"):
    dec = load_decision()
    return dec, VS.build_adapter(dec, f"src/{adapter}_adapter.safetensors", "cpu")


@torch.no_grad()
def merge_lora(net):
    """Replace every LoRALinear by an nn.Linear with the merged weight (a copy: the shared backbone tensors are untouched)."""
    n = 0
    for mod in list(net.modules()):
        for name, child in list(mod.named_children()):
            if isinstance(child, VS.LoRALinear):
                b = child.base
                lin = nn.Linear(b.in_features, b.out_features, bias=b.bias is not None)
                lin.weight.copy_(b.weight + child.scale * (child.lora_B @ child.lora_A))
                if b.bias is not None:
                    lin.bias.copy_(b.bias)
                setattr(mod, name, lin); n += 1
    return n


if __name__ == "__main__":
    adapter = sys.argv[3] if len(sys.argv) > 3 else "span_sets"
    dec, net = build(adapter)
    if sys.argv[1:2] == ["convert"]:
        import numpy as np, coremltools as ct
        import ct_patch  # noqa: F401
    ids = torch.randint(5, 49152, (1, 300)); ids[0, 0] = 0
    with torch.no_grad():
        ref = net.hidden(ids, torch.ones_like(ids))[0]
        print("merged LoRA layers:", merge_lora(net))
        if sys.argv[1:2] != ["convert"]:
            for L in (300, 512):
                g = Vela3EncoderExport(net.encoder, L).eval()
                x = torch.zeros(1, L, dtype=torch.int32); x[0, :300] = ids[0]; mk = torch.zeros(1, L, dtype=torch.int32); mk[0, :300] = 1
                h = g(x, mk)[0, :300]
                print(f"L={L} max|dh|={(h - ref).abs().max():.2e} min cos={nn.functional.cosine_similarity(h, ref, dim=-1).min():.6f}")
            sys.exit()
    L = int(sys.argv[2])
    g = Vela3EncoderExport(net.encoder, L).eval()
    x = torch.zeros(1, L, dtype=torch.int32); mk = torch.ones(1, L, dtype=torch.int32)
    with torch.no_grad():
        tr = torch.jit.trace(g, (x, mk))
    ml = ct.convert(tr, inputs=[ct.TensorType("input_ids", shape=(1, L), dtype=np.int32), ct.TensorType("attention_mask", shape=(1, L), dtype=np.int32)],
                    outputs=[ct.TensorType("hidden", dtype=np.float16)], minimum_deployment_target=ct.target.macOS15,
                    compute_precision=ct.precision.FLOAT16, convert_to="mlprogram")
    ml.short_description = f"vela3-fast UltraBERT-380M encoder + {adapter} adapter (LoRA merged), fixed length {L}"
    ml.save(f"coreml/vela3_{adapter}_L{L}.mlpackage"); print("saved", f"coreml/vela3_{adapter}_L{L}.mlpackage")