import argparse, json, os, re, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from tqdm import tqdm
from configs import get_config
from configs.paths import LOG_DIR, dim_paths, ensure_dirs
from src.detectors import BehaviorDetector
from src.interventions import generate_with_alpha
from src.utils import build_chat_prompt, get_device, load_model_and_tokenizer, read_json, read_jsonl, setup_logger, write_json


def extract_boxed_letter(text):
    if not text:
        return None
    matches, idx = [], 0
    while True:
        i = text.find("\\boxed", idx)
        if i < 0:
            break
        j = text.find("{", i)
        if j < 0:
            break
        depth, end = 0, -1
        for k in range(j, len(text)):
            if text[k] == "{":
                depth += 1
            elif text[k] == "}":
                depth -= 1
                if depth == 0:
                    end = k
                    break
        if end > j:
            matches.append(text[j + 1:end].strip())
            idx = end + 1
        else:
            break
    if not matches:
        return None
    last = matches[-1].strip().upper()
    m = re.match(r"\(?\s*([ABCD])", last)
    return m.group(1) if m else None


def repetition_score(text, tail_chars=400, ngram=30):
    tail = text[-tail_chars:] if len(text) > tail_chars else text
    if len(tail) < ngram * 2:
        return 0.0
    seen, repeated, total = {}, 0, 0
    for i in range(len(tail) - ngram):
        chunk = tail[i:i + ngram]
        total += 1
        if chunk in seen:
            repeated += 1
        else:
            seen[chunk] = 1
    return repeated / total if total else 0.0


def think_tokens(tok, cot):
    seg = cot.split("</think>")[0] if "</think>" in cot else cot
    return len(tok(seg, add_special_tokens=False)["input_ids"])


def load_alpha1_baseline(results_dir, seed):
    candidates = [
        os.path.join(results_dir, f"run_gpqa_d_gpqa_d_s{seed}.jsonl"),
        os.path.join(results_dir, "run_gpqa_d_gpqa_d_s64.jsonl"),
    ]
    baseline = {}
    for path in candidates:
        if not os.path.exists(path):
            continue
        for line in open(path, encoding="utf-8"):
            line = line.strip()
            if not line:
                continue
            try:
                r = json.loads(line)
            except Exception:
                continue
            if abs(float(r.get("alpha", -999)) - 1.0) < 1e-6:
                baseline[int(r["problem_idx"])] = r
        if baseline:
            return baseline, path
    return baseline, None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dimension", default="monitoring")
    ap.add_argument("--alpha", type=float, default=0.7)
    ap.add_argument("--seed", type=int, default=64)
    ap.add_argument("--sel-suffix", default="_allmonoV2")
    ap.add_argument("--gen-max-tokens", type=int, default=None)
    ap.add_argument("--data-path", default=None)
    ap.add_argument("--out-suffix", default="_alpha07_s64")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    ensure_dirs(args.dimension)
    cfg = get_config(args.dimension)
    p = dim_paths(args.dimension)

    gen_max = args.gen_max_tokens or cfg.GEN_MAX_NEW_TOKENS
    temperature = getattr(cfg, "DEFAULT_TEMPERATURE", 0.6)
    top_p = getattr(cfg, "DEFAULT_TOP_P", 0.95)

    data_path = args.data_path or os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        "data", "gpqa_d.jsonl"
    )

    log = setup_logger(
        "gpqa_layer_ablation_alpha07",
        os.path.join(LOG_DIR, f"run_gpqa_d_layer_ablation{args.out_suffix}.log"),
    )

    log.info("=" * 72)
    log.info("GPQA-D layer ablation: selected vs all candidate layers")
    log.info(f"alpha={args.alpha} seed={args.seed} gen_max={gen_max} temp={temperature} top_p={top_p}")
    log.info(f"data_path={data_path}")
    log.info("=" * 72)

    if not os.path.exists(p.DIRECTIONS):
        log.error(f"missing {p.DIRECTIONS}")
        sys.exit(1)

    dblob = torch.load(p.DIRECTIONS, map_location="cpu", weights_only=False)
    directions_all = {int(L): v for L, v in dblob["directions"].items()}
    all_layers = sorted(directions_all.keys())

    base, ext = os.path.splitext(p.SELECTED_LAYERS)
    sel_path = f"{base}{args.sel_suffix}{ext}"
    if not os.path.exists(sel_path):
        log.error(f"missing selected layer file: {sel_path}")
        sys.exit(2)

    selected_json = read_json(sel_path)
    selected_layers = [int(L) for L in selected_json["selected_layers"] if int(L) in directions_all]

    all_candidate_path = f"{base}_allcandidate{ext}"
    write_json(
        {
            "dimension": args.dimension,
            "selected_layers": all_layers,
            "n_selected": len(all_layers),
            "policy": "all_candidate_layers_no_monotonic_filter_for_gpqa_ablation",
            "note": "Auto-written by run_gpqa_d_layer_ablation_alpha07.py; uses every layer with a learned direction.",
        },
        all_candidate_path,
    )

    conditions = [
        {"name": "selected_alpha07", "layers": selected_layers},
        {"name": "allcandidate_alpha07", "layers": all_layers},
    ]

    for cond in conditions:
        log.info(f"{cond['name']} layers({len(cond['layers'])})={cond['layers']}")

    out_path = os.path.join(p.RESULTS_DIR, f"run_gpqa_d_layer_ablation{args.out_suffix}.jsonl")
    sum_path = os.path.join(p.RESULTS_DIR, f"run_gpqa_d_layer_ablation{args.out_suffix}_summary.json")

    items = read_jsonl(data_path)
    gt = {i: it["answer"].strip().upper() for i, it in enumerate(items)}

    baseline, baseline_path = load_alpha1_baseline(p.RESULTS_DIR, args.seed)
    log.info(f"alpha=1 baseline records for token-rise comparison: {len(baseline)} from {baseline_path}")

    if args.force and os.path.exists(out_path):
        os.remove(out_path)

    seen = set()
    if os.path.exists(out_path):
        for line in open(out_path, encoding="utf-8"):
            line = line.strip()
            if line:
                try:
                    seen.add(json.loads(line)["_key"])
                except Exception:
                    pass
        log.info(f"[resume] {len(seen)} records cached")

    todo = []
    for pi, it in enumerate(items):
        for cond in conditions:
            key = f"P{pi}_{cond['name']}"
            if key not in seen:
                todo.append((pi, it["problem"], cond, key))

    log.info(f"records to compute: {len(todo)} / {len(items) * len(conditions)}")

    detector = BehaviorDetector(cfg)
    device = get_device()

    model = tokenizer = None
    if todo:
        log.info("Loading model...")
        model, tokenizer = load_model_and_tokenizer(device=device)

    fh = open(out_path, "a", encoding="utf-8") if todo else None

    for pi, problem, cond, key in tqdm(todo, desc="gpqa_layer_ablation", dynamic_ncols=True, mininterval=10):
        prompt = build_chat_prompt(tokenizer, problem, enable_thinking=True, system="")
        cond_dirs = {L: directions_all[L] for L in cond["layers"]}
        eff = {L: float(args.alpha) for L in cond_dirs}
        gen_seed = args.seed * 1000 + pi

        t0 = time.time()
        cot = generate_with_alpha(
            model,
            tokenizer,
            prompt,
            cond_dirs,
            eff,
            device,
            max_new_tokens=gen_max,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            seed=gen_seed,
        )
        elapsed = time.time() - t0

        pred = extract_boxed_letter(cot)
        det = detector.detect(cot)
        rep = repetition_score(cot)
        ttok = think_tokens(tokenizer, cot)

        b = baseline.get(pi)
        b_tok = b.get("think_tokens") if b else None
        delta = (ttok - b_tok) if b_tok is not None else None

        rec = {
            "_key": key,
            "problem_idx": pi,
            "seed": args.seed,
            "alpha": float(args.alpha),
            "condition": cond["name"],
            "n_layers": len(cond["layers"]),
            "layers": cond["layers"],
            "problem": problem,
            "cot": cot,
            "pred": pred,
            "gt": gt.get(pi),
            "correct": (pred == gt.get(pi)) if pred and gt.get(pi) else False,
            "has_boxed": pred is not None,
            "has_think_end": "</think>" in cot,
            "think_tokens": ttok,
            "n_chars": len(cot),
            "mon_total": det["total"],
            "repetition_score": rep,
            "collapse": rep > 0.5,
            "baseline_think_tokens": b_tok,
            "delta_think_tokens_vs_alpha1": delta,
            "token_rise_vs_alpha1": (delta is not None and delta > 0),
            "elapsed_s": elapsed,
        }

        fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
        fh.flush()

        log.info(
            f"{key}: pred={pred} gt={gt.get(pi)} {'OK' if rec['correct'] else 'x'} "
            f"tok={ttok} mon={det['total']} rep={rep:.2f} delta_vs_a1={delta} t={elapsed:.0f}s"
        )

    if fh:
        fh.close()

    records = []
    if os.path.exists(out_path):
        for line in open(out_path, encoding="utf-8"):
            line = line.strip()
            if line:
                try:
                    records.append(json.loads(line))
                except Exception:
                    pass

    avg = lambda xs: sum(xs) / len(xs) if xs else 0.0
    summary = {}

    log.info("\n=== SUMMARY: GPQA-D layer ablation alpha=0.7 ===")
    log.info(f"baseline_for_token_rise={baseline_path}")
    log.info(
        f"{'condition':>18} {'layers':>6} {'n':>3} {'acc':>8} {'correct':>8} {'noBox':>6} "
        f"{'think_tok':>10} {'mon':>7} {'collapse':>9} {'rise_vs_a1':>10} {'delta_tok':>10}"
    )

    for cond in conditions:
        name = cond["name"]
        rs = [r for r in records if r.get("condition") == name]
        if not rs:
            continue

        n = len(rs)
        n_correct = sum(bool(r["correct"]) for r in rs)
        no_box = n - sum(bool(r["has_boxed"]) for r in rs)
        collapse_rate = sum(bool(r["collapse"]) for r in rs) / n

        comparable = [r for r in rs if r.get("baseline_think_tokens") is not None]
        rise_rate = (
            sum(bool(r.get("token_rise_vs_alpha1")) for r in comparable) / len(comparable)
            if comparable else None
        )
        mean_delta = (
            avg([r["delta_think_tokens_vs_alpha1"] for r in comparable])
            if comparable else None
        )

        summary[name] = {
            "n": n,
            "accuracy": n_correct / n,
            "n_correct": n_correct,
            "n_no_boxed": no_box,
            "mean_think_tokens": avg([r["think_tokens"] for r in rs]),
            "mean_chars": avg([r["n_chars"] for r in rs]),
            "mean_mon": avg([r["mon_total"] for r in rs]),
            "mean_repetition_score": avg([r["repetition_score"] for r in rs]),
            "collapse_rate": collapse_rate,
            "token_rise_vs_alpha1_rate": rise_rate,
            "mean_delta_think_tokens_vs_alpha1": mean_delta,
            "n_layers": len(cond["layers"]),
            "layers": cond["layers"],
        }

        rise_s = "NA" if rise_rate is None else f"{rise_rate*100:.1f}%"
        delta_s = "NA" if mean_delta is None else f"{mean_delta:.0f}"

        log.info(
            f"{name:>18} {len(cond['layers']):>6} {n:>3} {n_correct/n:>7.1%} {n_correct:>8} {no_box:>6} "
            f"{avg([r['think_tokens'] for r in rs]):>10.0f} "
            f"{avg([r['mon_total'] for r in rs]):>7.1f} "
            f"{collapse_rate*100:>8.1f}% {rise_s:>10} {delta_s:>10}"
        )

    write_json(
        {
            "dataset": "GPQA-Diamond",
            "seed": args.seed,
            "alpha": float(args.alpha),
            "temperature": temperature,
            "top_p": top_p,
            "gen_max": gen_max,
            "selected_layer_file": sel_path,
            "all_candidate_layer_file": all_candidate_path,
            "baseline_alpha1_file_for_token_rise": baseline_path,
            "conditions": conditions,
            "summary": summary,
        },
        sum_path,
    )

    log.info(f"Saved {out_path}\n      {sum_path}\nDone.")


if __name__ == "__main__":
    main()