"""
Gradient-ascent causal influence map: toilet → bathroom across layers.

Idea
────
Using the trained attention sequence probes (SequenceLayerProbes) for *toilet*
and *bathroom*, ask: if we nudge the residual stream at layer ``l`` in the direction
that increases the *toilet* readout, how much does the *bathroom* readout change at
every downstream layer ``l' >= l``?

Per image (run on TOILET-ONLY images: toilet=1 & bathroom=0):
  1. Generate a caption, forced-forward, capture resid_post at all layers; slice to
     the generated-caption tokens (the positions the probe pools over).
  2. score_bath_l(h_l): bathroom probe at layer l over the caption sequence (scalar).
  3. g = ∂ score_toilet_l / ∂ h_l   (gradient w.r.t. the layer-l caption activations;
     this is local to the probe — no model backward needed), normalized by its total
     magnitude: ĝ = g / ‖g‖  (single Frobenius norm over the whole (T, d_model) tensor).
  4. Intervene: h_l' = h_l + alpha * ‖h_l‖ * ĝ   (gradient ascent along the unit toilet
     direction, step scaled by the caption-block residual norm ‖h_l‖ so alpha is a
     *fraction* of residual magnitude — comparable across layers, since later layers have
     much larger ‖h_l‖ and a raw fixed step would be relatively tiny there).
  5. Forward with h_l patched; read bathroom score at every layer l' >= l:
        delta[l, l'] = score_bathroom_l'(intervened) - score_bathroom_l'(baseline)   (NOT /alpha)
  6. Average over images → heatmap (intervention layer l × readout layer l').

A positive band above the diagonal = pushing toilet-ness at l causally raises the
bathroom readout downstream → evidence for a toilet→bathroom mechanism.

Saves: heatmap PNG + the raw (L, L) matrix as JSON.
"""

import argparse
import json
import os
import random

import numpy as np
import torch as t
from PIL import Image
from transformers import AutoConfig, LlavaProcessor

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from mechanistic_interp.sequence_probe import sequence_layer_probes_from_checkpoint
from hallucination.mechanistic_interp.compare_baselines import apply_efuf_edit, apply_nullu_edit

_HOOK_SUFFIX = {"pre": "hook_resid_pre", "mid": "hook_resid_mid", "post": "hook_resid_post"}


# ── Model-variant loading (base / lora / nullu / efuf), per attribution_patching ──
def _decoder_layers(model):
    lm = getattr(model, "language_model", None) or getattr(getattr(model, "model", None), "language_model", None)
    inner = getattr(lm, "model", None)
    if inner is not None and hasattr(inner, "layers"):
        return inner.layers
    return lm.layers


def load_variant_model(args, dtype, device):
    """Build the chosen model variant. NOTE: the steering/readout PROBES are loaded
    separately (from base-trained checkpoints) — only the MODEL changes here."""
    v = args.variant
    if v == "lora":
        from model.llava.hooked_lora_llava import HookedLoRALlava
        model = HookedLoRALlava.from_pretrained(
            args.model_name, torch_dtype=dtype, device_map={"": device}).eval()
        lora_dir = args.lora_path if os.path.isdir(args.lora_path) else os.path.dirname(args.lora_path)
        model.load_lora_adapter(lora_dir, merge=True)
        print(f"[grad-ascent] model=lora (merged {lora_dir})")
    else:
        model = HookedSAELlavaConditionalGeneration.from_pretrained(
            args.model_name, torch_dtype=dtype, device_map={"": device}).eval()
        if v == "efuf":
            n = apply_efuf_edit(model, args.efuf_path); print(f"[grad-ascent] model=efuf ({n} proj tensors, {args.efuf_path})")
        elif v == "nullu":
            n = apply_nullu_edit(model, _decoder_layers(model), args.nullu_path, args.nullu_lowest, args.nullu_highest)
            print(f"[grad-ascent] model=nullu ({n} layers [{args.nullu_lowest},{args.nullu_highest}), {args.nullu_path})")
        else:
            print(f"[grad-ascent] model=base")
    return model


def hp_name(layer: int, hook_type: str) -> str:
    return f"model.language_model.layers.{layer}.{_HOOK_SUFFIX[hook_type]}"


def caption_slice(attn_mask, asst_text, processor, max_seq_tokens):
    """Return (start, end) indices of the generated-caption tokens in the sequence.

    Caption tokens sit flush at the tail of the real (unpadded) region; we keep the
    leading min(cap_len, max_seq_tokens) of them — matching training's a[:max_seq_tokens].
    """
    seq_len = int(attn_mask.sum().item())
    cap_len = len(processor.tokenizer(asst_text, add_special_tokens=False)["input_ids"])
    if cap_len <= 0:
        return None
    start = seq_len - cap_len
    end = start + min(cap_len, max_seq_tokens)
    if start < 0:
        return None
    return start, end


@t.no_grad()
def generate_caption(model, processor, image, question, device, max_new_tokens):
    prompt = f"USER: <image>\n{question}\nASSISTANT:"
    inp = processor(images=[image], text=[prompt], return_tensors="pt").to(device)
    out = model.generate(**inp, do_sample=False, num_beams=1, use_cache=True,
                         max_new_tokens=max_new_tokens)
    cap = processor.batch_decode(out, skip_special_tokens=True)[0]
    return cap.split("ASSISTANT:")[-1].strip()


def probe_logit(probe_module, layer, seq_feats):
    """Scalar logit for one image from one layer's attention probe. seq_feats: (1,T,d)."""
    kpm = t.zeros(seq_feats.shape[:2], dtype=t.bool, device=seq_feats.device)  # no padding
    return probe_module.probes[probe_module._idx[layer]](seq_feats, kpm).squeeze(0)


def probe_score(probe_module, layer, seq_feats):
    """Bounded probe score σ(logit) = P(concept present) ∈ [0,1] (logits are unbounded)."""
    return t.sigmoid(probe_logit(probe_module, layer, seq_feats))


def plot_heatmaps(deltas, alphas, out, title, mode="meannorm", cbar_label="Δ toilet score σ",
                  steer_name="bathroom", readout_name="toilet"):
    """Heatmap grid of {alpha: (L,L) Δ matrix} in one of two modes:
      'normal'   — each α panel on its OWN diverging scale + own colorbar (true per-α magnitude).
      'meannorm' — every cell ÷ mean|Δ| over all panels; ONE shared, robust (p99) scale
                   so panels are comparable and colour reads as '× typical effect'.
    """
    nA = len(alphas)
    ncol = min(3, nA)
    nrow = -(-nA // ncol)
    fig, axes = plt.subplots(nrow, ncol, figsize=(5.0 * ncol, 4.5 * nrow), squeeze=False)
    if mode == "meannorm":
        allabs = np.abs(np.concatenate([deltas[a][np.isfinite(deltas[a])].ravel() for a in alphas]))
        mean_mag = float(allabs.mean()) + 1e-12
        vmax = float(np.percentile(allabs / mean_mag, 99))
        print(f"[grad-ascent] mean|Δ|={mean_mag:.5f}; shared p99 cap={vmax:.2f}× mean "
              f"(max={allabs.max()/mean_mag:.1f}× mean)")
        im = None
        for k, a in enumerate(alphas):
            ax = axes[k // ncol][k % ncol]
            im = ax.imshow(deltas[a] / mean_mag, origin="upper", cmap="RdBu_r",
                           vmin=-vmax, vmax=vmax, aspect="auto")
            ax.set_title(f"α = {a}")
            ax.set_xlabel(f"readout layer l′ ({readout_name})")
            ax.set_ylabel(f"intervention layer l ({steer_name})")
        for k in range(nA, nrow * ncol):
            axes[k // ncol][k % ncol].axis("off")
        cbar = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.02)
        cbar.set_label(f"{cbar_label} / mean|Δ|  (× typical effect)")
    else:  # normal — per-panel scale
        for k, a in enumerate(alphas):
            ax = axes[k // ncol][k % ncol]
            d = deltas[a]
            vmax = float(np.nanmax(np.abs(d)))
            im = ax.imshow(d, origin="upper", cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
            ax.set_title(f"α = {a}")
            ax.set_xlabel(f"readout layer l′ ({readout_name})")
            ax.set_ylabel(f"intervention layer l ({steer_name})")
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=cbar_label)
        for k in range(nA, nrow * ncol):
            axes[k // ncol][k % ncol].axis("off")
        fig.tight_layout()
    fig.suptitle(title, fontsize=12, y=1.03)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"[grad-ascent] saved {out}")


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("--device_id", type=int, default=0)
    ap.add_argument("--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"])

    # Model variant — which residual stream to steer/measure on. Probes stay base-trained.
    ap.add_argument("--variant", default="base", choices=["base", "lora", "nullu", "efuf"],
                    help="Model variant the gradient ascent runs on (probes still load from base).")
    ap.add_argument("--lora_path", default="/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter",
                    help="LoRA(ours) adapter dir (variant=lora).")
    ap.add_argument("--efuf_path", default="/data/caotue/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_paper_10ep/epoch_002.pth",
                    help="EFUF .pth checkpoint (variant=efuf).")
    ap.add_argument("--nullu_path", default="/data/caotue/nullu/edited_models/LLaVA-7B-top4-0-32-bathroom_toilet",
                    help="Nullu edited-model HF dir (variant=nullu).")
    ap.add_argument("--nullu_lowest", type=int, default=8)
    ap.add_argument("--nullu_highest", type=int, default=32)

    # Probes — base-trained checkpoints (the readout/steer directions). NOT the variant.
    # REVERSED: steer = toilet probe, readout = bathroom probe.
    ap.add_argument("--bath_probe", default="/data/caotue/latent_probes/seqprobes_toilet/post/seqprobe.pth",
                    help="Base-trained SequenceLayerProbes for the steered direction (toilet in reversed mode).")
    ap.add_argument("--toilet_probe", default="/data/caotue/latent_probes/seqprobes_bathroom/post/seqprobe.pth",
                    help="Base-trained SequenceLayerProbes for the measured readout (bathroom in reversed mode).")

    ap.add_argument("--image_folder", default="/data/caotue/CC3M-Dataset/cc3m_images")
    ap.add_argument("--samples_json", default="mechanistic_interp/toilet_bathroom/samples.json",
                    help="Sample file with per-image base_mentions_object flags.")
    ap.add_argument("--base_prompt", default="Describe this image.",
                    help="Prompt whose base_mentions_object flag selects images.")
    ap.add_argument("--base_mentions", default="any", choices=["false", "true", "any"],
                    help="Filter toilet-only images by base_mentions_object: 'any' = no filter (default for reversed).")
    ap.add_argument("--id_col", default="image_id")
    ap.add_argument("--toilet_col", default="toilet")
    ap.add_argument("--question", default="Describe this image.")

    # Alternative image selection from an HF dataset (for relations with NO samples.json).
    # When --hf_dataset is set, scene-only images = rows with scene_col==1 & object_col==0
    # (the direct analog of bathroom-only), and --samples_json / --base_mentions are ignored.
    ap.add_argument("--hf_dataset", default=None,
                    help="If set, pick scene-only images (--scene_col==1 & --object_col==0) "
                         "from this HF dataset instead of --samples_json.")
    ap.add_argument("--split", default="validation",
                    help="HF split used for --hf_dataset selection (default validation).")
    ap.add_argument("--scene_col", default=None,
                    help="Steer concept column (present==1) in --hf_dataset.")
    ap.add_argument("--object_col", default=None,
                    help="Readout concept column (absent==0) in --hf_dataset.")
    ap.add_argument("--steer_name", default="toilet",
                    help="Display name of the steered concept (plot labels/title).")
    ap.add_argument("--readout_name", default="bathroom",
                    help="Display name of the readout concept (plot labels/title).")
    ap.add_argument("--forced_text", default=None,
                    help="If set, skip generation and force the ASSISTANT answer to exactly "
                         "this string (e.g. 'This image features a bathroom with a'). The probe "
                         "then reads these forced answer tokens — a controlled, constant context "
                         "across all images. If unset, use the model's freely-generated caption.")
    ap.add_argument("--num_images", type=int, default=100,
                    help="Random-subsample cap (0 = ALL). Default 100 (shuffled by --seed).")
    ap.add_argument("--seed", type=int, default=0, help="Seed for the random image subsample.")

    ap.add_argument("--alphas", type=float, nargs="+", default=[0.05, 0.1, 0.2, 0.4, 0.8],
                    help="Gradient-ascent step sizes to sweep. h_l' = h_l + alpha*‖h_l‖*ĝ, "
                         "so alpha is a FRACTION of the caption-block residual norm (alpha=1 "
                         "⇒ step magnitude == ‖h_l‖). g computed once per layer; only the "
                         "patched forward repeats per alpha.")
    ap.add_argument("--hook_type", default="post", choices=["pre", "mid", "post"])
    ap.add_argument("--plot_mode", default="meannorm", choices=["normal", "meannorm"],
                    help="normal = each α its own colorbar; meannorm = ÷mean|Δ|, shared p99 scale.")
    ap.add_argument("--max_new_tokens", type=int, default=64)
    ap.add_argument("--max_seq_tokens", type=int, default=64)

    ap.add_argument("--out", default="mechanistic_interp/graph/gradient_ascent_toilet2bath.png")
    ap.add_argument("--out_json", default="mechanistic_interp/graph/gradient_ascent_toilet2bath.json")
    args = ap.parse_args()

    dtype = {"float32": t.float32, "float16": t.float16, "bfloat16": t.bfloat16}[args.dtype]
    device = f"cuda:{args.device_id}" if t.cuda.is_available() else "cpu"

    n_layers = AutoConfig.from_pretrained(args.model_name).text_config.num_hidden_layers
    layers = list(range(n_layers))

    print(f"[grad-ascent] device={device} dtype={dtype} layers={n_layers} alphas={args.alphas}")
    print(f"[grad-ascent] steer={args.steer_name} ({args.bath_probe})")   # bath_probe = toilet probe in reverse
    print(f"[grad-ascent] readout={args.readout_name} ({args.toilet_probe})")  # toilet_probe = bathroom probe in reverse
    if args.forced_text:
        print(f"[grad-ascent] FORCED answer: '{args.forced_text}'")
    else:
        print(f"[grad-ascent] answer: freely-generated caption")

    # ── Select "scene-only" images = steer concept present, readout concept absent.
    #    Two backends:
    #      (a) --hf_dataset set: rows with scene_col==1 & object_col==0 (no samples.json
    #          / base_mentions filter — for relations that lack one).
    #      (b) else --samples_json: category==bathroom_only & toilet==0, filtered by
    #          base_mentions_object per --base_mentions (flagship bathroom→toilet study).
    stem_to_file = {}
    for root, _, files in os.walk(args.image_folder):
        for fn in files:
            if fn.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
                stem_to_file[os.path.splitext(fn)[0]] = os.path.join(root, fn)
    ids = []
    if args.hf_dataset:
        if not args.scene_col or not args.object_col:
            raise SystemExit("--hf_dataset requires --scene_col and --object_col.")
        from datasets import load_dataset
        ds = load_dataset(args.hf_dataset, split=args.split)
        for row in ds:
            if row[args.scene_col] == 1 and row[args.object_col] == 0:
                stem = os.path.splitext(os.path.basename(str(row[args.id_col])))[0]
                if stem in stem_to_file:
                    ids.append(stem)
        sel_desc = f"{args.scene_col}=1 & {args.object_col}=0 from {args.hf_dataset}[{args.split}]"
    else:
        want = {"false": False, "true": True, "any": None}[args.base_mentions]
        data = json.load(open(args.samples_json))
        for it in data:
            if it.get("category") != "toilet_only" or it.get("toilet") != 1:
                continue
            pr = it.get("prompt_results", {}).get(args.base_prompt, {})
            if want is not None and pr.get("base_mentions_object") is not want:
                continue
            stem = os.path.splitext(os.path.basename(it[args.id_col]))[0]
            if stem in stem_to_file:
                ids.append(stem)
        halluc = {"false": "NON-hallucinating", "true": "HALLUCINATING", "any": "all"}[args.base_mentions]
        sel_desc = f"{halluc} toilet-only"
    if args.num_images and args.num_images > 0:
        random.Random(args.seed).shuffle(ids)          # random subsample
        ids = ids[:args.num_images]
    print(f"[grad-ascent] scene-only images ({sel_desc}): {len(ids)} (variant={args.variant}, "
          f"probes=base, alphas={args.alphas})")

    # ── Model (variant) + probes (base-trained) ──────────────────────────────
    model = load_variant_model(args, dtype, device)
    processor = LlavaProcessor.from_pretrained(args.model_name)
    bath = sequence_layer_probes_from_checkpoint(args.bath_probe, device)
    toilet = sequence_layer_probes_from_checkpoint(args.toilet_probe, device)
    bath.eval(); toilet.eval()

    # delta[alpha][l, l'] accumulator over the upper triangle (l' >= l); NaN below diagonal.
    tri = np.triu(np.ones((n_layers, n_layers))) > 0
    delta_sum = {a: np.where(tri, 0.0, np.nan) for a in args.alphas}
    count = 0

    hps = [hp_name(l, args.hook_type) for l in layers]

    def make_cap(name, store):
        def _fn(act, hook):
            store[name] = act
        return _fn

    for i, stem in enumerate(ids):
        try:
            img = Image.open(stem_to_file[stem]).convert("RGB")
        except Exception:
            continue
        if args.forced_text:
            asst = args.forced_text
        else:
            asst = generate_caption(model, processor, img, args.question, device, args.max_new_tokens)
        forced = f"USER: <image>\n{args.question}\nASSISTANT: {asst}"
        fwd = processor(images=[img], text=[forced], return_tensors="pt").to(device)

        sl = caption_slice(fwd["attention_mask"], asst, processor, args.max_seq_tokens)
        if sl is None:
            continue
        s0, s1 = sl

        # 1. Baseline forward: capture resid_post at all layers.
        base_acts = {}
        with t.no_grad():
            model.run_with_hooks(fwd, fwd_hooks=[(hp, make_cap(hp, base_acts)) for hp in hps])

        # Baseline toilet scores per layer (bounded σ(logit) ∈ [0,1]).
        toi_base = {}
        with t.no_grad():
            for l in layers:
                feats = base_acts[hps[l]][:, s0:s1].float()
                toi_base[l] = float(probe_score(toilet, l, feats).item())

        # 2. For each intervention layer l: gradient ascent in bathroom dir, propagate.
        #    g depends only on l (not alpha) → compute once, reuse across alphas.
        for l in layers:
            h_l = base_acts[hps[l]]                       # (1, S, d_model)
            cap = h_l[:, s0:s1].float().detach().requires_grad_(True)
            score = probe_logit(bath, l, cap)             # scalar bathroom logit
            g, = t.autograd.grad(score, cap)              # (1, T, d_model), local probe grad
            g = g.detach()
            g = g / (g.norm() + 1e-8)                      # unit ascent direction (whole-tensor norm)
            g = g.to(h_l.dtype)

            h_norm = h_l[:, s0:s1].float().norm().to(h_l.dtype)   # ‖h_l‖ over caption block; makes alpha a fraction of residual magnitude (comparable across layers)
            for a in args.alphas:
                patched = h_l.clone()
                patched[:, s0:s1] = patched[:, s0:s1] + a * h_norm * g

                # l' == l: no propagation needed. Δ = score(steered) − score(baseline).
                with t.no_grad():
                    toi_int_l = float(probe_score(toilet, l, patched[:, s0:s1].float()).item())
                delta_sum[a][l, l] += (toi_int_l - toi_base[l])

                # l' > l: patch layer l, capture downstream.
                if l < n_layers - 1:
                    down = {}
                    def patch_fn(act, hook, p=patched):
                        return p
                    hooks = [(hps[l], patch_fn)] + \
                            [(hps[lp], make_cap(hps[lp], down)) for lp in range(l + 1, n_layers)]
                    with t.no_grad():
                        model.run_with_hooks(fwd, fwd_hooks=hooks)
                    with t.no_grad():
                        for lp in range(l + 1, n_layers):
                            feats = down[hps[lp]][:, s0:s1].float()
                            toi_int = float(probe_score(toilet, lp, feats).item())
                            delta_sum[a][l, lp] += (toi_int - toi_base[lp])

        count += 1
        if (i + 1) % 10 == 0 or i == 0:
            print(f"[grad-ascent] processed {count}/{len(ids)}")

    if count == 0:
        raise SystemExit("No images processed.")
    deltas = {a: delta_sum[a] / count for a in args.alphas}
    print(f"[grad-ascent] averaged over {count} images")

    # ── Save JSON (all alphas) ─────────────────────────────────────────────────
    os.makedirs(os.path.dirname(args.out_json), exist_ok=True)
    with open(args.out_json, "w") as f:
        json.dump({"alphas": args.alphas, "n_images": count, "hook_type": args.hook_type,
                   "delta": {str(a): np.where(np.isnan(d), None, d).tolist()
                             for a, d in deltas.items()}}, f)
    print(f"[grad-ascent] saved {args.out_json}")

    # ── Heatmap (plot_mode = normal | meannorm) ─────────────────────────────────
    sel_tag = (f"{args.scene_col}=1&{args.object_col}=0" if args.hf_dataset
               else f"base_mentions={args.base_mentions}")
    plot_heatmaps(
        deltas, args.alphas, args.out,
        title=f"{args.steer_name}→{args.readout_name} causal influence ({args.plot_mode}) — "
              f"model={args.variant}, {count} imgs [{sel_tag}], {args.hook_type}",
        mode=args.plot_mode,
        cbar_label=f"Δ {args.readout_name} score σ",
        steer_name=args.steer_name, readout_name=args.readout_name)


if __name__ == "__main__":
    main()