"""Full supervised training using the pinned HF token-classification recipe."""
import argparse
import json
import math
import random
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification
from pii.core import BIO, MODEL_REVISION, load_jsonl, encode_windows

def prepare(records, tokenizer, max_length, stride):
    chunks=[]
    stats={"documents":len(records),"chunks":0,"tokens":0,"ignored_boundary_tokens":0,"partial_window_entity_tokens":0}
    for record in records:
        encoded=encode_windows(tokenizer,record["text"],max_length,stride)
        for ids, mask, offsets in zip(encoded["input_ids"],encoded["attention_mask"],encoded["offset_mapping"]):
            content=[(a,b) for a,b in offsets if a!=b]
            left,right=content[0][0] if content else 0,content[-1][1] if content else 0
            labels=[]
            for a,b in offsets:
                label=0
                if a==b:
                    labels.append(-100)
                    continue
                for span in record["spans"]:
                    s,e=span["start"],span["end"]
                    if a<e and s<b:
                        if not (s<=a and b<=e):
                            label=-100
                            stats["ignored_boundary_tokens"]+=1
                        elif s<left or e>right:
                            label=-100
                            stats["partial_window_entity_tokens"]+=1
                        else:
                            # First representable token of the entity gets B.
                            earlier=any(s<=aa<a and bb<=e and aa!=bb for aa,bb in offsets)
                            label=BIO.index(("I-" if earlier else "B-")+span["label"])
                        break
                labels.append(label)
            chunks.append({"input_ids":ids,"attention_mask":mask,"labels":labels})
            stats["tokens"]+=len(ids)
    stats["chunks"]=len(chunks)
    return chunks,stats

def main():
    p=argparse.ArgumentParser()
    p.add_argument("--stride",type=int,required=True)
    p.add_argument("--device",default="mps")
    p.add_argument("--resume",type=Path)
    p.add_argument("--prepare-only",action="store_true")
    args=p.parse_args()
    out=Path("artifacts/training");out.mkdir(parents=True,exist_ok=True)
    seed=42
    random.seed(seed);np.random.seed(seed);torch.manual_seed(seed)
    tokenizer=AutoTokenizer.from_pretrained("artifacts/backbone",local_files_only=True)
    config={"seed":seed,"batch_size":16,"epochs":2,"learning_rate":2e-5,"weight_decay":0.01,"max_length":512,"stride":args.stride,"device":args.device,"precision":"float32","model_revision":MODEL_REVISION,"recipe":"https://huggingface.co/docs/transformers/tasks/token_classification","stopping_rule":"Two epochs from official recipe; select lowest validation token loss. This is a prototype recipe hypothesis, not an optimum.","checkpoint_cadence":"Each epoch per official recipe; measured probe checkpoint cost 0.733s.","selection_metric":"validation token cross entropy"}
    cache=out/f"chunks-manual-stride-{args.stride}.pt"
    if cache.exists():
        prepared=torch.load(cache,weights_only=False)
    else:
        prepared={}
        for name in ("train","validation"):
            prepared[name],prepared[name+"_stats"]=prepare(load_jsonl(f"data/prepared/{name}.jsonl"),tokenizer,512,args.stride)
            print(json.dumps({name:prepared[name+"_stats"]}),flush=True)
        torch.save(prepared,cache)
    config["data_stats"]={name:prepared[name+"_stats"] for name in ("train","validation")}
    config["estimated_training_seconds_max_length_probe"]=math.ceil(len(prepared["train"])/16)*2*0.558059167
    (out/"run-config.json").write_text(json.dumps(config,indent=2)+"\n")
    if args.prepare_only:
        print(json.dumps(config,indent=2),flush=True);return
    if config["estimated_training_seconds_max_length_probe"]>24*3600:
        raise RuntimeError("Estimated run exceeds user-approved 24h limit; owner decision required")
    model=AutoModelForTokenClassification.from_pretrained(args.resume / "model" if args.resume else "artifacts/backbone",num_labels=len(BIO),id2label=dict(enumerate(BIO)),label2id={v:k for k,v in enumerate(BIO)},local_files_only=True,attn_implementation="eager").to(args.device)
    optimizer=torch.optim.AdamW(model.parameters(),lr=config["learning_rate"],weight_decay=config["weight_decay"])
    steps_per_epoch=math.ceil(len(prepared["train"])/16)
    scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda step:max(0.0,1-step/(steps_per_epoch*config["epochs"])))
    collate=DataCollatorForTokenClassification(tokenizer,padding=True,return_tensors="pt")
    start_epoch,best_loss,global_step=0,float("inf"),0
    if args.resume:
        state=torch.load(args.resume/"state.pt",map_location="cpu",weights_only=False)
        optimizer.load_state_dict(state["optimizer"]);scheduler.load_state_dict(state["scheduler"])
        random.setstate(state["random"]);np.random.set_state(state["numpy"]);torch.set_rng_state(state["torch"])
        if args.device=="mps":torch.mps.set_rng_state(state["mps"])
        start_epoch,best_loss,global_step=state["epoch"],state["best_loss"],state["step"]
    history=[]
    started=time.perf_counter()
    for epoch in range(start_epoch,config["epochs"]):
        model.train()
        order=list(range(len(prepared["train"])));random.shuffle(order)
        running=0.0
        for step,begin in enumerate(range(0,len(order),16)):
            batch=collate([prepared["train"][i] for i in order[begin:begin+16]])
            batch={k:v.to(args.device) for k,v in batch.items()}
            optimizer.zero_grad(set_to_none=True)
            loss=model(**batch).loss
            if not torch.isfinite(loss):raise RuntimeError("nonfinite loss")
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(),1.0,error_if_nonfinite=True) # Trainer default.
            optimizer.step();scheduler.step()
            running+=loss.item();global_step+=1
            # Persist current progress every batch; no logging interval assumption.
            (out/"progress.json").write_text(json.dumps({"epoch":epoch+1,"step_in_epoch":step+1,"steps_per_epoch":steps_per_epoch,"global_step":global_step,"loss":loss.item(),"elapsed_seconds":time.perf_counter()-started})+"\n")
        model.eval();total_loss,total_tokens=0.0,0
        with torch.no_grad():
            for begin in range(0,len(prepared["validation"]),16):
                batch=collate(prepared["validation"][begin:begin+16]);batch={k:v.to(args.device) for k,v in batch.items()}
                count=(batch["labels"]!=-100).sum().item()
                total_loss+=model(**batch).loss.item()*count;total_tokens+=count
        val_loss=total_loss/total_tokens
        selected=val_loss<best_loss
        if selected:
            best_loss=val_loss
            model.save_pretrained(out/"best");tokenizer.save_pretrained(out/"best")
        checkpoint=out/f"epoch-{epoch+1}"
        model.save_pretrained(checkpoint/"model")
        torch.save({"optimizer":optimizer.state_dict(),"scheduler":scheduler.state_dict(),"random":random.getstate(),"numpy":np.random.get_state(),"torch":torch.get_rng_state(),"mps":torch.mps.get_rng_state() if args.device=="mps" else None,"epoch":epoch+1,"best_loss":best_loss,"step":global_step},checkpoint/"state.pt")
        event={"epoch":epoch+1,"train_mean_batch_loss":running/steps_per_epoch,"validation_token_loss":val_loss,"selected":selected,"elapsed_seconds":time.perf_counter()-started}
        history.append(event)
        with (out/"history.jsonl").open("a") as handle:handle.write(json.dumps(event)+"\n")
        print(json.dumps(event),flush=True)
    (out/"complete.json").write_text(json.dumps({"history":history,"best_validation_loss":best_loss,"elapsed_seconds":time.perf_counter()-started},indent=2)+"\n")

if __name__=="__main__":main()