"""Single-head ablation sweep over the 144 (block, head) pairs.

For each head, zero the columns of its block's attention output projection,
score the calibration pool with the Stage 0 classifier, and record the F1 delta
and the L2 deviation of the 40 classifier-relevant output dims. Heads are then
ranked by ascending F1 drop, so the head whose removal helps most sorts first,
and the cumulative curve prunes that prefix.

Writes head_importance.json and pruning_curve.json.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # repo root, for `common`
from common import (BACKBONE, N_BLOCKS, N_HEADS, device, f1_at, heads_masked,  # noqa: E402
                    load_pool, score_pool, write_artifact)
from common.models import load_backbone  # noqa: E402
from common.pools import CALIB1000, by_name  # noqa: E402

HERE = Path(__file__).resolve().parent
CLASSIFIER = HERE.parent / 'stage_0' / 'classifier.json'
CURVE_K = [1, 5, 10, 15, 20, 30, 40, 50, 60, 80, 100, 120, 144]
RANKING = 'ascending F1_drop; the smallest drop is the most prunable head'


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--backbone', default=BACKBONE)
    ap.add_argument('--pool', default=CALIB1000.name)
    args = ap.parse_args()

    dev = device()
    c = json.loads(CLASSIFIER.read_text())
    pos = torch.tensor(c['pos_dims'], dtype=torch.long, device=dev)
    neg = torch.tensor(c['neg_dims'], dtype=torch.long, device=dev)
    target_dims = torch.cat([pos, neg]).unique()
    thr = float(c['threshold'])
    print(f'[init] |target_dims|={len(target_dims)}  threshold={thr:.3f}', flush=True)

    backbone = load_backbone(args.backbone).to(dev)
    pool = by_name(args.pool)
    print(f'[pool] {pool.name}, preloading', flush=True)
    loaded = load_pool(pool, dev, preload=True)
    print(f'  person rate {loaded.positive_rate:.3f}', flush=True)

    base_scores, base_targets = score_pool(backbone, loaded, pos, neg, target_dims)
    base = f1_at(base_scores, loaded.labels, thr)
    print(f'[baseline] F1={base.f1:.4f}  P={base.precision:.4f}  R={base.recall:.4f}',
          flush=True)

    results = []
    t0 = time.time()
    for b in range(N_BLOCKS):
        for h in range(N_HEADS):
            with heads_masked(backbone, [(b, h)]):
                scores, targets = score_pool(backbone, loaded, pos, neg, target_dims)
            m = f1_at(scores, loaded.labels, thr)
            l2 = (targets - base_targets).pow(2).sum(dim=1).sqrt().mean().item()
            results.append({'block': b, 'head': h, 'F1': m.f1, 'precision': m.precision,
                            'recall': m.recall, 'F1_drop': base.f1 - m.f1, 'target_L2': l2})
            print(f'  B{b:>2}H{h:>2}  F1={m.f1:.4f}  drop={base.f1 - m.f1:+.4f}  '
                  f'L2={l2:.3f}  {time.time() - t0:.1f}s', flush=True)

    ranked = sorted(results, key=lambda r: r['F1_drop'])

    curve = []
    for k in CURVE_K:
        with heads_masked(backbone, [(r['block'], r['head']) for r in ranked[:k]]):
            scores, _ = score_pool(backbone, loaded, pos, neg)
        m = f1_at(scores, loaded.labels, thr)
        curve.append({'heads_pruned': k, 'F1': m.f1, 'F1_drop': base.f1 - m.f1,
                      'precision': m.precision, 'recall': m.recall})
        print(f'  K={k:>3}  F1={m.f1:.4f}  drop={base.f1 - m.f1:+.4f}', flush=True)

    stamp = dict(generator='stage_2/ablate.py', classifier=CLASSIFIER,
                 pool_info=loaded.provenance())
    write_artifact(HERE / 'head_importance.json', {
        'ranking': RANKING,
        'baseline': base.asdict(),
        'per_head': results,
        'ranked_most_prunable_first': [(r['block'], r['head'], r['F1_drop'])
                                       for r in ranked],
    }, **stamp)
    write_artifact(HERE / 'pruning_curve.json',
                   {'baseline': base.asdict(), 'curve': curve}, **stamp)
    print('[done]', flush=True)


if __name__ == '__main__':
    main()