import json, math, random, time, argparse, os, sys
from pathlib import Path
import torch, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from .model import EXLLM, EXLLMConfig
from .tokenizer import HybridTokenizer

class J(Dataset):
    def __init__(self,path,tok,maxlen):
        self.rows=[]; self.loss_tokens=0; self.processed_tokens=0
        for line in open(path,encoding='utf-8'):
            r=json.loads(line); seq=tok.encode_example(r['prompt'],r['answer'],maxlen)
            k=seq.index(tok.ASSIST); x=seq[:-1]; y=seq[1:]
            y=[-100 if i<k else z for i,z in enumerate(y)]
            self.rows.append((x,y)); self.processed_tokens+=len(x); self.loss_tokens+=sum(v!=-100 for v in y)
    def __len__(self): return len(self.rows)
    def __getitem__(self,i): return self.rows[i]

def collate(batch,pad):
    mx=max(len(x) for x,_ in batch)
    X=torch.full((len(batch),mx),pad,dtype=torch.long); Y=torch.full((len(batch),mx),-100,dtype=torch.long)
    for i,(x,y) in enumerate(batch):
        X[i,:len(x)]=torch.tensor(x); Y[i,:len(y)]=torch.tensor(y)
    return X,Y

@torch.inference_mode()
def ev(m,l):
    m.eval(); s=n=0
    for x,y in l:
        with torch.autocast('cpu',dtype=torch.bfloat16):
            z=m(x); loss=F.cross_entropy(z.reshape(-1,z.size(-1)),y.reshape(-1),ignore_index=-100,reduction='sum')
        s+=float(loss); n+=(y!=-100).sum().item()
    return s/max(n,1)

def run_stage(m,tok,train_path,valid_path,epochs,batch,lr,threads,stage,history):
    tr=J(train_path,tok,m.cfg.max_seq_len); va=J(valid_path,tok,m.cfg.max_seq_len)
    g=torch.Generator(); g.manual_seed(20260929+stage)
    tl=DataLoader(tr,batch_size=batch,shuffle=True,generator=g,collate_fn=lambda b:collate(b,tok.PAD),num_workers=0)
    vl=DataLoader(va,batch_size=batch,shuffle=False,collate_fn=lambda b:collate(b,tok.PAD),num_workers=0)
    opt=torch.optim.AdamW(m.parameters(),lr=lr,betas=(.9,.95),weight_decay=.04,foreach=True)
    total=epochs*len(tl); warm=max(10,int(total*.03)); step=0; t0=time.time(); best=1e9
    for ep in range(1,epochs+1):
        m.train(); ss=nn=0
        for x,y in tl:
            step+=1; q=0 if step<=warm else (step-warm)/max(1,total-warm)
            scale=step/warm if step<=warm else .15+.85*.5*(1+math.cos(math.pi*q))
            for g0 in opt.param_groups:g0['lr']=lr*scale
            opt.zero_grad(set_to_none=True)
            with torch.autocast('cpu',dtype=torch.bfloat16):
                z=m(x); loss=F.cross_entropy(z.reshape(-1,z.size(-1)),y.reshape(-1),ignore_index=-100)
            loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(),1.0); opt.step()
            n=(y!=-100).sum().item(); ss+=float(loss)*n; nn+=n
            if step%50==0:
                print(json.dumps({'stage':stage,'epoch':ep,'step':step,'loss':round(float(loss),5),'lr':g0['lr'],'sec':round(time.time()-t0,1)}),flush=True)
        v=ev(m,vl)
        rec={'stage':stage,'epoch':ep,'train_loss':ss/max(nn,1),'valid_loss':v,'sec':time.time()-t0,'examples':len(tr),'processed_tokens_per_epoch':tr.processed_tokens,'loss_tokens_per_epoch':tr.loss_tokens}
        history.append(rec); print(json.dumps(rec),flush=True)
        if v<best:
            best=v
            torch.save({'model':m.state_dict(),'config':m.cfg.__dict__,'meta':{'name':'EXLLM-ONI5M','version':'1.0.0','origin':'random-initialization','params':m.num_parameters(),'tokenizer':'EXLLM v1.1 fixed hybrid-char-byte-fallback','data':'project-generated-only'},'history':history},f'weights/oni5m-stage{stage}-best.pt')
    return tr.processed_tokens*epochs, tr.loss_tokens*epochs

def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--threads',type=int,default=5); ap.add_argument('--batch',type=int,default=64); a=ap.parse_args()
    torch.manual_seed(20260929); random.seed(20260929); torch.set_num_threads(a.threads)
    tok=HybridTokenizer.load('tokenizer.json'); cfg=EXLLMConfig(); assert tok.vocab_size==cfg.vocab_size==868
    m=EXLLM(cfg); history=[]; total_proc=total_loss=0
    plan=[
      (1,'data/stage1.jsonl',3,8e-4),
      (2,'data/train_balanced.jsonl',7,4e-4),
      (3,'data/stage3.jsonl',3,1.5e-4),
    ]
    for stage,path,epochs,lr in plan:
        p,l=run_stage(m,tok,path,'data/valid.jsonl',epochs,a.batch,lr,a.threads,stage,history); total_proc+=p; total_loss+=l
    final={'model':m.state_dict(),'config':cfg.__dict__,'meta':{'name':'EXLLM-ONI5M','version':'1.0.0','origin':'random-initialization','params':m.num_parameters(),'tokenizer':'EXLLM v1.1 fixed hybrid-char-byte-fallback','data':'project-generated-only','seed':20260929,'processed_tokens':total_proc,'loss_bearing_tokens':total_loss},'history':history}
    torch.save(final,'weights/EXLLM-ONI5M-v1.0.pt')
    Path('training/train_history.json').write_text(json.dumps(history,ensure_ascii=False,indent=2),encoding='utf-8')
    print(json.dumps({'final':'weights/EXLLM-ONI5M-v1.0.pt','params':m.num_parameters(),'processed_tokens':total_proc,'loss_bearing_tokens':total_loss}),flush=True)
if __name__=='__main__': main()