/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_fullmodel_cuda.cu — Phase-2 completion gate.  Parity of the WHOLE
 * device-resident RWKV-x training step (cuda_kernels/fullmodel.cu) vs the CPU
 * oracle rx_pretrain_forward / rx_pretrain_backward (rwkvx/rx_pretrain.c) on
 * GPU 0: embedding + interleaved RWKV-7 / MoBA-attention blocks + ln_out + LM
 * head + cross-entropy, forward AND backward.
 *
 * Checks the scalar loss AND every parameter gradient (embedding, every block
 * param, head) at inf-norm rel-err = max|gpu-cpu|/max|cpu| < 1e-3.  Then measures
 * the full fwd+bwd training-step throughput with WEIGHTS RESIDENT on-device
 * (uploaded once, loop without re-uploading) — tokens/sec on one K80 die.
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include "cuda_kernels/fullmodel.cu"

static uint64_t sm_state;
static uint64_t sm_next(void){ uint64_t z=(sm_state+=0x9E3779B97F4A7C15ULL);
    z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
static int sm_tok(int V){ return (int)(sm_next()%(uint64_t)V); }

static int g_fail=0;
static double chk(const char*tname,const float*gpu,const float*cpu,size_t n,double bar,int verbose){
    double maxabs=0.0,scale=0.0;
    for(size_t i=0;i<n;i++){ double d=fabs((double)gpu[i]-(double)cpu[i]); if(d>maxabs)maxabs=d;
        double c=fabs((double)cpu[i]); if(c>scale)scale=c; }
    double rel=maxabs/(scale+1e-30); const char*st=(rel<bar)?"PASS":"FAIL"; if(rel>=bar)g_fail=1;
    if(verbose||rel>=bar)
        printf("    %-22s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
    return rel;
}

static void run_parity(const char*name, rx_pt_config pc, int T, uint64_t seed, int ignore_index){
    printf("\n=== parity [%s]  C=%d H=%d N=%d rwkv=%d attn=%d V=%d ffn=%d T=%d chunk=%d topk=%d ===\n",
        name,pc.n_embd,pc.n_head,pc.head_size,pc.n_rwkv_layer,pc.n_attn_layer,pc.vocab_size,
        pc.ffn_dim,T,pc.chunk_size,pc.moba_top_k);
    rx_ptmodel*cm=rx_model_new_random(&pc,seed);
    if(!cm){ printf("  model build FAIL\n"); g_fail=1; return; }
    int*tokens=(int*)malloc((size_t)T*sizeof(int)),*targets=(int*)malloc((size_t)T*sizeof(int));
    sm_state=seed^0xABCDEF;
    for(int t=0;t<T;t++){ tokens[t]=sm_tok(pc.vocab_size); targets[t]=sm_tok(pc.vocab_size); }

    rx_pt_acts*acts=rx_pt_acts_new(cm,T);
    if(rx_pretrain_forward(cm,tokens,T,acts)){ printf("  CPU fwd FAIL\n"); g_fail=1; return; }
    double loss_c=rx_pretrain_backward(cm,acts,targets,T,ignore_index);

    fm_model*m=fm_build(cm,T);
    fm_forward(m,tokens,T);
    double loss_g=fm_backward(m,targets,T,ignore_index);

    double lrel=fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30);
    printf("  loss  cpu=%.8f gpu=%.8f  rel=%.3e [%s]\n",loss_c,loss_g,lrel,(lrel<1e-3)?"PASS":"FAIL");
    if(lrel>=1e-3) g_fail=1;

    /* compare every param grad by name */
    double worst=0.0; char worstn[48]="";
    int np=rx_ptmodel_num_params(cm);
    for(int i=0;i<np;i++){
        const char*nm; float*data,*grad; int n;
        rx_ptmodel_param(cm,i,&nm,&data,&grad,&n);
        int gn; float*dg=fm_gfind(m,nm,&gn);
        if(!dg){ printf("    %-22s NO GPU GRAD (name mismatch!)\n",nm); g_fail=1; continue; }
        if(gn!=n){ printf("    %-22s SIZE MISMATCH cpu=%d gpu=%d\n",nm,n,gn); g_fail=1; continue; }
        float*hg=(float*)malloc((size_t)n*sizeof(float));
        cudaMemcpy(hg,dg,(size_t)n*sizeof(float),cudaMemcpyDeviceToHost);
        double rel=chk(nm,hg,grad,(size_t)n,1e-3,0);
        if(rel>worst){ worst=rel; snprintf(worstn,48,"%s",nm); }
        free(hg);
    }
    printf("  worst param-grad rel=%.3e (%s)   [%d params checked]\n",worst,worstn,np);

    free(tokens);free(targets);rx_pt_acts_free(acts);rx_ptmodel_free(cm);
    /* (device model left allocated; process exits soon) */
}

static void run_throughput(rx_pt_config pc, int T, int iters, uint64_t seed){
    printf("\n=== throughput (full fwd+bwd, WEIGHTS RESIDENT, one K80 die) ===\n");
    printf("  config: C=%d H=%d N=%d rwkv=%d attn=%d V=%d ffn=%d T=%d B=1 chunk=%d topk=%d\n",
        pc.n_embd,pc.n_head,pc.head_size,pc.n_rwkv_layer,pc.n_attn_layer,pc.vocab_size,
        pc.ffn_dim,T,pc.chunk_size,pc.moba_top_k);
    rx_ptmodel*cm=rx_model_new_random(&pc,seed);
    int*tokens=(int*)malloc((size_t)T*sizeof(int)),*targets=(int*)malloc((size_t)T*sizeof(int));
    sm_state=seed;
    for(int t=0;t<T;t++){ tokens[t]=sm_tok(pc.vocab_size); targets[t]=sm_tok(pc.vocab_size); }
    fm_model*m=fm_build(cm,T);   /* upload ALL weights ONCE (resident) */
    /* warmup */
    fm_forward(m,tokens,T); fm_backward(m,targets,T,-1);
    cudaDeviceSynchronize();
    struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
    for(int it=0;it<iters;it++){          /* NO re-upload of weights inside the loop */
        fm_forward(m,tokens,T);
        fm_backward(m,targets,T,-1);
    }
    cudaDeviceSynchronize(); clock_gettime(CLOCK_MONOTONIC,&t1);
    double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9;
    double toks=(double)T*iters;
    printf("  %d steps in %.4fs  ->  %.1f tokens/sec  (full fwd+bwd, resident weights)\n",
        iters,sec,toks/sec);
    printf("  (per-step %.2f ms; %d params)\n",sec/iters*1e3,rx_ptmodel_num_params(cm));
    free(tokens);free(targets);rx_ptmodel_free(cm);
}

int main(void){
    int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
    printf("=== FULL RWKV-x training step — GPU parity vs CPU oracle (Phase-2 gate) ===\n");
    printf("device %d: %s cc %d.%d\n",dev,pr.name,pr.major,pr.minor);

    /* small config: 3 RWKV + 1 attention (interval=3), N=64, small vocab/T. */
    rx_pt_config pc=(rx_pt_config){ .n_embd=128,.n_head=2,.head_size=64,
        .n_rwkv_layer=3,.n_attn_layer=1,.vocab_size=64,
        .Dw=32,.Da=32,.Dg=32,.Dv=32,.ffn_dim=256,.chunk_size=4,.moba_top_k=0 };
    run_parity("dense causal", pc, 16, 12345, -1);

    /* genuine MoBA-sparse selection active (topk<ncand) */
    pc.moba_top_k=1;
    run_parity("moba sparse", pc, 16, 22222, -1);

    /* ignore-index in the CE (mask some targets) — dense */
    pc.moba_top_k=0;
    /* larger throughput config (still one die, B=1) */
    rx_pt_config tc=(rx_pt_config){ .n_embd=512,.n_head=8,.head_size=64,
        .n_rwkv_layer=6,.n_attn_layer=2,.vocab_size=8000,
        .Dw=64,.Da=64,.Dg=128,.Dv=64,.ffn_dim=2048,.chunk_size=64,.moba_top_k=0 };
    run_throughput(tc, 512, 30, 777);

    printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
    return g_fail;
}