"""Optional GPT-2 hook example. Syntax checked only; local weights required."""
import argparse
import json
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--model-dir',type=Path,required=True)
    ap.add_argument('--clean',required=True)
    ap.add_argument('--corrupt',required=True)
    ap.add_argument('--answer',required=True)
    ap.add_argument('--alternative',required=True)
    ap.add_argument('--module',default='transformer.h.5')
    ap.add_argument('--position',type=int,required=True)
    ap.add_argument('--output',type=Path,default=Path('local_patch_results.json'))
    args=ap.parse_args();torch.manual_seed(76)
    tok=AutoTokenizer.from_pretrained(args.model_dir,local_files_only=True)
    model=AutoModelForCausalLM.from_pretrained(args.model_dir,local_files_only=True).to('cpu').eval()
    if model.config.model_type!='gpt2':raise ValueError('This example supports GPT-2 blocks only.')
    module=model.get_submodule(args.module)
    a=tok(args.clean,return_tensors='pt');b=tok(args.corrupt,return_tensors='pt')
    print('clean tokens:',list(enumerate(tok.convert_ids_to_tokens(a['input_ids'][0].tolist()))))
    print('corrupt tokens:',list(enumerate(tok.convert_ids_to_tokens(b['input_ids'][0].tolist()))))
    if a['input_ids'].shape!=b['input_ids'].shape:
        raise ValueError('Unequal token lengths: manually align tokens or choose another pair.')
    ids=[tok.encode(s,add_special_tokens=False) for s in [args.answer,args.alternative]]
    if any(len(v)!=1 for v in ids):raise ValueError('Both answer strings must be single tokens.')
    if ids[0]==ids[1]:raise ValueError('Answer and alternative must differ.')
    pos=args.position
    if not 0<=pos<a['input_ids'].shape[1]:raise ValueError('Invalid token position.')
    cache={}
    def hidden(output):return output[0] if isinstance(output,tuple) else output
    def capture(_,inputs,output):cache['clean']=hidden(output).detach().clone()
    with torch.no_grad():
        handle=module.register_forward_hook(capture)
        try:zc=model(**a).logits
        finally:handle.remove()
        zb=model(**b).logits
        def run_patch(full=False):
            def patch(_,inputs,output):
                h=hidden(output).clone()
                if full:h[:]=cache['clean']
                else:h[:,pos,:]=cache['clean'][:,pos,:]
                return (h,)+output[1:] if isinstance(output,tuple) else h
            handle=module.register_forward_hook(patch)
            try:return model(**b).logits
            finally:handle.remove()
        zp=run_patch();zf=run_patch(full=True);za=model(**b).logits
    assert torch.allclose(zb,za,atol=1e-5),'Hook leaked into later run'
    assert torch.allclose(zf,zc,atol=1e-4),'Full block restoration failed'
    def score(z):return float(z[0,-1,ids[0][0]]-z[0,-1,ids[1][0]])
    cl,bl,pl=map(score,[zc,zb,zp]);denom=cl-bl
    out=dict(clean=args.clean,corrupt=args.corrupt,module=args.module,position=pos,
             clean_tokens=a['input_ids'].tolist(),corrupt_tokens=b['input_ids'].tolist(),
             cache_shape=list(cache['clean'].shape),logit_shape=list(zc.shape),
             clean_ld=cl,corrupt_ld=bl,patched_ld=pl,delta_ld=pl-bl,
             recovery=(pl-bl)/denom if denom>1e-6 else None,
             torch_version=torch.__version__,model_path=str(args.model_dir.resolve()),
             note='Single paired prompt diagnostic; inspect baselines before interpretation.')
    args.output.write_text(json.dumps(out,indent=2)+'\n');print(json.dumps(out,indent=2))


if __name__=='__main__':main()