/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_embed_ce_cuda.cu — isolated parity for the three new leaf kernels of the
 * full-model step, vs the CPU oracle (rwkvx/rx_pretrain.c), on GPU 0:
 *   - embedding gather fwd + scatter-add bwd  (rx_embedding_forward/backward)
 *   - LM head fwd + bwd (cuBLAS)              (rx_head_forward/backward)
 *   - fused log-softmax cross-entropy fwd/bwd (rx_cross_entropy_*)
 * Reuses the device kernels from cuda_kernels/fullmodel.cu.  inf-norm bar 1e-3.
 * Small vocab so tokens REPEAT — exercises the scatter-add row collisions.
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>

#include "cuda_kernels/fullmodel.cu"

static uint64_t S; static uint64_t nx(void){ uint64_t z=(S+=0x9E3779B97F4A7C15ULL);
    z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
static float uf(void){ return (float)((nx()>>11)*(1.0/9007199254740992.0)); }
static float sy(float s){ return (2.0f*uf()-1.0f)*s; }
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=sy(s); return p; }
static int g_fail=0;
static void chk(const char*t,const float*g,const float*c,size_t n,double bar){
    double ma=0,sc=0; for(size_t i=0;i<n;i++){ double d=fabs((double)g[i]-c[i]); if(d>ma)ma=d;
        double a=fabs((double)c[i]); if(a>sc)sc=a; }
    double rel=ma/(sc+1e-30); if(rel>=bar)g_fail=1;
    printf("    %-10s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",t,rel,ma,sc,rel<bar?"PASS":"FAIL"); }

int main(void){
    cudaSetDevice(0); cublasHandle_t H; cublasCreate(&H);
    int T=24,C=48,V=16;   /* V<T so tokens repeat -> scatter-add collisions */
    size_t TC=(size_t)T*C, TV=(size_t)T*V, VC=(size_t)V*C, CV=(size_t)C*V;
    printf("=== embedding / head / cross-entropy isolated parity (GPU 0) ===\n");
    S=7;
    float*emb=rnd(VC,0.5f),*head=rnd(CV,0.3f),*x=rnd(TC,1.0f);
    int*tok=(int*)malloc(T*4),*tgt=(int*)malloc(T*4);
    for(int t=0;t<T;t++){ tok[t]=(int)(nx()%V); tgt[t]=(int)(nx()%V); }

    /* ---- embedding fwd ---- */
    float*eo_c=(float*)malloc(TC*4); rx_embedding_forward(emb,tok,T,C,eo_c);
    float *demb,*eo_g; int *dtok; cudaMalloc(&demb,VC*4);cudaMalloc(&dtok,T*4);cudaMalloc(&eo_g,TC*4);
    cudaMemcpy(demb,emb,VC*4,cudaMemcpyHostToDevice); cudaMemcpy(dtok,tok,T*4,cudaMemcpyHostToDevice);
    fm_k_emb_fwd<<<fm_nblk(TC),256>>>(demb,dtok,eo_g,T,C); cudaDeviceSynchronize();
    float*tmp=(float*)malloc(TC*4); cudaMemcpy(tmp,eo_g,TC*4,cudaMemcpyDeviceToHost);
    printf("  [embedding]\n"); chk("emb_fwd",tmp,eo_c,TC,1e-3);
    /* ---- embedding bwd (scatter-add) ---- */
    float*dout=rnd(TC,1.0f); float*demb_c=(float*)malloc(VC*4);
    rx_embedding_backward(tok,T,C,V,dout,demb_c);
    float*ddout,*dgemb; cudaMalloc(&ddout,TC*4);cudaMalloc(&dgemb,VC*4);
    cudaMemcpy(ddout,dout,TC*4,cudaMemcpyHostToDevice); cudaMemset(dgemb,0,VC*4);
    fm_k_emb_bwd<<<fm_nblk(TC),256>>>(dtok,ddout,dgemb,T,C); cudaDeviceSynchronize();
    float*dembg=(float*)malloc(VC*4); cudaMemcpy(dembg,dgemb,VC*4,cudaMemcpyDeviceToHost);
    chk("emb_bwd",dembg,demb_c,VC,1e-3);

    /* ---- head fwd/bwd ---- */
    float*lg_c=(float*)malloc(TV*4); rx_head_forward(x,head,T,C,V,lg_c);
    float*dx_c=(float*)malloc(TC*4),*dh_c=(float*)malloc(CV*4);
    float*dl=rnd(TV,1.0f); rx_head_backward(x,head,dl,T,C,V,dx_c,dh_c);
    float *dxg,*dhead,*dxdev,*dlg,*dlogit;
    cudaMalloc(&dxg,TC*4);cudaMalloc(&dhead,CV*4);cudaMalloc(&dxdev,TC*4);cudaMalloc(&dlg,TV*4);cudaMalloc(&dlogit,TV*4);
    cudaMemcpy(dxdev,x,TC*4,cudaMemcpyHostToDevice); cudaMemcpy(dhead,head,CV*4,cudaMemcpyHostToDevice);
    cudaMemcpy(dlg,dl,TV*4,cudaMemcpyHostToDevice);
    rm_gemm_tm(H,0,0,T,V,C,1.0f,dxdev,dhead,0.0f,dlogit);           /* logits */
    float*lg_g=(float*)malloc(TV*4); cudaMemcpy(lg_g,dlogit,TV*4,cudaMemcpyDeviceToHost);
    printf("  [head]\n"); chk("head_fwd",lg_g,lg_c,TV,1e-3);
    float*dheadg2; cudaMalloc(&dheadg2,CV*4);
    rm_gemm_tm(H,1,0,C,V,T,1.0f,dxdev,dlg,0.0f,dheadg2);            /* d_head=x^T@dl */
    rm_gemm_tm(H,0,1,T,C,V,1.0f,dlg,dhead,0.0f,dxg);               /* dx=dl@head^T */
    float*dxh=(float*)malloc(TC*4),*dhh=(float*)malloc(CV*4);
    cudaMemcpy(dxh,dxg,TC*4,cudaMemcpyDeviceToHost); cudaMemcpy(dhh,dheadg2,CV*4,cudaMemcpyDeviceToHost);
    chk("head_dx",dxh,dx_c,TC,1e-3); chk("head_dW",dhh,dh_c,CV,1e-3);

    /* ---- cross-entropy fwd/bwd (ignore=-1 and ignore case) ---- */
    for(int trial=0;trial<2;trial++){
        int ign = trial? tgt[3] : -1;
        float*dlc=(float*)malloc(TV*4);
        double loss_c=rx_cross_entropy_backward(lg_c,tgt,T,V,ign,dlc);
        int nv=0; for(int t=0;t<T;t++) if(!(ign>=0&&tgt[t]==ign)) nv++;
        int*dtgt; cudaMalloc(&dtgt,T*4); cudaMemcpy(dtgt,tgt,T*4,cudaMemcpyHostToDevice);
        float*dlg2; cudaMalloc(&dlg2,TV*4); double*dlt; cudaMalloc(&dlt,(size_t)T*8);
        float*dlogit_c; cudaMalloc(&dlogit_c,TV*4); cudaMemcpy(dlogit_c,lg_c,TV*4,cudaMemcpyHostToDevice);
        fm_k_ce_bwd<<<fm_nblk(T),256>>>(dlogit_c,dtgt,dlg2,dlt,T,V,ign,nv); cudaDeviceSynchronize();
        double*lth=(double*)malloc((size_t)T*8); cudaMemcpy(lth,dlt,(size_t)T*8,cudaMemcpyDeviceToHost);
        double loss_g=0; for(int t=0;t<T;t++) loss_g+=lth[t]; loss_g=nv>0?loss_g/nv:0;
        float*dlgh=(float*)malloc(TV*4); cudaMemcpy(dlgh,dlg2,TV*4,cudaMemcpyDeviceToHost);
        printf("  [cross-entropy ignore=%d nv=%d] loss cpu=%.8f gpu=%.8f rel=%.3e\n",
            ign,nv,loss_c,loss_g,fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30));
        if(fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30)>=1e-3) g_fail=1;
        chk("ce_dlogit",dlgh,dlc,TV,1e-3);
        free(dlc);free(lth);free(dlgh);
    }
    printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
    return g_fail;
}