"""Experiment 2: causal swap interventions gated by the J-lens.

At each (layer, position) where the J-lens reads the source concept, transfer
the activation content source -> target (embedding write basis; see core.py).

(a) Two-hop: swap spider<->ant in the workspace -> answer flips 8 -> 6.
(b) Broadcast: swap France->China under different templates.
(c) Capital-question grid over country pairs.
"""

import itertools, os, sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import JLens, chat_ids, load_model

RES = Path(os.environ.get("JLENS_RESULTS", "results"))


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt",
                    map_location="cuda")
    jl = JLens(model, tok, ck["Jbar"].cuda())
    print(f"Jbar from {ck.get('n_prompts')} prompts")

    # (a) spider -> ant
    SP = [(" spider", " ant"), (" spiders", " ants"), ("spider", "ant")]
    q = "The animal that spins webs has how many legs? Answer with just a number."
    ids = chat_ids(tok, q)
    print("\n--- two-hop swap: spider -> ant (baseline:",
          repr(jl.generate(ids, 6)), ") ---")
    for thr in (0.002, 0.005, 0.01):
        for alpha in (1.0, 1.5, 2.0):
            out, fired = jl.generate_swapped(ids, SP, thr=thr, alpha=alpha,
                                             max_new_tokens=6)
            print(f"thr={thr} a={alpha}: {out!r} (fired {len(fired)} slots)")
    # paper write-basis for comparison
    out, _ = jl.generate_swapped(ids, SP, thr=0.005, alpha=1.0, write="jlens",
                                 max_new_tokens=6)
    print(f"paper jlens-write thr=0.005 a=1.0: {out!r}")
    # control: unrelated swap
    out, _ = jl.generate_swapped(ids, [(" spider", " piano")], thr=0.005,
                                 alpha=1.0, max_new_tokens=6)
    print(f"control spider->piano: {out!r}")

    # reverse: ant prompt -> spider
    q2 = ("The insect that builds colonies and lifts many times its own "
          "weight has how many legs? Answer with just a number.")
    ids2 = chat_ids(tok, q2)
    print("\nant prompt baseline:", repr(jl.generate(ids2, 6)))
    for alpha in (1.0, 1.5, 2.0):
        out, fired = jl.generate_swapped(
            ids2, [(" ant", " spider"), (" ants", " spiders"), ("ant", "spider")],
            thr=0.005, alpha=alpha, max_new_tokens=6)
        print(f"swap ant->spider a={alpha}: {out!r} (fired {len(fired)})")

    # (b) broadcast France -> China across templates
    print("\n--- broadcast: France -> China (thr=0.01, a=1.0) ---")
    FR = [(" France", " China"), ("France", "China"), (" French", " Chinese")]
    templates = [
        ("capital", "What is the capital of France? Answer with just the city name."),
        ("language", "What language is spoken in France? Answer with one word."),
        ("continent", "Which continent is France on? Answer with one word."),
        ("currency", "What currency is used in France? Answer with one word."),
        ("river", "Name a famous river in France. Answer with one word."),
        ("food", "Name a famous dish from France. Answer with a short phrase."),
    ]
    for tag, qq in templates:
        ids = chat_ids(tok, qq)
        base = jl.generate(ids, 8)
        sw, fired = jl.generate_swapped(ids, FR, thr=0.01, alpha=1.0,
                                        max_new_tokens=8)
        print(f"[{tag:9}] base={base!r:30} swapped={sw!r} ({len(fired)} slots)")

    # (c) capital grid over country pairs
    print("\n--- capital-question swap grid (thr=0.01, a=1.0) ---")
    countries = {"France": ("Paris", "French"), "China": ("Beijing", "Chinese"),
                 "Japan": ("Tokyo", "Japanese"), "Egypt": ("Cairo", "Egyptian"),
                 "Brazil": ("Brasília", "Brazilian"), "Canada": ("Ottawa", "Canadian")}
    hits = tries = 0
    rows = []
    for (src, (scap, sadj)), (tgt, (tcap, tadj)) in itertools.permutations(
            countries.items(), 2):
        qq = f"What is the capital of {src}? Answer with just the city name."
        ids = chat_ids(tok, qq)
        pairs = [(" " + src, " " + tgt), (src, tgt), (" " + sadj, " " + tadj)]
        sw, fired = jl.generate_swapped(ids, pairs, thr=0.01, alpha=1.0,
                                        max_new_tokens=8)
        ok = tcap.lower().replace("í", "i").split()[0][:5] in \
            sw.lower().replace("í", "i")
        hits += ok
        tries += 1
        rows.append((src, tgt, tcap, sw, ok))
        print(f"{src:>7}->{tgt:<7} expect {tcap:<9} got {sw!r} {'OK' if ok else ''}")
    print(f"\nswap grid success: {hits}/{tries}")
    torch.save(rows, RES / "swap_grid.pt")


if __name__ == "__main__":
    main()