/* dragon.c — the streaming 560B engine (Phase 3a: native get_tensor + one layer via mmap).
 * Reads any weight from mmap'd shards by name (manifest bsearch), dequantizes bf16/fp8 on the fly.
 * Validates layer-0 (2f math, native reads) vs layer0_target.f32.
 * Build: gcc -O3 -fopenmp -D_GNU_SOURCE -o dragon dragon.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>

#define HID 6144
#define T 5
#define NEXP_TOT 768
#define NR 512

/* ---- manifest ---- */
#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 long shard_size[128]; static int nshard;

static void load_manifest(void){
    FILE* f=fopen("manifest.bin","rb");
    if(!f){perror("manifest");exit(1);}
    fread(&nrec,4,1,f); recs=malloc((long)nrec*sizeof(Rec)); fread(recs,sizeof(Rec),nrec,f); 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);}
        shard_size[nshard]=st.st_size; close(fd); nshard++;
    } fclose(s);
    fprintf(stderr,"manifest: %d tensors, %d shards mmapped\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 tensor: %s\n",nm); exit(2);
}
static float bf16f(uint16_t b){uint32_t u=((uint32_t)b)<<16;float f;memcpy(&f,&u,4);return f;}
static 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;}
/* dequant tensor <nm> into out (float), returns d0,d1 */
static void get_tensor(const char* nm, float* out, long* pd0, long* pd1){
    Rec* r=find(nm); long d0=r->d0,d1=r->d1; *pd0=d0;*pd1=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<d0*d1;i++) out[i]=bf16f(w[i]); }
    else if(r->dt==1){ memcpy(out,p,d0*d1*4); }
    else { /* fp8 + 128x128 block scale */
        const float* sc=(const float*)(shard_base[r->sshard]+r->soff); long s1=r->s1;
        #pragma omp parallel for
        for(long i=0;i<d0;i++){ long sr=(i/128)*s1; for(long j=0;j<d1;j++) out[i*d1+j]=e4m3f(p[i*d1+j])*sc[sr+j/128]; }
    }
}
/* ---- ops (from 2f, OpenMP matvec) ---- */
static void mv(const float*W,const float*x,float*o,int M,int Kd){
    #pragma omp parallel for
    for(int i=0;i<M;i++){double a=0;const float*w=W+(long)i*Kd;for(int k=0;k<Kd;k++)a+=(double)w[k]*x[k];o[i]=a;}
}
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));}

#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
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 *bufA,*bufB; /* scratch for big weight loads */
static void mla(int layer,int blk,const float*hn,const float*cosT,const float*sinT,float*out){
    char nm[160]; long a,b; const float SQ=2.0f,SKV=3.4641016151377544f,SC=0.07216878364870322f;
    #define GW(w) (snprintf(nm,160,"model.layers.%d.self_attn.%d.%s.weight",layer,blk,w),nm)
    static float Q[H][T][HD],KK[H][T][HD],V[H][T][VD];
    static float qa[(long)QLORA*HID],qb[(long)H*HD*QLORA],kva[(long)(KVLORA+ROPE)*HID],kvb[(long)H*(NOPE+VD)*KVLORA],op[(long)HID*H*VD];
    float qaln[QLORA],kvaln[KVLORA];
    get_tensor(GW("q_a_proj"),qa,&a,&b); get_tensor(GW("q_a_layernorm"),qaln,&a,&b); get_tensor(GW("q_b_proj"),qb,&a,&b);
    get_tensor(GW("kv_a_proj_with_mqa"),kva,&a,&b); get_tensor(GW("kv_a_layernorm"),kvaln,&a,&b); get_tensor(GW("kv_b_proj"),kvb,&a,&b); get_tensor(GW("o_proj"),op,&a,&b);
    for(int t=0;t<T;t++){const float*h=hn+(long)t*HID;
        float ql[QLORA],qln[QLORA],q[H*HD]; mv(qa,h,ql,QLORA,HID);rmsn(ql,qaln,qln,QLORA);mv(qb,qln,q,H*HD,QLORA);
        float c[KVLORA+ROPE];mv(kva,h,c,KVLORA+ROPE,HID);float kl[KVLORA];rmsn(c,kvaln,kl,KVLORA);for(int i=0;i<KVLORA;i++)kl[i]*=SKV;
        float kp[H*(NOPE+VD)];mv(kvb,kl,kp,H*(NOPE+VD),KVLORA);
        float kr[ROPE],ki[ROPE];il(c+KVLORA,ki,ROPE);memcpy(kr,ki,ROPE*4);rope(kr,cosT+t*ROPE,sinT+t*ROPE);
        for(int hd=0;hd<H;hd++){const float*qh=q+hd*HD;for(int i=0;i<NOPE;i++)Q[hd][t][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,cosT+t*ROPE,sinT+t*ROPE);
            for(int i=0;i<ROPE;i++)Q[hd][t][NOPE+i]=qr[i];
            const float*kk=kp+hd*(NOPE+VD);for(int i=0;i<NOPE;i++)KK[hd][t][i]=kk[i];for(int i=0;i<ROPE;i++)KK[hd][t][NOPE+i]=kr[i];for(int i=0;i<VD;i++)V[hd][t][i]=kk[NOPE+i];}
    }
    static float ctx[T][H*VD];
    #pragma omp parallel for
    for(int hd=0;hd<H;hd++)for(int qi=0;qi<T;qi++){float sc[T],mx=-1e30f;for(int ki=0;ki<=qi;ki++){double aa=0;for(int d=0;d<HD;d++)aa+=(double)Q[hd][qi][d]*KK[hd][ki][d];sc[ki]=(float)aa*SC;if(sc[ki]>mx)mx=sc[ki];}double sm=0;for(int ki=0;ki<=qi;ki++){sc[ki]=expf(sc[ki]-mx);sm+=sc[ki];}for(int d=0;d<VD;d++){double aa=0;for(int ki=0;ki<=qi;ki++)aa+=sc[ki]*V[hd][ki][d];ctx[qi][hd*VD+d]=aa/sm;}}
    for(int t=0;t<T;t++)mv(op,ctx[t],out+(long)t*HID,HID,H*VD);
}
static void dense(int layer,int blk,const float*hn,float*out){
    char nm[160];long a,b; static float g[(long)FFN*HID],u[(long)FFN*HID],d[(long)HID*FFN];
    snprintf(nm,160,"model.layers.%d.mlps.%d.gate_proj.weight",layer,blk);get_tensor(nm,g,&a,&b);
    snprintf(nm,160,"model.layers.%d.mlps.%d.up_proj.weight",layer,blk);get_tensor(nm,u,&a,&b);
    snprintf(nm,160,"model.layers.%d.mlps.%d.down_proj.weight",layer,blk);get_tensor(nm,d,&a,&b);
    for(int t=0;t<T;t++){float gp[FFN],up[FFN],ac[FFN];mv(g,hn+(long)t*HID,gp,FFN,HID);mv(u,hn+(long)t*HID,up,FFN,HID);for(int i=0;i<FFN;i++)ac[i]=silu(gp[i])*up[i];mv(d,ac,out+(long)t*HID,HID,FFN);}
}
static void moe(int layer,const float*hn,float*out){
    char nm[160];long a,b; static float cw[(long)NEXP_TOT*HID]; float bias[NEXP_TOT];
    snprintf(nm,160,"model.layers.%d.mlp.router.classifier.weight",layer);get_tensor(nm,cw,&a,&b);
    snprintf(nm,160,"model.layers.%d.mlp.router.e_score_correction_bias",layer);get_tensor(nm,bias,&a,&b);
    memset(out,0,(long)T*HID*4);
    static float g[(long)EFFN*HID],u[(long)EFFN*HID],d[(long)HID*EFFN];
    for(int t=0;t<T;t++){float lg[NEXP_TOT];mv(cw,hn+(long)t*HID,lg,NEXP_TOT,HID);
        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)t*HID+i]+=hn[(long)t*HID+i]*w;continue;}
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,e);get_tensor(nm,g,&a,&b);
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.up_proj.weight",layer,e);get_tensor(nm,u,&a,&b);
            snprintf(nm,160,"model.layers.%d.mlp.experts.%d.down_proj.weight",layer,e);get_tensor(nm,d,&a,&b);
            float gp[EFFN],up[EFFN],ac[EFFN],dp[HID];mv(g,hn+(long)t*HID,gp,EFFN,HID);mv(u,hn+(long)t*HID,up,EFFN,HID);for(int i=0;i<EFFN;i++)ac[i]=silu(gp[i])*up[i];mv(d,ac,dp,HID,EFFN);
            for(int i=0;i<HID;i++)out[(long)t*HID+i]+=dp[i]*w;
        }
    }
}
static void layer(int L,float*x,const float*cosT,const float*sinT){
    static float tmp[(long)T*HID],xn[(long)T*HID],shortcut[(long)T*HID]; char nm[160];long a,b;float ln[HID];
    for(int i=0;i<2;i++){
        snprintf(nm,160,"model.layers.%d.input_layernorm.%d.weight",L,i);get_tensor(nm,ln,&a,&b);
        for(int t=0;t<T;t++)rmsn(x+(long)t*HID,ln,xn+(long)t*HID,HID);
        mla(L,i,xn,cosT,sinT,tmp); for(long j=0;j<(long)T*HID;j++)x[j]+=tmp[j];
        snprintf(nm,160,"model.layers.%d.post_attention_layernorm.%d.weight",L,i);get_tensor(nm,ln,&a,&b);
        for(int t=0;t<T;t++)rmsn(x+(long)t*HID,ln,xn+(long)t*HID,HID);
        if(i==0)moe(L,xn,shortcut);
        dense(L,i,xn,tmp); for(long j=0;j<(long)T*HID;j++)x[j]+=tmp[j];
        if(i==1)for(long j=0;j<(long)T*HID;j++)x[j]+=shortcut[j];
    }
}
static float* Lf(const char*n,long c){char p[512];snprintf(p,sizeof p,"ref/%s.f32",n);FILE*f=fopen(p,"rb");float*b=malloc(c*4);fread(b,4,c,f);fclose(f);return b;}


int main(int argc,char**argv){
    int NLAYERS = argc>1?atoi(argv[1]):1;
    load_manifest();
    float* x=Lf("embed_out",(long)T*HID);
    float* cosT=Lf("mla_cos",(long)T*ROPE),*sinT=Lf("mla_sin",(long)T*ROPE);
    for(int L=0;L<NLAYERS;L++){ layer(L,x,cosT,sinT); fprintf(stderr,"layer %d done\n",L); }
    printf("out[0,:6]: ");for(int i=0;i<6;i++)printf("%.4f ",x[i]);printf("\n");
    if(NLAYERS==1){ float* ref=Lf("layer0_target",(long)T*HID);double dm=0;for(long j=0;j<(long)T*HID;j++){double dd=fabs((double)x[j]-ref[j]);if(dd>dm)dm=dd;}
        printf("PHASE 3a (native mmap layer-0) max diff vs Python: %.3e\nVERDICT: %s\n",dm,dm<0.1?"PASS":"FAIL"); }
    if(NLAYERS>=28){
        long a,b; float nw[HID]; get_tensor("model.norm.weight",nw,&a,&b);
        float last[HID]; rmsn(x+(long)(T-1)*HID,nw,last,HID);   /* final norm on last token */
        /* lm_head [131072,6144] bf16 — stream row by row, track top-5 */
        Rec* r=find("lm_head.weight"); long V=r->d0; const uint16_t* w=(const uint16_t*)(shard_base[r->shard]+r->off);
        fprintf(stderr,"lm_head V=%ld, streaming argmax...\n",V);
        float top[5]; int tid[5]; for(int i=0;i<5;i++){top[i]=-1e30f;tid[i]=-1;}
        #pragma omp parallel
        { float lt[5]; int li[5]; for(int i=0;i<5;i++){lt[i]=-1e30f;li[i]=-1;}
          #pragma omp for
          for(long v=0;v<V;v++){ const uint16_t* wr=w+v*HID; double s=0; for(int k=0;k<HID;k++)s+=(double)bf16f(wr[k])*last[k];
            float f=s; for(int i=0;i<5;i++) if(f>lt[i]){ for(int j=4;j>i;j--){lt[j]=lt[j-1];li[j]=li[j-1];} lt[i]=f;li[i]=v; break; } }
          #pragma omp critical
          { for(int i=0;i<5;i++) for(int j=0;j<5;j++) if(lt[i]>top[j]){ for(int m=4;m>j;m--){top[m]=top[m-1];tid[m]=tid[m-1];} top[j]=lt[i];tid[j]=li[i]; break; } }
        }
        printf("\n===== dragon.c FULL FORWARD — top-5 next token IDs =====\n");
        for(int i=0;i<5;i++) printf("  id %d  logit %.2f\n",tid[i],top[i]);
    }
    return 0;
}