"""Local optimized full-weight RLVR baseline, adapted from rasbt/reasoning-from-scratch (Apache-2.0).
Same objective: mean group-standardized reward * summed completion log probability,
8 rollouts, temperature .8, nucleus .9, no KL, AdamW 1e-5, grad clip 1, seed 42.
Changes: SDPA, batched sampling, last-position-only generation head, per-rollout backward.
BF16 floating-point and sampling ordering differ from the reference; not bitwise reproduction.
"""
import argparse, copy, json, time, sys, types, gc
from pathlib import Path
import torch
import torch.nn.functional as F
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'vendor'))
from reasoning_from_scratch.qwen3_batched import Qwen3Model, QWEN_CONFIG_06_B, apply_rope_with_pos_ids
from reasoning_from_scratch.qwen3 import KVCache, Qwen3Tokenizer
from reasoning_from_scratch.ch03 import render_prompt, extract_final_candidate, grade_answer
from reasoning_from_scratch.ch04 import top_p_filter

def attention(self,x,mask,cos,sin,cache=None,pos_ids=None):
    b,t,_=x.shape
    q=self.W_query(x).view(b,t,self.num_heads,self.head_dim).transpose(1,2)
    k=self.W_key(x).view(b,t,self.num_kv_groups,self.head_dim).transpose(1,2)
    v=self.W_value(x).view(b,t,self.num_kv_groups,self.head_dim).transpose(1,2)
    if self.q_norm:q=self.q_norm(q)
    if self.k_norm:k=self.k_norm(k)
    q=apply_rope_with_pos_ids(q,cos,sin,pos_ids)
    k=apply_rope_with_pos_ids(k,cos,sin,pos_ids)
    if cache is not None:
        k=torch.cat((cache[0],k),2);v=torch.cat((cache[1],v),2)
    nxt=(k,v)
    k=k.repeat_interleave(self.group_size,1);v=v.repeat_interleave(self.group_size,1)
    z=F.scaled_dot_product_attention(q,k,v,attn_mask=mask,dropout_p=0.)
    return self.out_proj(z.transpose(1,2).reshape(b,t,self.d_out)),nxt

class FastModel(Qwen3Model):
    def __init__(self,cfg):
        super().__init__(cfg)
        for block in self.trf_blocks:block.att.forward=types.MethodType(attention,block.att)
    def forward(self,ids,cache=None,attn_mask=None,last_only=False):
        x=self.tok_emb(ids);b,t=ids.shape
        start=cache.get(0)[0].shape[2] if cache and cache.get(0) is not None else 0
        end=start+t
        assert end<=self.cfg['context_length']
        # Boolean SDPA mask: True means allowed. Absolute positions for left padding;
        # RoPE is invariant to a shared positional offset over each sequence.
        pos=torch.arange(start,end,device=x.device)
        allowed=(torch.arange(end,device=x.device)[None,:]<=pos[:,None])[None,None]
        if attn_mask is not None:
            allowed=allowed & attn_mask[:,None,None,:end]
            x=x*attn_mask[:,start:end,None].to(x.dtype)
        pos=pos[None,:].expand(b,-1)
        for i,block in enumerate(self.trf_blocks):
            x,nxt=block(x,allowed,self.cos,self.sin,cache=cache.get(i) if cache else None,pos_ids=pos)
            if cache is not None:cache.update(i,nxt)
        if last_only:x=x[:,-1:]
        return self.out_head(self.final_norm(x).to(self.cfg['dtype']))

@torch.no_grad()
def generate(model,tokenizer,prompts,max_tokens,sample=False):
    model.eval()
    dev=next(model.parameters()).device
    width=max(map(len,prompts));b=len(prompts)
    pad=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    ids=torch.tensor([[pad]*(width-len(p))+p for p in prompts],device=dev)
    mask=torch.tensor([[False]*(width-len(p))+[True]*len(p) for p in prompts],device=dev)
    cache=KVCache(model.cfg['n_layers'])
    logits=model(ids,cache,mask,last_only=True)[:,-1]
    done=torch.zeros(b,device=dev,dtype=torch.bool);out=[]
    for step in range(max_tokens):
        if sample:
            prob=top_p_filter(torch.softmax(logits/.8,-1),.9)
            nxt=torch.multinomial(prob,1).squeeze(1)
        else:nxt=logits.argmax(-1)
        nxt=torch.where(done,tokenizer.eos_token_id,nxt)
        out.append(nxt)
        done=done | (nxt==tokenizer.eos_token_id)
        if done.all().item() or step+1==max_tokens:break
        mask=torch.cat((mask,torch.ones((b,1),device=dev,dtype=torch.bool)),1)
        logits=model(nxt[:,None],cache,mask,last_only=True)[:,-1]
    raw=torch.stack(out,1).cpu().tolist();result=[]
    for row in raw:
        if tokenizer.eos_token_id in row:row=row[:row.index(tokenizer.eos_token_id)+1]
        result.append(row)
    return result

class EvalModel(FastModel):
    def forward(self,ids,cache=None,attn_mask=None,last_only=True):
        return super().forward(ids,cache,attn_mask,last_only=last_only)

@torch.no_grad()
def generate_eval(model,tokenizer,prompts,max_tokens):
    from reasoning_from_scratch.qwen3_batched import generate_text_basic_batched_cache_stop
    pad=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    width=max(map(len,prompts));dev=next(model.parameters()).device
    ids=torch.tensor([[pad]*(width-len(p))+p for p in prompts],device=dev)
    mask=torch.tensor([[False]*(width-len(p))+[True]*len(p) for p in prompts],device=dev)
    result=generate_text_basic_batched_cache_stop(model,ids,max_tokens,
        eos_token_id=tokenizer.eos_token_id,pad_id=pad,attn_mask=mask).cpu().tolist()
    return [row[:row.index(tokenizer.eos_token_id)+1] if tokenizer.eos_token_id in row else row for row in result]

def logprob(model,seq,prompt_len):
    logits=model(seq[None,:])[:,prompt_len-1:-1,:].float()
    # Same selected log-softmax sum as upstream, without an extra full probability tensor.
    return -F.cross_entropy(logits.reshape(-1,logits.shape[-1]),seq[prompt_len:],reduction='sum')

def memory():
    return dict(peak_allocated_gib=torch.cuda.max_memory_allocated()/2**30,
                peak_reserved_gib=torch.cuda.max_memory_reserved()/2**30)

def emit(obj):print(json.dumps(obj),flush=True)

def verify():
    torch.manual_seed(1)
    cfg=dict(QWEN_CONFIG_06_B,vocab_size=128,emb_dim=32,n_heads=4,n_kv_groups=2,head_dim=8,n_layers=2,hidden_dim=64,context_length=64,dtype=torch.float32)
    ref=Qwen3Model(cfg).cuda();fast=FastModel(cfg).cuda();fast.load_state_dict(ref.state_dict())
    ids=torch.randint(3,128,(2,9),device='cuda');mask=torch.ones_like(ids,dtype=torch.bool);mask[0,:3]=False
    with torch.no_grad():
        torch.testing.assert_close(fast(ids),ref(ids),rtol=1e-4,atol=1e-5)
        torch.testing.assert_close(fast(ids,attn_mask=mask),ref(ids,attn_mask=mask),rtol=1e-4,atol=1e-5)
        torch.testing.assert_close(fast(ids,attn_mask=mask)[0,3:],fast(ids[0:1,3:])[0],rtol=1e-4,atol=1e-5)
        ca=KVCache(2);cb=KVCache(2)
        fast(ids[:,:8],ca);ref(ids[:,:8],cb)
        torch.testing.assert_close(fast(ids[:,8:],ca),ref(ids[:,8:],cb),rtol=1e-4,atol=1e-5)
    clone=copy.deepcopy(fast)
    advantages=torch.tensor([-.7,.7],device='cuda')
    loss=sum(-advantages[i]*logprob(fast,ids[i],3)/2 for i in range(2));loss.backward()
    for i in range(2):(-advantages[i]*logprob(clone,ids[i],3)/2).backward()
    for p,q in zip(fast.parameters(),clone.parameters()):torch.testing.assert_close(p.grad,q.grad,rtol=1e-4,atol=1e-5)
    tok=types.SimpleNamespace(pad_token_id=0,eos_token_id=1)
    ev=EvalModel(cfg).cuda();ev.load_state_dict(fast.state_dict())
    prompts=[[5,7,9],[3,4],[7,8,9,10]]
    assert generate(fast,tok,prompts,12)==generate_eval(ev,tok,prompts,12)
    emit({'verification':'PASS: SDPA logits, padding, cache, sequential-gradient equivalence, compacted evaluation generation'})

def train(args):
    torch.manual_seed(args.seed)
    model=FastModel(dict(QWEN_CONFIG_06_B,context_length=4096)).cuda()
    model.load_state_dict(torch.load(ROOT/'data/qwen3-0.6B-base.pth',map_location='cpu',weights_only=True))
    tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'data/tokenizer-base.json')
    data=json.loads((ROOT/'data/math_train.json').read_text())
    if not 1 <= args.steps <= len(data):
        raise ValueError(f'--steps must be between 1 and {len(data)}')
    optimizer=torch.optim.AdamW(model.parameters(),lr=1e-5)
    out=ROOT/'results'/args.name;out.mkdir(parents=True,exist_ok=False)
    (out/'config.json').write_text(json.dumps(vars(args),indent=2))
    torch.cuda.reset_peak_memory_stats();start=time.perf_counter()
    with (out/'training.jsonl').open('w') as f,(out/'rollouts.jsonl').open('w') as samples:
        for step,row in enumerate(data[:args.steps],1):
            ts=time.perf_counter();prompt=tok.encode(render_prompt(row['problem']))
            optimizer.zero_grad(set_to_none=True)
            gens=generate(model,tok,[prompt]*8,args.tokens,sample=True)
            rewards=[]
            for g in gens:
                text=tok.decode(g);ext=extract_final_candidate(text,fallback=None)
                reward=float(bool(ext) and grade_answer(ext,row['answer']));rewards.append(reward)
                samples.write(json.dumps(dict(step=step,problem=row['problem'],answer=row['answer'],text=text,reward=reward))+'\n')
            samples.flush()
            r=torch.tensor(rewards,device='cuda');adv=(r-r.mean())/(r.std()+1e-4)
            model.train();total_loss=0.
            # Still perform AdamW on zero-advantage groups, exactly as reference default.
            for i,g in enumerate(gens):
                seq=torch.tensor(prompt+g,device='cuda')
                loss=-adv[i].detach()*logprob(model,seq,len(prompt))/8
                loss.backward();total_loss+=loss.item()
            norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.)
            assert torch.isfinite(norm),'Nonfinite gradient'
            optimizer.step();torch.cuda.synchronize()
            record=dict(step=step,seconds=time.perf_counter()-ts,elapsed_seconds=time.perf_counter()-start,
                        reward=sum(rewards)/8,loss=total_loss,grad_norm=float(norm),generated_tokens=sum(map(len,gens)),**memory())
            f.write(json.dumps(record)+'\n');f.flush();emit(record)
    torch.save(model.state_dict(),out/'model.pth')
    total=time.perf_counter()-start
    result=dict(mode='train',steps=args.steps,seconds=total,**memory())
    (out/'training_summary.json').write_text(json.dumps(result,indent=2));emit(result)

def evaluate(args):
    model=EvalModel(dict(QWEN_CONFIG_06_B,context_length=4096)).cuda()
    path=Path(args.checkpoint) if args.checkpoint else ROOT/'data/qwen3-0.6B-base.pth'
    model.load_state_dict(torch.load(path,map_location='cpu',weights_only=True))
    tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'data/tokenizer-base.json')
    data=json.loads((ROOT/'data/math500_test.json').read_text())[:args.eval_n]
    (ROOT/'results').mkdir(exist_ok=True)
    # Length sorting reduces padding; keep original indexes in output.
    entries=sorted([(i,r,tok.encode(render_prompt(r['problem']))) for i,r in enumerate(data)],key=lambda z:len(z[2]))
    out=ROOT/'results'/args.name;out.mkdir(parents=True,exist_ok=False)
    (out/'evaluation_config.json').write_text(json.dumps(vars(args),indent=2))
    torch.cuda.reset_peak_memory_stats();start=time.perf_counter();correct=0;tokens=0
    with (out/'evaluation.jsonl').open('w') as f:
        for k in range(0,len(entries),args.batch):
            batch=entries[k:k+args.batch];gens=generate_eval(model,tok,[p for i,r,p in batch],args.eval_tokens)
            for (i,row,p),g in zip(batch,gens):
                text=tok.decode(g);ext=extract_final_candidate(text);ok=bool(grade_answer(ext,row['answer']))
                correct+=ok;tokens+=len(g)
                f.write(json.dumps(dict(index=i,problem=row['problem'],answer=row['answer'],text=text,extracted=ext,correct=ok,tokens=len(g)))+'\n')
            f.flush();emit(dict(evaluated=min(k+args.batch,len(entries)),correct=correct,seconds=time.perf_counter()-start,**memory()))
    result=dict(mode='eval',n=len(data),correct=correct,accuracy=correct/len(data),avg_tokens=tokens/len(data),seconds=time.perf_counter()-start,checkpoint=str(path),**memory())
    (out/'evaluation_summary.json').write_text(json.dumps(result,indent=2));emit(result)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('mode',choices=['verify','train','eval'])
    p.add_argument('--steps',type=int,default=50);p.add_argument('--tokens',type=int,default=512)
    p.add_argument('--seed',type=int,default=42);p.add_argument('--name',default='train_seed42')
    p.add_argument('--checkpoint');p.add_argument('--batch',type=int,default=16)
    p.add_argument('--eval-n',type=int,default=500);p.add_argument('--eval-tokens',type=int,default=2048)
    args=p.parse_args();torch.set_num_threads(4)
    if any(v <= 0 for v in (args.steps, args.tokens, args.batch, args.eval_n, args.eval_tokens)):
        p.error('Step, token, batch and evaluation counts must be positive')
    if not torch.cuda.is_available():
        p.error('An NVIDIA CUDA GPU is required for training, evaluation and verify')
    emit(dict(device=torch.cuda.get_device_name(),torch=torch.__version__,cuda=torch.version.cuda,mode=args.mode))
    {'verify':lambda a:verify(),'train':train,'eval':evaluate}[args.mode](args)
