import { numpy as np, nn, jit, valueAndGrad, tree, blockUntilReady } from '@jax-js/jax';
import {sampleCharacter} from './sampling.js';

export function rng(seed = 42) {
  let a = seed >>> 0;
  return () => { a += 0x6D2B79F5; let t = a; t = Math.imul(t ^ t >>> 15, t | 1); t ^= t + Math.imul(t ^ t >>> 7, t | 61); return ((t ^ t >>> 14) >>> 0) / 4294967296; };
}
export function tokenize(text) {
  const chars = [...new Set([...text])].sort();
  if (chars.length < 2 || chars.length > 128) throw new Error('The dataset must contain 2 to 128 distinct characters.');
  const index = new Map(chars.map((c, i) => [c, i]));
  const letters=[...text],paragraphStarts=[0];
  for(let i=1;i<letters.length-1;i++)if(letters[i-1]==='\n'&&letters[i]==='\n'&&letters[i+1]!=='\n')paragraphStarts.push(i+1);
  return { chars, paragraphStarts, ids: [...text].map(c => index.get(c)), encode: s => [...s].map(c => { if (!index.has(c)) throw new Error(`Character missing from the dataset: ${JSON.stringify(c)}`); return index.get(c); }) };
}
export function initParams(c, seed = 42) {
  const random = rng(seed);
  const weight = (shape, scale) => np.array(Float32Array.from({ length: shape.reduce((a,b) => a*b,1) }, () => (random()*2-1)*scale)).reshape(shape);
  if(c.architecture==='nanogpt') {
    // NanoGPT bias=False, dropout=0: LayerNorm, exact GELU, tied embedding/head.
    // Same distribution and residual scaling; JS/Python RNG streams differ.
    const normal=(shape,std=.02)=>np.array(Float32Array.from({length:shape.reduce((a,b)=>a*b,1)},()=>Math.sqrt(-2*Math.log(Math.max(1e-9,random())))*Math.cos(2*Math.PI*random())*std)).reshape(shape);
    const p={embed:normal([c.vocab,c.dim]),pos:normal([c.context,c.dim]),norm:np.ones([c.dim])};
    for(let l=0;l<c.layers;l++) {
      for(const k of ['q','k','v'])p[`${l}.${k}`]=normal([c.dim,c.dim]);
      p[`${l}.o`]=normal([c.dim,c.dim],.02/Math.sqrt(2*c.layers));
      p[`${l}.up`]=normal([c.dim,4*c.dim]);
      p[`${l}.down`]=normal([4*c.dim,c.dim],.02/Math.sqrt(2*c.layers));
      p[`${l}.n1`]=np.ones([c.dim]);p[`${l}.n2`]=np.ones([c.dim]);
    }
    return p;
  }
  if(c.architecture==='modern') {
    // Modded-inspired FP32 variant: Gaussian fan-in init, zero residual projections,
    // RMSNorm, QK norm, RoPE, squared ReLU, soft-capped untied head. AdamW below.
    const normal=(shape,std)=>np.array(Float32Array.from({length:shape.reduce((a,b)=>a*b,1)},()=>Math.sqrt(-2*Math.log(Math.max(1e-9,random())))*Math.cos(2*Math.PI*random())*std)).reshape(shape);
    const zeroOrNormal=shape=>c.nonzeroInit?normal(shape,0.1):np.zeros(shape);
    const p={embed:normal([c.vocab,c.dim],1),norm:np.ones([c.dim]),head:zeroOrNormal([c.dim,c.vocab])};
    for(let l=0;l<c.layers;l++) {
      for(const k of ['q','k','v'])p[`${l}.${k}`]=normal([c.dim,c.dim],Math.sqrt(.33/c.dim));
      p[`${l}.o`]=zeroOrNormal([c.dim,c.dim]);
      p[`${l}.up`]=normal([c.dim,4*c.dim],Math.sqrt(.33/c.dim));
      p[`${l}.down`]=zeroOrNormal([4*c.dim,c.dim]);
      p[`${l}.n1`]=np.ones([c.dim]);p[`${l}.n2`]=np.ones([c.dim]);
    }
    return p;
  }
  const p = { embed: weight([c.vocab,c.dim],0.2), pos: weight([c.context,c.dim],0.02) };
  for(let l=0;l<c.layers;l++) {
    for(const k of ['q','k','v','o']) p[`${l}.${k}`] = weight([c.dim,c.dim],Math.sqrt(3/c.dim));
    p[`${l}.up`] = weight([c.dim,4*c.dim],Math.sqrt(3/c.dim));
    p[`${l}.down`] = weight([4*c.dim,c.dim],Math.sqrt(3/(4*c.dim)));
    p[`${l}.n1`] = np.ones([c.dim]); p[`${l}.n2`] = np.ones([c.dim]);
  }
  p.norm = np.ones([c.dim]); p.head = weight([c.dim,c.vocab],0.02);
  // Keep the same RNG draws for other weights, for the positional ablation.
  if(c.architecture==='rotary'){p.pos.dispose();delete p.pos;}
  return p;
}
const norm = (x,w) => x.ref.div(np.sqrt(np.mean(np.square(x),-1,{keepdims:true}).add(1e-5))).mul(w);
const layerNorm=(x,w)=>{const centered=x.ref.sub(np.mean(x,-1,{keepdims:true}));return norm(centered,w);};
const rms = x => x.ref.div(np.sqrt(np.mean(np.square(x),-1,{keepdims:true}).add(1e-5)));
function rotary(x,t,hd,full=false) {
  const angles=Array.from({length:t*hd/2},(_,i)=>{
    const j=i%(hd/2),pos=Math.floor(i/(hd/2));
    return full?pos*(1/10000)**(2*j/hd):j<hd/4?pos*(1/1024)**(j/Math.max(1,hd/4-1)):0;
  });
  const cs=np.array(angles.map(Math.cos)).reshape([1,t,1,hd/2]);
  const sn=np.array(angles.map(Math.sin)).reshape([1,t,1,hd/2]);
  const [a,b]=np.split(x,2,-1);
  return np.concatenate([a.ref.mul(cs.ref).add(b.ref.mul(sn.ref)),b.mul(cs).sub(a.mul(sn))],-1);
}
// Each argument is consumed exactly once, including every parameter leaf.
export function forward(p, ids, c) {
  const [b,t] = ids.shape, d = c.dim, h = c.heads;
  const modern=c.architecture==='modern',rope=c.architecture==='rotary',nano=c.architecture==='nanogpt',normalize=nano?layerNorm:norm;
  let x = np.matmul(nn.oneHot(ids,c.vocab),nano?p.embed.ref:p.embed);
  if(!modern&&!rope)x=x.add(t===c.context?p.pos:p.pos.slice([0,t],[]));
  for(let l=0;l<c.layers;l++) {
    const n = normalize(x.ref,p[`${l}.n1`]);
    let q = np.matmul(n.ref,p[`${l}.q`]).reshape([b,t,h,d/h]);
    let k = np.matmul(n.ref,p[`${l}.k`]).reshape([b,t,h,d/h]);
    const v = np.matmul(n,p[`${l}.v`]).reshape([b,t,h,d/h]);
    if(modern){q=rotary(rms(q),t,d/h);k=rotary(rms(k),t,d/h);}
    if(rope){q=rotary(q,t,d/h,true);k=rotary(k,t,d/h,true);}
    const a = nn.dotProductAttention(q,k,v,{isCausal:true,...(modern?{scale:.12}:{})}).reshape([b,t,d]);
    x = x.add(np.matmul(a,p[`${l}.o`]));
    const z = normalize(x.ref,p[`${l}.n2`]);
    const u=np.matmul(z,p[`${l}.up`]);
    x = x.add(np.matmul(modern?np.square(nn.relu(u)):nn.gelu(u,{approximate:!nano}),p[`${l}.down`]));
  }
  const logits=np.matmul(normalize(x,p.norm),nano?p.embed.transpose():p.head);
  return modern?logits.ref.mul(15).div(np.sqrt(np.square(logits).add(225))):logits;
}
export function loss(p,x,y,c) {
  return nn.logSoftmax(forward(p,x,c),-1).mul(nn.oneHot(y,c.vocab)).sum(-1).mean().neg();
}
export function batch(ids,c,random) {
  const x = [], y = [];
  for(let b=0;b<c.batch;b++) {
    const start = Math.floor(random()*(ids.length-c.context));
    x.push(...ids.slice(start,start+c.context)); y.push(...ids.slice(start+1,start+c.context+1));
  }
  return [np.array(x,{dtype:np.int32}).reshape([c.batch,c.context]),np.array(y,{dtype:np.int32}).reshape([c.batch,c.context])];
}
// Paragraph starts use Unicode code-point offsets, like token ids.
// Circular continuation includes the last paragraph without padding fake targets.
export function paragraphBatch(ids,c,random,starts) {
  if(!starts?.length)throw new Error('No paragraph starts available.');
  const x=[],y=[];
  for(let b=0;b<c.batch;b++){
    const start=starts[Math.floor(random()*starts.length)];
    for(let i=0;i<c.context;i++){x.push(ids[(start+i)%ids.length]);y.push(ids[(start+i+1)%ids.length]);}
  }
  return [np.array(x,{dtype:np.int32}).reshape([c.batch,c.context]),np.array(y,{dtype:np.int32}).reshape([c.batch,c.context])];
}
export class Trainer {
  constructor(c,seed=42) {
    this.c=c; this.p=initParams(c,seed); this.step=0;
    this.m=tree.map(p=>np.zeros(p.shape),this.p); this.v=tree.map(p=>np.zeros(p.shape),this.p);
    this.lossGrad=jit(valueAndGrad((p,x,y)=>loss(p,x,y,c)));
    this.infer=jit((p,x)=>forward(p,x,c));
    this.inferLast=jit((p,x,i)=>np.take(forward(p,x,c),i,1));
    this.evaluate=jit((p,x,y)=>nn.logSoftmax(forward(p,x,c),-1).mul(nn.oneHot(y,c.vocab)).sum(-1).slice([], [Math.max(0,c.context-16),c.context]).mean().neg());
    this.beta2=c.beta2??.999;
    this.update=jit((p,m,v,g,lr,b1,b2)=>{
      // Global norm clipping before Adam.
      const total=Object.values(g).map(x=>np.square(x.ref).sum()).reduce((a,b)=>a.add(b));
      const scale=np.minimum(1,np.reciprocal(np.sqrt(total).add(1e-8)));
      const nextP={},nextM={},nextV={};
      for(const k of Object.keys(p)) {
        const grad=g[k].mul(scale.ref);
        const mm=m[k].mul(0.9).add(grad.ref.mul(0.1));
        const vv=v[k].mul(this.beta2).add(np.square(grad).mul(1-this.beta2));
        const decayed=p[k].mul(np.array(1).sub(lr.ref.mul(p[k].shape.length===2?(c.weightDecay??0):0)));
        nextP[k]=decayed.sub(mm.ref.div(b1.ref).div(np.sqrt(vv.ref.div(b2.ref)).add(1e-8)).mul(lr.ref));
        nextM[k]=mm; nextV[k]=vv;
      }
      scale.dispose(); lr.dispose(); b1.dispose(); b2.dispose();
      return [nextP,nextM,nextV];
    });
  }
  async train(x,y,rate=0.003) {
    const [l,g]=this.lossGrad(tree.ref(this.p),x,y);
    this.step++;
    [this.p,this.m,this.v]=this.update(this.p,this.m,this.v,g,np.array(rate),np.array(1-0.9**this.step),np.array(1-this.beta2**this.step));
    await blockUntilReady([this.p,this.m,this.v]);
    const value=(await l.data())[0];
    if(!Number.isFinite(value)) throw new Error('Non-finite loss. Reset and lower the learning rate.');
    return value;
  }
  async generate(prefix,tokenizer,{length=80,temperature=0.8,minP=0,seed=123,bucketed=this.c.fastSampling??false,cancel=()=>false,onToken=()=>{}}={}) {
    const ids=tokenizer.encode(prefix); if(!ids.length) throw new Error('The prompt must not be empty.');
    const random=rng(seed); let text='';
    for(let i=0;i<length&&!cancel();i++) {
      const ctx=ids.slice(-this.c.context), pos=ctx.length-1;
      const size=bucketed?Math.min(this.c.context,Math.max(16,2**Math.ceil(Math.log2(ctx.length)))):this.c.context;
      const padded=[...ctx,...Array(size-ctx.length).fill(0)];
      const x=np.array(padded,{dtype:np.int32}).reshape([1,size]);
      const logits=bucketed?this.inferLast(tree.ref(this.p),x,np.array(pos,{dtype:np.int32})):this.infer(tree.ref(this.p),x);
      const raw=await logits.data(),values=bucketed?raw:raw.slice(pos*this.c.vocab,(pos+1)*this.c.vocab); 
      const choice=sampleCharacter(values,temperature,random,minP);
      ids.push(choice); text+=tokenizer.chars[choice]; onToken(text);
      if(i%8===0) await new Promise(r=>setTimeout(r,0));
    }
    return text;
  }
  dispose(){tree.dispose([this.p,this.m,this.v]); this.lossGrad.dispose(); this.infer.dispose(); this.inferLast.dispose(); this.evaluate.dispose(); this.update.dispose();}
}