/* dragon_native.c — 560B engine: NATIVE-PRECISION compute + MLA KV cache.
 * Matmuls read the mmap'd bf16/fp8 bytes DIRECTLY (inline dequant per element).
 * No f32 weight buffers, no LRU — the OS page cache holds the native bytes (2-4x smaller than f32).
 * KV cache validated in dragon_fast.c (reproduces recompute output exactly).
 * Build: gcc -O3 -march=native -fopenmp -D_GNU_SOURCE -o dragon_native dragon_native.c -lm
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <time.h>

#define HID 6144
#define H 64
#define NOPE 128
#define ROPE 64
#define HD 192
#define VD 128
#define QLORA 1536
#define KVLORA 512
#define EFFN 2048
#define FFN 12288
#define NEXP_TOT 768
#define NR 512
#define SEQMAX 64
#define NSUB 56

#pragma pack(push,1)
typedef struct { char name[128]; int32_t shard, dt; int64_t off, d0, d1; int32_t sshard; int64_t soff; int32_t s0,s1; } Rec;
#pragma pack(pop)
static Rec* recs; static int nrec;
static const uint8_t* shard_base[128]; static int nshard;
static void load_manifest(void){
    FILE* f=fopen("manifest.bin","rb"); if(!f){perror("manifest");exit(1);}
    if(fread(&nrec,4,1,f)!=1)exit(1); recs=malloc((long)nrec*sizeof(Rec)); if(fread(recs,sizeof(Rec),nrec,f)!=(size_t)nrec)exit(1); fclose(f);
    FILE* s=fopen("shards.txt","r"); char line[1024];
    while(fgets(line,sizeof line,s)){ line[strcspn(line,"\n")]=0; int fd=open(line,O_RDONLY); struct stat st; fstat(fd,&st);
        shard_base[nshard]=mmap(NULL,st.st_size,PROT_READ,MAP_PRIVATE,fd,0); if(shard_base[nshard]==MAP_FAILED){perror(line);exit(1);} close(fd); nshard++; } fclose(s);
    fprintf(stderr,"manifest: %d tensors, %d shards\n",nrec,nshard);
}
static Rec* find(const char* nm){ int lo=0,hi=nrec-1; while(lo<=hi){int m=(lo+hi)/2;int c=strcmp(recs[m].name,nm);if(c==0)return &recs[m];if(c<0)lo=m+1;else hi=m-1;} fprintf(stderr,"MISSING %s\n",nm);exit(2); }
static inline float bf16f(uint16_t b){uint32_t u=((uint32_t)b)<<16;float f;memcpy(&f,&u,4);return f;}
static inline float e4m3f(uint8_t b){int s=(b>>7)&1,e=(b>>3)&0xF,m=b&7;float v;if(e==0xF&&m==7)return NAN;if(e==0)v=ldexpf(m/8.0f,-6);else v=ldexpf(1.0f+m/8.0f,e-7);return s?-v:v;}

/* ---- NATIVE matvec: o[M] = W @ x, W read directly from mmap in its own dtype ----
 * THE MATMUL: fp8 via 256-entry LUT (e4m3 has only 256 values — zero ldexpf calls);
 * bf16 via AVX2+FMA (8 weights/instruction: widen u16, <<16, that IS the f32). */
static float E4M3_LUT[256];
static void lut_init(void){ for(int b=0;b<256;b++) E4M3_LUT[b]=e4m3f((uint8_t)b); }
#ifdef __AVX2__
#include <immintrin.h>
static inline float hsum8(__m256 v){ __m128 lo=_mm256_castps256_ps128(v), hi=_mm256_extractf128_ps(v,1);
    lo=_mm_add_ps(lo,hi); lo=_mm_hadd_ps(lo,lo); lo=_mm_hadd_ps(lo,lo); return _mm_cvtss_f32(lo); }
static inline float dot_bf16(const uint16_t* w, const float* x, long K){
    __m256 a0=_mm256_setzero_ps(), a1=_mm256_setzero_ps(); long k=0;
    for(; k+16<=K; k+=16){
        __m256i w0=_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i*)(w+k))),16);
        __m256i w1=_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i*)(w+k+8))),16);
        a0=_mm256_fmadd_ps(_mm256_castsi256_ps(w0),_mm256_loadu_ps(x+k),a0);
        a1=_mm256_fmadd_ps(_mm256_castsi256_ps(w1),_mm256_loadu_ps(x+k+8),a1);
    }
    float s=hsum8(_mm256_add_ps(a0,a1));
    for(; k<K; k++) s+=bf16f(w[k])*x[k];
    return s;
}
static inline float dot_fp8_blk(const uint8_t* w, const float* x, long n){ /* one 128-col block */
    __m256 acc=_mm256_setzero_ps(); long k=0;
    for(; k+8<=n; k+=8){
        __m256i idx=_mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i*)(w+k)));
        __m256 wf=_mm256_i32gather_ps(E4M3_LUT,idx,4);
        acc=_mm256_fmadd_ps(wf,_mm256_loadu_ps(x+k),acc);
    }
    float s=hsum8(acc);
    for(; k<n; k++) s+=E4M3_LUT[w[k]]*x[k];
    return s;
}
#else
static inline float dot_bf16(const uint16_t* w, const float* x, long K){ double a=0; for(long k=0;k<K;k++)a+=(double)bf16f(w[k])*x[k]; return (float)a; }
static inline float dot_fp8_blk(const uint8_t* w, const float* x, long n){ double a=0; for(long k=0;k<n;k++)a+=(double)E4M3_LUT[w[k]]*x[k]; return (float)a; }
#endif
static void mv_rec(Rec* r, const float* x, float* o){
    long M=r->d0, K=r->d1; const uint8_t* base=shard_base[r->shard]+r->off;
    if(r->dt==0){ const uint16_t* W=(const uint16_t*)base;
        #pragma omp parallel for schedule(static)
        for(long i=0;i<M;i++) o[i]=dot_bf16(W+i*K,x,K);
    } else if(r->dt==1){ const float* W=(const float*)base;
        #pragma omp parallel for schedule(static)
        for(long i=0;i<M;i++){ const float* w=W+i*K; double a=0;
            for(long k=0;k<K;k++) a+=(double)w[k]*x[k]; o[i]=(float)a; }
    } else { const float* sc=(const float*)(shard_base[r->sshard]+r->soff); long s1=r->s1;
        #pragma omp parallel for schedule(static)
        for(long i=0;i<M;i++){ const uint8_t* w=base+i*K; const float* srow=sc+(i/128)*s1; double a=0;
            for(long kb=0;kb<K;kb+=128){ long ke=kb+128<K?kb+128:K;
                a+=(double)dot_fp8_blk(w+kb,x+kb,ke-kb)*srow[kb/128]; }
            o[i]=(float)a; }
    }
}
static void mv_name(const char* nm, const float* x, float* o){ mv_rec(find(nm),x,o); }
/* small f32 copy for norm vectors etc. */
static void vec_name(const char* nm, float* out){ Rec* r=find(nm); long n=r->d0*r->d1; const uint8_t* p=shard_base[r->shard]+r->off;
    if(r->dt==0){const uint16_t* w=(const uint16_t*)p; for(long i=0;i<n;i++)out[i]=bf16f(w[i]);}
    else memcpy(out,p,n*4); }

static void rmsn(const float*x,const float*w,float*o,int d){double ss=0;for(int i=0;i<d;i++)ss+=(double)x[i]*x[i];float iv=1.0f/sqrtf((float)(ss/d)+1e-5f);for(int i=0;i<d;i++)o[i]=x[i]*iv*w[i];}
static inline float silu(float z){return z/(1.0f+expf(-z));}
static void il(const float*in,float*o,int d){int h=d/2;for(int i=0;i<h;i++){o[i]=in[2*i];o[h+i]=in[2*i+1];}}
static void rope(float*t,const float*co,const float*si){int h=ROPE/2;float tmp[ROPE];for(int i=0;i<ROPE;i++){float rh=(i<h)?-t[h+i]:t[i-h];tmp[i]=t[i]*co[i]+rh*si[i];}memcpy(t,tmp,ROPE*4);}

static float (*KC)[H][SEQMAX][HD];
static float (*VC)[H][SEQMAX][VD];
static float *G_COS,*G_SIN;
static double P_attn,P_dense,P_moe,P_head;
static double now(void){struct timespec t;clock_gettime(CLOCK_MONOTONIC,&t);return t.tv_sec+t.tv_nsec/1e9;}

static void mla(int layer,int blk,int sb,const float* hn,int npos,int startpos,float* out){
    char nm[160]; const float SQ=2.0f,SKV=3.4641016151377544f,SC=0.07216878364870322f;
    #define AW(w) (snprintf(nm,160,"model.layers.%d.self_attn.%d.%s.weight",layer,blk,w),nm)
    Rec *rqa,*rqb,*rkva,*rkvb,*rop; float qaln[QLORA],kvaln[KVLORA];
    rqa=find(AW("q_a_proj")); vec_name(AW("q_a_layernorm"),qaln); rqb=find(AW("q_b_proj"));
    rkva=find(AW("kv_a_proj_with_mqa")); vec_name(AW("kv_a_layernorm"),kvaln); rkvb=find(AW("kv_b_proj")); rop=find(AW("o_proj"));
    float ctxbuf[H*VD];
    for(int p=0;p<npos;p++){ int pos=startpos+p; const float* h=hn+(long)p*HID;
        float ql[QLORA],qln[QLORA],q[H*HD]; mv_rec(rqa,h,ql); rmsn(ql,qaln,qln,QLORA); mv_rec(rqb,qln,q);
        float c[KVLORA+ROPE]; mv_rec(rkva,h,c); float kl[KVLORA]; rmsn(c,kvaln,kl,KVLORA); for(int i=0;i<KVLORA;i++)kl[i]*=SKV;
        float kp[H*(NOPE+VD)]; mv_rec(rkvb,kl,kp);
        float kr[ROPE],ki[ROPE]; il(c+KVLORA,ki,ROPE); memcpy(kr,ki,ROPE*4); rope(kr,G_COS+pos*ROPE,G_SIN+pos*ROPE);
        #pragma omp parallel for schedule(static)
        for(int hd=0;hd<H;hd++){ const float* qh=q+hd*HD; float Q[HD];
            for(int i=0;i<NOPE;i++)Q[i]=qh[i]*SQ;
            float qr[ROPE],qi[ROPE]; for(int i=0;i<ROPE;i++)qr[i]=qh[NOPE+i]*SQ; il(qr,qi,ROPE);memcpy(qr,qi,ROPE*4);rope(qr,G_COS+pos*ROPE,G_SIN+pos*ROPE);
            for(int i=0;i<ROPE;i++)Q[NOPE+i]=qr[i];
            const float* kk=kp+hd*(NOPE+VD);
            for(int i=0;i<NOPE;i++)KC[sb][hd][pos][i]=kk[i]; for(int i=0;i<ROPE;i++)KC[sb][hd][pos][NOPE+i]=kr[i];
            for(int i=0;i<VD;i++)VC[sb][hd][pos][i]=kk[NOPE+i];
            float sc[SEQMAX],mx=-1e30f; for(int k2=0;k2<=pos;k2++){double a=0;for(int d=0;d<HD;d++)a+=(double)Q[d]*KC[sb][hd][k2][d];sc[k2]=(float)a*SC;if(sc[k2]>mx)mx=sc[k2];}
            double sm=0;for(int k2=0;k2<=pos;k2++){sc[k2]=expf(sc[k2]-mx);sm+=sc[k2];}
            for(int d=0;d<VD;d++){double a=0;for(int k2=0;k2<=pos;k2++)a+=sc[k2]*VC[sb][hd][k2][d];ctxbuf[hd*VD+d]=a/sm;}
        }
        mv_rec(rop,ctxbuf,out+(long)p*HID);
    }
}
static void dense(int layer,int blk,const float*hn,int npos,float*out){
    char nm[160]; Rec *rg,*ru,*rd;
    snprintf(nm,160,"model.layers.%d.mlps.%d.gate_proj.weight",layer,blk); rg=find(nm);
    snprintf(nm,160,"model.layers.%d.mlps.%d.up_proj.weight",layer,blk);   ru=find(nm);
    snprintf(nm,160,"model.layers.%d.mlps.%d.down_proj.weight",layer,blk); rd=find(nm);
    for(int p=0;p<npos;p++){ float gp[FFN],up[FFN],ac[FFN];
        mv_rec(rg,hn+(long)p*HID,gp); mv_rec(ru,hn+(long)p*HID,up);
        for(int i=0;i<FFN;i++)ac[i]=silu(gp[i])*up[i];
        mv_rec(rd,ac,out+(long)p*HID); }
}
static void moe(int layer,const float*hn,int npos,float*out){
    char nm[160]; Rec* rcw; float bias[NEXP_TOT];
    snprintf(nm,160,"model.layers.%d.mlp.router.classifier.weight",layer); rcw=find(nm);
    snprintf(nm,160,"model.layers.%d.mlp.router.e_score_correction_bias",layer); vec_name(nm,bias);
    memset(out,0,(long)npos*HID*4);
    for(int p=0;p<npos;p++){ float lg[NEXP_TOT]; mv_rec(rcw,hn+(long)p*HID,lg);
        float mx=-1e30f;for(int i=0;i<NEXP_TOT;i++)if(lg[i]>mx)mx=lg[i];double sm=0;float sf[NEXP_TOT];for(int i=0;i<NEXP_TOT;i++){sf[i]=expf(lg[i]-mx);sm+=sf[i];}for(int i=0;i<NEXP_TOT;i++)sf[i]/=sm;
        int idx[12];float ch[NEXP_TOT];for(int i=0;i<NEXP_TOT;i++)ch[i]=sf[i]+bias[i];
        for(int j=0;j<12;j++){int bi=-1;float bv=-1e30f;for(int i=0;i<NEXP_TOT;i++){int tk=0;for(int q=0;q<j;q++)if(idx[q]==i)tk=1;if(!tk&&ch[i]>bv){bv=ch[i];bi=i;}}idx[j]=bi;}
        for(int j=0;j<12;j++){int e=idx[j];float w=sf[e]*6.0f;
            if(e>=NR){for(int i=0;i<HID;i++)out[(long)p*HID+i]+=hn[(long)p*HID+i]*w;continue;}
            Rec *rg,*ru,*rd;
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,e); rg=find(nm);
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.up_proj.weight",layer,e);   ru=find(nm);
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.down_proj.weight",layer,e); rd=find(nm);
            float gp[EFFN],up[EFFN],ac[EFFN],dp[HID];
            mv_rec(rg,hn+(long)p*HID,gp); mv_rec(ru,hn+(long)p*HID,up);
            for(int i=0;i<EFFN;i++)ac[i]=silu(gp[i])*up[i];
            mv_rec(rd,ac,dp);
            for(int i=0;i<HID;i++)out[(long)p*HID+i]+=dp[i]*w;
        }
    }
}
static void layer(int L,float* x,int npos,int startpos){
    static float tmp[(long)SEQMAX*HID],xn[(long)SEQMAX*HID],shortcut[(long)SEQMAX*HID]; char nm[160]; float ln[HID];
    for(int i=0;i<2;i++){ int sb=L*2+i;
        snprintf(nm,160,"model.layers.%d.input_layernorm.%d.weight",L,i); vec_name(nm,ln);
        for(int p=0;p<npos;p++) rmsn(x+(long)p*HID,ln,xn+(long)p*HID,HID);
        { double t=now(); mla(L,i,sb,xn,npos,startpos,tmp); P_attn+=now()-t; } for(long j=0;j<(long)npos*HID;j++)x[j]+=tmp[j];
        snprintf(nm,160,"model.layers.%d.post_attention_layernorm.%d.weight",L,i); vec_name(nm,ln);
        for(int p=0;p<npos;p++) rmsn(x+(long)p*HID,ln,xn+(long)p*HID,HID);
        if(i==0){ double t=now(); moe(L,xn,npos,shortcut); P_moe+=now()-t; }
        { double t=now(); dense(L,i,xn,npos,tmp); P_dense+=now()-t; } for(long j=0;j<(long)npos*HID;j++)x[j]+=tmp[j];
        if(i==1) for(long j=0;j<(long)npos*HID;j++)x[j]+=shortcut[j];
    }
}
static char* VOCAB; static int* VOFF; static int VN;
static void load_vocab(void){FILE*f=fopen("vocab.bin","rb");if(fread(&VN,4,1,f)!=1)exit(1);VOCAB=malloc(1L<<24);VOFF=malloc((VN+1)*4);long o=0;for(int i=0;i<VN;i++){unsigned char L;if(fread(&L,1,1,f)!=1)break;VOFF[i]=o;if(fread(VOCAB+o,1,L,f)!=L)break;o+=L;VOCAB[o++]=0;}VOFF[VN]=o;fclose(f);}
static void embed_ids(const int*ids,int n,float*out){Rec*r=find("model.embed_tokens.weight");const uint16_t*w=(const uint16_t*)(shard_base[r->shard]+r->off);for(int t=0;t<n;t++){const uint16_t*row=w+(long)ids[t]*HID;for(int k=0;k<HID;k++)out[(long)t*HID+k]=bf16f(row[k]);}}
static int lm_argmax(const float* xrow){
    float nw[HID]; vec_name("model.norm.weight",nw);
    float last[HID]; rmsn(xrow,nw,last,HID);
    Rec*r=find("lm_head.weight"); long V=r->d0; const uint16_t*w=(const uint16_t*)(shard_base[r->shard]+r->off);
    int best=-1;float bv=-1e30f;
    #pragma omp parallel
    { int lb=-1;float lv=-1e30f;
      #pragma omp for
      for(long v=0;v<V;v++){float s=dot_bf16(w+v*HID,last,HID);if(s>lv){lv=s;lb=v;}}
      #pragma omp critical
      {if(lv>bv){bv=lv;best=lb;}} }
    return best;
}
static float* Lf(const char*n,long c){char p[512];snprintf(p,sizeof p,"ref/%s.f32",n);FILE*f=fopen(p,"rb");if(!f){perror(p);exit(1);}float*b=malloc(c*4);if(fread(b,4,c,f)!=(size_t)c)exit(1);fclose(f);return b;}

int main(int argc,char**argv){
    int NGEN=argc>1?atoi(argv[1]):8;
    load_manifest(); load_vocab(); lut_init();
    KC=malloc((long)NSUB*sizeof(*KC)); VC=malloc((long)NSUB*sizeof(*VC));
    G_COS=Lf("rope_cos64",64*64); G_SIN=Lf("rope_sin64",64*64);
    FILE* pf=fopen("prompt_ids.i32","rb"); fseek(pf,0,SEEK_END);int pn=ftell(pf)/4;fseek(pf,0,SEEK_SET);
    int ids[SEQMAX]; if(fread(ids,4,pn,pf)!=(size_t)pn)exit(1); fclose(pf);
    static float x[(long)SEQMAX*HID];
    struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
    embed_ids(ids,pn,x);
    for(int L=0;L<28;L++) layer(L,x,pn,0);
    int nx=lm_argmax(x+(long)(pn-1)*HID);
    clock_gettime(CLOCK_MONOTONIC,&t1); double pt=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)/1e9;
    printf("PROMPT: "); for(int i=0;i<pn;i++)printf("%s",VOCAB+VOFF[ids[i]]); printf("\n");
    fprintf(stderr,"[prefill %d tok in %.1fs (%.1fs/tok)] first=%d\n",pn,pt,pt/pn,nx);
    printf("GEN: %s",VOCAB+VOFF[nx]); fflush(stdout);
    int n=pn; ids[n++]=nx;
    for(int g=0; g<NGEN-1 && n<SEQMAX; g++){
        clock_gettime(CLOCK_MONOTONIC,&t0);
        embed_ids(&ids[n-1],1,x);
        for(int L=0;L<28;L++) layer(L,x,1,n-1);
        double th=now(); int t=lm_argmax(x); P_head+=now()-th; clock_gettime(CLOCK_MONOTONIC,&t1);
        double dt=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)/1e9;
        if(t==2){printf("[EOS]");break;}
        printf("%s",VOCAB+VOFF[t]); fflush(stdout);
        fprintf(stderr,"[tok '%s' %.1fs | attn %.1f dense %.1f moe %.1f head %.1f]\n",VOCAB+VOFF[t],dt,P_attn,P_dense,P_moe,P_head); P_attn=P_dense=P_moe=P_head=0;
        ids[n++]=t;
    }
    printf("\n");
    return 0;
}