#!/usr/bin/env python3
"""Build a DS4 directional-steering vector from paired prompt sets.

The extractor asks ds4 to dump one 4096-wide activation row per layer, averages
the target and control rows, and writes a flat f32 file with 43 layer vectors.
At runtime ds4 applies:

    y = y - scale * direction[layer] * dot(direction[layer], y)

Positive scale suppresses the target direction.  Negative scale amplifies it.
"""

import argparse
import array
import json
import math
import os
import subprocess
import tempfile
from pathlib import Path


N_LAYER = 43
N_EMBD = 4096

SPECIALS = {
    "bos": "<｜begin▁of▁sentence｜>",
    "user": "<｜User｜>",
    "assistant": "<｜Assistant｜>",
    "think": "<think>",
    "nothink": "</think>",
}


def read_prompt_file(path: Path) -> list[str]:
    """Read one prompt per non-empty line, ignoring shell-style comments."""
    prompts: list[str] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        prompts.append(line)
    if not prompts:
        raise SystemExit(f"{path}: no prompts found")
    return prompts


def render_ds4_prompt(system: str, user: str, think: bool) -> str:
    """Render the minimal DS4 chat prefix used for activation capture."""
    pieces = [SPECIALS["bos"]]
    if system:
        pieces.append(system)
    pieces += [
        SPECIALS["user"],
        user,
        SPECIALS["assistant"],
        SPECIALS["think"] if think else SPECIALS["nothink"],
    ]
    return "".join(pieces)


def normalize(v: list[float]) -> list[float]:
    n2 = sum(x * x for x in v)
    if n2 <= 0.0:
        return v
    inv = 1.0 / math.sqrt(n2)
    return [x * inv for x in v]


def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def run_capture(
    ds4: Path,
    model: Path,
    prompt: str,
    system: str,
    think: bool,
    ctx: int,
    component: str,
    work: Path,
) -> list[list[float]]:
    """Run ds4 once and return the last prompt-row dump for every layer."""
    prompt_path = work / "prompt.txt"
    prompt_path.write_text(render_ds4_prompt(system, prompt, think), encoding="utf-8")
    dump_prefix = work / "dump"

    env = os.environ.copy()
    env["DS4_METAL_GRAPH_DUMP_PREFIX"] = str(dump_prefix)
    env["DS4_METAL_GRAPH_DUMP_NAME"] = component
    env["DS4_METAL_GRAPH_DUMP_POS"] = "0"

    cmd = [
        str(ds4),
        "-m", str(model),
        "--ctx", str(ctx),
        "--prompt-file", str(prompt_path),
        "-n", "1",
    ]
    subprocess.run(cmd, cwd=ds4.parent, env=env, check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)

    rows: list[list[float]] = []
    for layer in range(N_LAYER):
        path = work / f"dump_{component}-{layer}_pos0.bin"
        data = array.array("f")
        with path.open("rb") as f:
            data.fromfile(f, path.stat().st_size // 4)
        if len(data) < N_EMBD or len(data) % N_EMBD != 0:
            raise RuntimeError(f"bad dump shape for {path}: {len(data)} floats")
        rows.append(list(data[-N_EMBD:]))
    return rows


def add_rows(total: list[list[float]], rows: list[list[float]]) -> None:
    for layer in range(N_LAYER):
        dst = total[layer]
        src = rows[layer]
        for i, value in enumerate(src):
            dst[i] += value


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ds4", default="./ds4", help="path to the ds4 CLI")
    ap.add_argument("--model", default="ds4flash.gguf", help="GGUF model path")
    ap.add_argument("--good-file", required=True,
                    help="desired/target prompts, one per line")
    ap.add_argument("--bad-file", required=True,
                    help="contrast/control prompts, one per line")
    ap.add_argument("--out", default="dir-steering/out/direction.json",
                    help="metadata JSON path; .f32 is written next to it")
    ap.add_argument("--ctx", type=int, default=512)
    ap.add_argument("--system", default="You are a helpful assistant.")
    ap.add_argument("--component", default="ffn_out",
                    choices=("ffn_out", "attn_out"),
                    help="runtime-editable 4096-wide activation stream")
    ap.add_argument("--think", action="store_true",
                    help="capture after <think>; default captures direct answers")
    ap.add_argument("--pair-normalize", action="store_true",
                    help="average normalized per-pair differences")
    ap.add_argument("--no-orthogonalize", action="store_true",
                    help="do not remove the component parallel to the control mean")
    args = ap.parse_args()

    ds4 = Path(args.ds4).resolve()
    model = Path(args.model).resolve()
    good_prompts = read_prompt_file(Path(args.good_file))
    bad_prompts = read_prompt_file(Path(args.bad_file))
    n = min(len(good_prompts), len(bad_prompts))
    good_prompts = good_prompts[:n]
    bad_prompts = bad_prompts[:n]

    good_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
    bad_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
    pair_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]

    with tempfile.TemporaryDirectory(prefix="ds4-dir-steer-") as td:
        root = Path(td)
        for i, (good, bad) in enumerate(zip(good_prompts, bad_prompts), 1):
            print(f"pair {i}/{n}", flush=True)
            gw = root / f"good-{i}"
            bw = root / f"bad-{i}"
            gw.mkdir()
            bw.mkdir()
            good_rows = run_capture(ds4, model, good, args.system, args.think,
                                    args.ctx, args.component, gw)
            bad_rows = run_capture(ds4, model, bad, args.system, args.think,
                                   args.ctx, args.component, bw)
            add_rows(good_sum, good_rows)
            add_rows(bad_sum, bad_rows)
            if args.pair_normalize:
                for layer in range(N_LAYER):
                    diff = normalize([
                        good_rows[layer][j] - bad_rows[layer][j]
                        for j in range(N_EMBD)
                    ])
                    for j, value in enumerate(diff):
                        pair_sum[layer][j] += value

    layers = []
    for layer in range(N_LAYER):
        good_mean = [x / n for x in good_sum[layer]]
        bad_mean = [x / n for x in bad_sum[layer]]
        if args.pair_normalize:
            direction = normalize([x / n for x in pair_sum[layer]])
        else:
            direction = normalize([
                good_mean[i] - bad_mean[i]
                for i in range(N_EMBD)
            ])
        if not args.no_orthogonalize:
            base = normalize(bad_mean)
            projection = dot(direction, base)
            direction = normalize([
                direction[i] - projection * base[i]
                for i in range(N_EMBD)
            ])
        layers.append(direction)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "format": "ds4-directional-steering-v1",
        "shape": [N_LAYER, N_EMBD],
        "component": args.component,
        "thinking": bool(args.think),
        "pair_normalize": bool(args.pair_normalize),
        "orthogonalize_control_mean": not args.no_orthogonalize,
        "good_file": str(Path(args.good_file)),
        "bad_file": str(Path(args.bad_file)),
        "model": str(model),
        "note": "runtime positive scale suppresses this direction; negative scale amplifies it",
    }
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    flat = array.array("f")
    for direction in layers:
        flat.extend(direction)
    f32_out = out.with_suffix(".f32")
    with f32_out.open("wb") as f:
        flat.tofile(f)
    print(f"wrote {out}")
    print(f"wrote {f32_out}")


if __name__ == "__main__":
    main()