"""Merge N rank-64 LoRA adapters into one rank-64 adapter, offline on CPU.

WHY NOT A PER-KEY AVERAGE
-------------------------
A LoRA changes a weight matrix by a PRODUCT of two learned factors,

    dW_i = (alpha / r) * B_i @ A_i

so averaging the factors is not averaging the behaviour. For two adapters,

    (0.5*B1 + 0.5*B2) @ (0.5*A1 + 0.5*A2)
        = 0.25 * (B1@A1 + B1@A2 + B2@A1 + B2@A2)

carries cross terms pairing one adapter's input projection with another's output
projection. Those products belong to no adapter and to no blend of them.

WHAT THIS DOES
--------------
Targets the correct weighted sum of the deltas,

    M = sum_i  w_i * B_i @ A_i

(all parents here share alpha and r, so the (alpha/r) scaling cancels and the
merged adapter, saved at the same alpha and r, needs exactly B_new @ A_new = M).

M has rank up to 64*N, so for N > 1 it cannot be held exactly at rank 64. The
best rank-64 representation in the Frobenius sense is its truncated SVD -- what
peft's add_weighted_adapter does for combination_type="svd", computed here
straight from the adapter files so no base model is ever loaded.

Done efficiently: M = Bc @ Ac for the sqrt(w)-scaled concatenated parent factors,
so QR both sides and take the SVD of the small (64N x 64N) core rather than of
the full out x in matrix.

Truncation cost is measured per module against the exact untruncated sum and
reported, so the price of staying at rank 64 is on the record. Expect it to grow
with N: three parents have more to fit into the same 64 directions than two do.

Usage:
    merge_adapters.py OUT_DIR NAME=PATH:WEIGHT [NAME=PATH:WEIGHT ...]
"""

import json
import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = ""   # weight merge: CPU only, touch no card

import torch
from safetensors.torch import load_file, save_file

RANK = 64


def merge_module(pairs):
    """Return rank-64 (A, B) approximating sum_i w_i * B_i @ A_i, plus the exact target."""
    bcs, acs = [], []
    for a, b, w in pairs:
        s = w ** 0.5
        bcs.append(b * s)
        acs.append(a * s)
    bc = torch.cat(bcs, dim=1)              # [out, 64N]
    ac = torch.cat(acs, dim=0)              # [64N, in]

    qb, rb = torch.linalg.qr(bc, mode="reduced")
    qa, ra = torch.linalg.qr(ac.T, mode="reduced")
    u, s, vh = torch.linalg.svd(rb @ ra.T)

    sq = torch.sqrt(s[:RANK])
    b_new = qb @ (u[:, :RANK] * sq)
    a_new = (sq[:, None] * vh[:RANK, :]) @ qa.T
    return a_new, b_new, bc @ ac


def main() -> int:
    """Merge the adapters named on the command line into one adapter directory."""
    out_dir, specs = sys.argv[1], sys.argv[2:]
    parents = []
    for spec in specs:
        name, _, rest = spec.partition("=")
        path, _, weight = rest.rpartition(":")
        parents.append((name, path, float(weight)))

    total_w = sum(w for _, _, w in parents)
    print(f"merging {len(parents)} adapters (weights sum to {total_w:.4f}):")
    for n, p, w in parents:
        print(f"   {w:.4f}  {n}  {p}")

    tensors = [load_file(os.path.join(p, "adapter_model.safetensors")) for _, p, _ in parents]
    keys = set(tensors[0])
    for t in tensors[1:]:
        assert set(t) == keys, "parent key sets differ"

    modules = sorted({k.rsplit(".lora_", 1)[0] for k in keys})
    print(f"{len(keys)} tensors, {len(modules)} modules", flush=True)

    out_t, errs, worst = {}, [], (0.0, None)
    for i, m in enumerate(modules, 1):
        ka, kb = f"{m}.lora_A.weight", f"{m}.lora_B.weight"
        pairs = [(t[ka].float(), t[kb].float(), w)
                 for t, (_, _, w) in zip(tensors, parents)]
        for a, b, _ in pairs[1:]:
            assert a.shape == pairs[0][0].shape and b.shape == pairs[0][1].shape, f"shape mismatch at {m}"

        a_new, b_new, target = merge_module(pairs)
        err = float(torch.linalg.matrix_norm(b_new @ a_new - target)
                    / torch.linalg.matrix_norm(target))
        errs.append(err)
        if err > worst[0]:
            worst = (err, m)

        out_t[ka] = a_new.to(torch.bfloat16)
        out_t[kb] = b_new.to(torch.bfloat16)
        if i % 64 == 0:
            print(f"  {i}/{len(modules)}, running mean rel-err {sum(errs)/len(errs):.4f}", flush=True)

    os.makedirs(out_dir, exist_ok=True)
    save_file(out_t, os.path.join(out_dir, "adapter_model.safetensors"), metadata={"format": "pt"})
    cfg = json.load(open(os.path.join(parents[0][1], "adapter_config.json")))
    json.dump(cfg, open(os.path.join(out_dir, "adapter_config.json"), "w"), indent=2)

    mean_err = sum(errs) / len(errs)
    open(os.path.join(out_dir, "README.md"), "w").write(
        "# " + os.path.basename(out_dir) + " -- INFORMAL / NON-GATE\n\n"
        "Offline CPU merge of " + str(len(parents)) + " rank-64 LoRA adapters:\n"
        + "".join(f"  {w}  {n}  {p}\n" for n, p, w in parents) +
        "\nMethod: targets M = sum_i w_i * B_i @ A_i exactly, then takes the best\n"
        "rank-64 representation of it by truncated SVD (equivalent to peft\n"
        "add_weighted_adapter combination_type='svd', computed directly from the\n"
        "adapter files so no base model is loaded). A per-key average of the LoRA\n"
        "factors was rejected: it introduces cross terms and is not a blend.\n\n"
        f"Rank-64 truncation cost vs the exact rank-{64*len(parents)} sum:\n"
        f"  mean {mean_err:.4f}   min {min(errs):.4f}   max {max(errs):.4f} (at {worst[1]})\n")

    print(f"\nwrote {out_dir}")
    print(f"tensors: {len(out_t)}")
    print(f"TRUNCATION rel-err vs exact rank-{64*len(parents)} sum: "
          f"mean {mean_err:.4f}  min {min(errs):.4f}  max {max(errs):.4f}")
    print(f"worst module: {worst[1]}")
    return 0


if __name__ == "__main__":
    sys.exit(main())