"""
Cross-relation LoRA-B alignment: are the steering directions of different
scene→object hallucination adapters the SAME axis or relation-specific?

For each relation's adapter, take the per-layer LoRA-B direction of a
residual-writing module (default mlp.down_proj, B (d_model,r) reduced over rank →
unit v̂_l). Then per layer compute pairwise cosine similarity between relations:
  bathroom_toilet ↔ kitchen_oven ↔ livingroom_tv.

Random-direction baseline ≈ 1/sqrt(d_model) (≈0.016 for d=4096): |cos| above that
= real alignment. High |cos| across relations ⇒ a shared "scene→object / hallucinated-
object" axis; near-zero ⇒ each relation has its own direction.

Saves: per-layer pairwise cosine JSON + line plot.
"""
import argparse, json, os, itertools
import numpy as np, torch as t
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt

from mechanistic_interp.alternative_steering import load_lora_b_dirs

REL = {
    "bathroom_toilet": "/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter",
    "kitchen_oven":    "/data/caotue/multilayer-sae/adv_gen_outputs/run_kitchen_oven_v2/lora_adapter",
    "livingroom_tv":   "/data/caotue/multilayer-sae/adv_gen_outputs/run_livingroom_tv_v2/step_450",
}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--lora_module", default="down_proj", choices=["down_proj", "o_proj"])
    ap.add_argument("--b_reduce", default="mean", choices=["mean", "svd"])
    ap.add_argument("--n_layers", type=int, default=32)
    ap.add_argument("--out_prefix", default="mechanistic_interp/graph/lora_relations_sim")
    args = ap.parse_args()

    from transformers import AutoConfig
    d_model = AutoConfig.from_pretrained(args.model_name).text_config.hidden_size
    layers = list(range(args.n_layers))
    rng_base = 1.0 / np.sqrt(d_model)

    dirs = {r: load_lora_b_dirs(p, layers, args.lora_module, args.b_reduce, d_model, "cpu", t.float32)
            for r, p in REL.items()}
    rels = list(REL)
    pairs = list(itertools.combinations(rels, 2))

    cos = {f"{a}|{b}": [] for a, b in pairs}
    for l in layers:
        for a, b in pairs:
            va, vb = dirs[a][l], dirs[b][l]
            cos[f"{a}|{b}"].append(float(t.dot(va, vb) / (va.norm() * vb.norm() + 1e-8)))

    print(f"random-direction baseline |cos| ~ {rng_base:.4f} (d={d_model})\n")
    print("pair                              mean|cos|  max|cos|  signed-mean")
    for k, v in cos.items():
        v = np.array(v)
        print(f"  {k:32s}  {np.abs(v).mean():.3f}     {np.abs(v).max():.3f}     {v.mean():+.3f}")

    os.makedirs(os.path.dirname(args.out_prefix), exist_ok=True)
    json.dump({"layers": layers, "rng_baseline": rng_base, "module": args.lora_module,
               "b_reduce": args.b_reduce, "cosine": cos}, open(f"{args.out_prefix}.json", "w"), indent=2)

    fig, axes = plt.subplots(1, 2, figsize=(15, 5))
    for k, v in cos.items():
        axes[0].plot(layers, v, "o-", ms=3, label=k)
        axes[1].plot(layers, np.abs(v), "o-", ms=3, label=k)
    axes[0].axhline(0, color="gray", ls=":"); axes[0].set_ylabel("signed cosine")
    axes[1].axhline(rng_base, color="red", ls="--", lw=0.8, label=f"random {rng_base:.3f}")
    axes[1].set_ylabel("|cosine| (alignment)")
    for ax, ttl in ((axes[0], "signed"), (axes[1], "|cos|")):
        ax.set_xlabel("layer"); ax.set_title(f"LoRA-B cross-relation cosine ({ttl})"); ax.grid(alpha=0.3); ax.legend(fontsize=8)
    fig.suptitle(f"LoRA-B direction alignment across relations ({args.lora_module}/{args.b_reduce})")
    fig.tight_layout(); fig.savefig(f"{args.out_prefix}.png", dpi=150, bbox_inches="tight")
    print(f"\nsaved {args.out_prefix}.png / .json")


if __name__ == "__main__":
    main()