/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_model_bwd.c — WHOLE-MODEL gradient check for the full RWKV-x
 * forward+backward assembly (rx_pretrain_forward / rx_pretrain_backward).
 * This is the integration proof: embedding -> interleaved RWKV-7 + sparse-
 * attention blocks -> final LayerNorm -> head -> cross-entropy.
 *
 * Method (same discipline as test_block_bwd.c / test_attn_bwd.c):
 *   - A DOUBLE-PRECISION mirror of the ENTIRE model forward -> CE loss is the
 *     finite-difference reference (fp32 forward noise would swamp the small LN
 *     / decay / a-b grads and false-fail).  Every reduction (LN, softmax,
 *     log-sum-exp, WKV state) is accumulated in double.
 *   - The analytic backward runs the normal fp32 path (rx_pretrain_backward).
 *   - Central finite differences, PER-TENSOR normalised infinity-norm metric
 *     (max|g-fd| / max|fd|).  Pass: < 1e-3 for EVERY parameter tensor.
 *   - Large matrices are element-SUBSAMPLED (evenly spaced, capped) to keep the
 *     whole-model FD tractable; small tensors are checked exhaustively.
 *
 * Toy configs carry BOTH block types (RWKV-x interleave) and, in the 2-RWKV
 * case, the v_first coupling (block-0 produces it, block-1 consumes it).
 * MoBA moba_top_k is set >= #candidate chunks => full causal attention, which
 * is also what foundry's rx_forward runs, so the assembly can be compared to it
 * (see test_model_forward_matches_rx_forward, guarded by RX_LINK_ENGINE).
 */
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#ifdef RX_LINK_ENGINE
#include "rwkvx/rwkvx.h"
/* Bridge implemented in rwkvx.c: build a real rx_model from the ptmodel. */
rx_model* rx_model_from_ptmodel(rx_ptmodel *pt);
#endif

/* ======================= fp64 model mirror ======================= */
static double dsig(double z){ return 1.0/(1.0+exp(-z)); }

static void ln_fp64(const double *x, const float *w, const float *b,
                    double *out, int C){
    double mean=0.0; for(int i=0;i<C;i++) mean+=x[i]; mean/=C;
    double var=0.0; for(int i=0;i<C;i++){double d=x[i]-mean;var+=d*d;} var/=C;
    double rstd=1.0/sqrt(var+1e-5);
    for(int i=0;i<C;i++) out[i]=(x[i]-mean)*rstd*(double)w[i]+(double)b[i];
}

/* RWKV-7 block fp64 forward: out[T*C].  vfirst[T*C] WRITTEN at layer 0 (=v),
 * READ at layer>0.  (Transcribed from test_block_bwd.c loss_fp64.) */
static void block_fwd_fp64(const rx_block_params *P, const double *x,
                           double *vfirst, double *out, int T){
    const rx_tm_params *tm=&P->tm;
    int C=tm->c,H=tm->n_head,N=tm->head_size;
    int Dw=tm->Dw,Da=tm->Da,Dg=tm->Dg,Dv=tm->Dv,F=P->cm.ffn_dim;
    int L0=(tm->layer_id==0);
    double *h1=malloc((size_t)T*C*sizeof(double));
    double *tmo=malloc((size_t)T*C*sizeof(double));
    double *x1=malloc((size_t)T*C*sizeof(double));
    double *h2=malloc((size_t)T*C*sizeof(double));
    double *S=calloc((size_t)H*N*N,sizeof(double));
    double *sh=malloc(C*sizeof(double)),*shv=malloc(C*sizeof(double));
    double *r=malloc(C*sizeof(double)),*k=malloc(C*sizeof(double)),*v=malloc(C*sizeof(double));
    double *w=malloc(C*sizeof(double)),*a=malloc(C*sizeof(double)),*g=malloc(C*sizeof(double));
    double *kk=malloc(C*sizeof(double)),*kmod=malloc(C*sizeof(double)),*vf=malloc(C*sizeof(double));
    int maxD=Dw>Da?(Dw>Dg?Dw:Dg):(Da>Dg?Da:Dg); if(Dv>maxD)maxD=Dv; if(maxD<1)maxD=1;
    double *lo=malloc((size_t)(maxD+1)*sizeof(double));
    double *lo2=malloc((size_t)((Dv>0?Dv:1))*sizeof(double));
    double *y=malloc(C*sizeof(double)),*ygn=malloc(C*sizeof(double));
    double *xd=malloc(C*sizeof(double));
    double *fk=malloc((size_t)F*sizeof(double));

    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        for(int i=0;i<C;i++) xd[i]=x[base+i];
        ln_fp64(xd,P->ln1_w,P->ln1_b,&h1[base],C);
    }
    for (int e=0;e<H*N*N;e++) S[e]=0.0;
    for (int t=0;t<T;t++){
        size_t base=(size_t)t*C;
        const double *xx=&h1[base];
        const double *prev=(t==0)?NULL:&h1[base-C];
        #define SHIFT(dst,lerp) for(int i=0;i<C;i++){ \
            double pv=prev?prev[i]:0.0; dst[i]=xx[i]+(pv-xx[i])*(double)lerp[i]; }
        SHIFT(sh,tm->x_r); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->receptance_w[i*C+j]; r[j]=s;}
        SHIFT(sh,tm->x_k); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->key_w[i*C+j]; k[j]=s;}
        SHIFT(shv,tm->x_v);
        for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*tm->value_w[i*C+j]; v[j]=s;}
        SHIFT(sh,tm->x_w);
        for(int d=0;d<Dw;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->w1[i*Dw+d]; lo[d]=tanh(s);}
        for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dw;d++)s+=lo[d]*tm->w2[d*C+j];
            double sg=dsig((double)tm->w0[j]+s); w[j]=exp(-0.606531*sg);}
        SHIFT(sh,tm->x_a);
        for(int d=0;d<Da;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->a1[i*Da+d]; lo[d]=s;}
        for(int j=0;j<C;j++){double s=0;for(int d=0;d<Da;d++)s+=lo[d]*tm->a2[d*C+j];
            a[j]=dsig((double)tm->a0[j]+s);}
        SHIFT(sh,tm->x_g);
        for(int d=0;d<Dg;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->g1[i*Dg+d]; lo[d]=dsig(s);}
        for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dg;d++)s+=lo[d]*tm->g2[d*C+j]; g[j]=s;}
        for(int i=0;i<C;i++) kk[i]=k[i]*(double)tm->k_k[i];
        for(int h=0;h<H;h++){int off=h*N;double ss=0;for(int i=0;i<N;i++)ss+=kk[off+i]*kk[off+i];
            double inv=1.0/(sqrt(ss)+1e-12); for(int i=0;i<N;i++) kk[off+i]*=inv;}
        for(int i=0;i<C;i++) kmod[i]=k[i]*(1.0+(a[i]-1.0)*(double)tm->k_a[i]);
        if (L0){ for(int i=0;i<C;i++){ vf[i]=v[i]; vfirst[base+i]=v[i]; } }
        else {
            for(int d=0;d<Dv;d++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*tm->v1[i*Dv+d]; lo2[d]=s;}
            for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dv;d++)s+=lo2[d]*tm->v2[d*C+j];
                double bl=dsig((double)tm->v0[j]+s);
                vf[j]=v[j]+(vfirst[base+j]-v[j])*bl;}
        }
        for(int h=0;h<H;h++){int off=h*N;
            for(int i=0;i<N;i++){
                double *Si=&S[(size_t)(h*N+i)*N];
                double sa=0; for(int j=0;j<N;j++) sa+=(-kk[off+j])*Si[j];
                double vi=vf[off+i], yy=0;
                for(int j=0;j<N;j++){ Si[j]=Si[j]*w[off+j]+kmod[off+j]*vi+sa*(kk[off+j]*a[off+j]);
                    yy+=Si[j]*r[off+j]; }
                y[off+i]=yy;
            }
        }
        for(int h=0;h<H;h++){int off=h*N;double mean=0;for(int i=0;i<N;i++)mean+=y[off+i];mean/=N;
            double var=0;for(int i=0;i<N;i++){double d=y[off+i]-mean;var+=d*d;}var/=N;
            double rstd=1.0/sqrt(var+64e-5);
            for(int i=0;i<N;i++) ygn[off+i]=(y[off+i]-mean)*rstd*(double)tm->ln_x_w[off+i]+(double)tm->ln_x_b[off+i];}
        for(int h=0;h<H;h++){int off=h*N;double s=0;
            for(int j=0;j<N;j++)s+=r[off+j]*kmod[off+j]*(double)tm->r_k[off+j];
            for(int j=0;j<N;j++)ygn[off+j]+=s*vf[off+j];}
        for(int i=0;i<C;i++) ygn[i]*=g[i];   /* g gate */
        for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=ygn[i]*tm->output_w[i*C+j]; tmo[base+j]=s;}
        #undef SHIFT
    }
    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        for(int i=0;i<C;i++) x1[base+i]=x[base+i]+tmo[base+i];
        ln_fp64(&x1[base],P->ln2_w,P->ln2_b,&h2[base],C);
    }
    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        const double *xx=&h2[base]; const double *prev=(t==0)?NULL:&h2[base-C];
        for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0;
            sh[i]=xx[i]+(pv-xx[i])*(double)P->cm.ffn_x_k[i]; }
        for(int mm=0;mm<F;mm++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*P->cm.ffn_key_w[i*F+mm];
            double rr=s>0.0?s:0.0; fk[mm]=rr*rr;}
        for(int j=0;j<C;j++){double s=0;for(int mm=0;mm<F;mm++)s+=fk[mm]*P->cm.ffn_value_w[mm*C+j];
            out[base+j]=x1[base+j]+s;}
    }
    free(h1);free(tmo);free(x1);free(h2);free(S);free(sh);free(shv);
    free(r);free(k);free(v);free(w);free(a);free(g);free(kk);free(kmod);free(vf);
    free(lo);free(lo2);free(y);free(ygn);free(xd);free(fk);
}

/* Sparse-attention block fp64 forward: out[T*C].  (Transcribed from
 * test_attn_bwd.c loss_fp64; full causal when moba_top_k >= #candidates.) */
static void attn_fwd_fp64(const rx_attn_params *P, const double *x,
                          double *out, int T){
    int C=P->c,H=P->n_head,N=P->head_size,F=P->cm.ffn_dim;
    int cs=P->chunk_size<1?1:P->chunk_size;
    double scale=1.0/sqrt((double)N);
    double *h1=malloc((size_t)T*C*sizeof(double));
    double *q=malloc((size_t)T*C*sizeof(double));
    double *k=malloc((size_t)T*C*sizeof(double));
    double *v=malloc((size_t)T*C*sizeof(double));
    double *y=malloc((size_t)T*C*sizeof(double));
    double *x1=malloc((size_t)T*C*sizeof(double));
    double *h2=malloc((size_t)T*C*sizeof(double));
    double *xd=malloc((size_t)C*sizeof(double));
    double *sh=malloc((size_t)C*sizeof(double));
    double *fk=malloc((size_t)F*sizeof(double));

    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        for(int i=0;i<C;i++) xd[i]=x[base+i];
        ln_fp64(xd,P->ln1_w,P->ln1_b,&h1[base],C);
        for(int j=0;j<C;j++){ double sq=0,sk=0,sv=0;
            for(int i=0;i<C;i++){ double hv=h1[base+i];
                sq+=hv*P->receptance_w[i*C+j]; sk+=hv*P->key_w[i*C+j]; sv+=hv*P->value_w[i*C+j]; }
            q[base+j]=sq; k[base+j]=sk; v[base+j]=sv; }
    }
    for (int h=0;h<H;h++){ int off=h*N;
        for (int tq=0;tq<T;tq++){
            size_t qb=(size_t)tq*C+off;
            int local_chunk=tq/cs, ncand=local_chunk;
            unsigned char *csel=calloc((size_t)(T>0?T:1),1);
            int keepall=(P->moba_top_k<=0 || P->moba_top_k>=ncand);
            if (keepall){ for(int cc=0;cc<ncand;cc++) csel[cc]=1; }
            else {
                double *gate=malloc((size_t)(T>0?T:1)*sizeof(double));
                for(int c2=0;c2<ncand;c2++){ double gg=0;
                    for(int j=0;j<N;j++){ double mkj=0;
                        for(int s=c2*cs;s<(c2+1)*cs;s++) mkj+=k[(size_t)s*C+off+j];
                        mkj/=(double)cs; gg+=q[qb+j]*mkj; }
                    gate[c2]=gg; }
                for(int it=0;it<P->moba_top_k;it++){ int arg=-1; double best=-1e300;
                    for(int c2=0;c2<ncand;c2++) if(!csel[c2]&&gate[c2]>best){best=gate[c2];arg=c2;}
                    if(arg<0)break; csel[arg]=1; }
                free(gate);
            }
            double *sc=malloc((size_t)(T>0?T:1)*sizeof(double)); double maxs=-1e300;
            for(int s=0;s<=tq;s++){ int c2=s/cs;
                int inc=(c2==local_chunk)||(c2<ncand&&csel[c2]);
                if(!inc){ sc[s]=-1e300; continue; }
                double d=0; for(int j=0;j<N;j++) d+=q[qb+j]*k[(size_t)s*C+off+j];
                d*=scale; sc[s]=d; if(d>maxs)maxs=d; }
            double sum=0;
            for(int s=0;s<=tq;s++){ if(sc[s]<=-1e299)continue; double e=exp(sc[s]-maxs); sc[s]=e; sum+=e; }
            for(int j=0;j<N;j++){ double acc=0;
                for(int s=0;s<=tq;s++){ if(sc[s]<=-1e299)continue;
                    acc+=(sc[s]/sum)*v[(size_t)s*C+off+j]; }
                y[qb+j]=acc; }
            free(sc); free(csel);
        }
    }
    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        for(int j=0;j<C;j++){ double s=0; for(int i=0;i<C;i++) s+=y[base+i]*P->output_w[i*C+j];
            x1[base+j]=x[base+j]+s; }
        ln_fp64(&x1[base],P->ln2_w,P->ln2_b,&h2[base],C);
    }
    for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
        const double *xx=&h2[base]; const double *prev=(t==0)?NULL:&h2[base-C];
        for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0;
            sh[i]=xx[i]+(pv-xx[i])*(double)P->cm.ffn_x_k[i]; }
        for(int mm=0;mm<F;mm++){ double s=0; for(int i=0;i<C;i++) s+=sh[i]*P->cm.ffn_key_w[i*F+mm];
            double rr=s>0.0?s:0.0; fk[mm]=rr*rr; }
        for(int j=0;j<C;j++){ double s=0; for(int mm=0;mm<F;mm++) s+=fk[mm]*P->cm.ffn_value_w[mm*C+j];
            out[base+j]=x1[base+j]+s; }
    }
    free(h1);free(q);free(k);free(v);free(y);free(x1);free(h2);free(xd);free(sh);free(fk);
}

/* Whole-model fp64 forward -> mean CE loss.  Reads (possibly perturbed) weights
 * straight from the model's parameter views.  logits_out optional [T*vocab]. */
static double model_loss_fp64(rx_ptmodel *m, const int *tokens,
                              const int *targets, int T, int ignore,
                              double *logits_out){
    rx_pt_config pc=rx_ptmodel_config(m);
    int C=pc.n_embd, V=pc.vocab_size;
    int n_blocks=pc.n_rwkv_layer+pc.n_attn_layer;
    /* rebuild the interleave order (same rule as rx_model_new_random) */
    int *btype=malloc(n_blocks*sizeof(int)), *bidx=malloc(n_blocks*sizeof(int));
    if (pc.n_attn_layer==0){ for(int i=0;i<n_blocks;i++){btype[i]=0;bidx[i]=i;} }
    else { int interval=pc.n_rwkv_layer/pc.n_attn_layer, pos=0;
        for(int a=0;a<pc.n_attn_layer;a++){ int st=a*interval,en=(a+1)*interval;
            for(int j=st;j<en;j++){btype[pos]=0;bidx[pos]=j;pos++;}
            btype[pos]=1;bidx[pos]=a;pos++; }
        for(int j=pc.n_attn_layer*interval;j<pc.n_rwkv_layer;j++){btype[pos]=0;bidx[pos]=j;pos++;} }

    size_t TC=(size_t)T*C;
    double *x=malloc(TC*sizeof(double));
    double *tmp=malloc(TC*sizeof(double));
    double *vfirst=calloc(TC,sizeof(double));
    {   /* emb gather — read emb from the registry (name "emb") */
        float *emb=NULL;
        for (int i=0;i<rx_ptmodel_num_params(m);i++){ const char*nm; float*d;int nn;
            if (rx_ptmodel_param(m,i,&nm,&d,NULL,&nn)==0 && strcmp(nm,"emb")==0){emb=d;break;} }
        for (int t=0;t<T;t++) for (int i=0;i<C;i++) x[(size_t)t*C+i]=(double)emb[(size_t)tokens[t]*C+i];
    }
    for (int pos=0;pos<n_blocks;pos++){
        if (btype[pos]==0) block_fwd_fp64(rx_ptmodel_rwkv(m,bidx[pos]), x, vfirst, tmp, T);
        else               attn_fwd_fp64(rx_ptmodel_attn(m,bidx[pos]), x, tmp, T);
        memcpy(x,tmp,TC*sizeof(double));
    }
    /* ln_out + head + CE */
    float *lnw=NULL,*lnb=NULL,*head=NULL;
    for (int i=0;i<rx_ptmodel_num_params(m);i++){ const char*nm; float*d;int nn;
        if (rx_ptmodel_param(m,i,&nm,&d,NULL,&nn)!=0) continue;
        if (!strcmp(nm,"ln_out_w")) lnw=d; else if(!strcmp(nm,"ln_out_b")) lnb=d;
        else if(!strcmp(nm,"head")) head=d; }
    double *h=malloc(C*sizeof(double));
    double loss=0.0; int nv=0;
    for (int t=0;t<T;t++){
        ln_fp64(&x[(size_t)t*C],lnw,lnb,h,C);
        double *lg = logits_out ? &logits_out[(size_t)t*V] : malloc(V*sizeof(double));
        for (int j=0;j<V;j++){ double s=0; for(int i=0;i<C;i++) s+=h[i]*(double)head[(size_t)i*V+j]; lg[j]=s; }
        int tg=targets[t];
        if (!(ignore>=0 && tg==ignore)){
            double mx=-1e300; for(int j=0;j<V;j++) if(lg[j]>mx)mx=lg[j];
            double se=0; for(int j=0;j<V;j++) se+=exp(lg[j]-mx);
            loss += (mx+log(se)) - lg[tg]; nv++;
        }
        if (!logits_out) free(lg);
    }
    free(x);free(tmp);free(vfirst);free(h);free(btype);free(bidx);
    return nv>0?loss/nv:0.0;
}

/* ======================= FD driver ======================= */
#define FD_CAP 40   /* max element samples per tensor */

static float fd_tensor(rx_ptmodel *m, const int *tok, const int *tgt, int T,
                       int ignore, float *data, const float *grad, int n,
                       const char *nm, int verbose){
    int cap = n<FD_CAP ? n : FD_CAP;
    int step = n/cap; if (step<1) step=1;
    float maxdiff=0.0f, maxfd=0.0f, maxg=0.0f;
    for (int c=0;c<cap;c++){
        int e=c*step; if (e>=n) break;
        float o=data[e]; float hh=1e-3f*(1.0f+fabsf(o));
        data[e]=o+hh; double Lp=model_loss_fp64(m,tok,tgt,T,ignore,NULL);
        data[e]=o-hh; double Lm=model_loss_fp64(m,tok,tgt,T,ignore,NULL);
        data[e]=o;
        float fd=(float)((Lp-Lm)/(2.0*(double)hh));
        float d=fabsf(grad[e]-fd);
        if (d>maxdiff) maxdiff=d;
        if (fabsf(fd)>maxfd) maxfd=fabsf(fd);
        if (fabsf(grad[e])>maxg) maxg=fabsf(grad[e]);
    }
    float re=maxdiff/(maxfd+1e-9f);
    if (verbose) printf("      %-22s rel=%.2e  (max|g-fd|=%.2e max|fd|=%.2e)\n",
                        nm,re,maxdiff,maxfd);
    return re;
}

static int run_case(const char *name, rx_pt_config pc, uint64_t seed, int T){
    printf("    [%s] C=%d H=%d N=%d rwkv=%d attn=%d vocab=%d F=%d "
           "Dw=%d Da=%d Dg=%d Dv=%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,pc.Dw,pc.Da,pc.Dg,pc.Dv,T,pc.chunk_size,pc.moba_top_k);

    rx_ptmodel *m=rx_model_new_random(&pc,seed);
    if (!m){ printf("      model_new_random FAILED\n"); return 0; }

    /* tokens + targets */
    int *tok=malloc(T*sizeof(int)), *tgt=malloc(T*sizeof(int));
    srand((unsigned)seed);
    for (int t=0;t<T;t++){ tok[t]=rand()%pc.vocab_size; tgt[t]=rand()%pc.vocab_size; }

    rx_pt_acts *acts=rx_pt_acts_new(m,T);
    if (rx_pretrain_forward(m,tok,T,acts)!=0){ printf("      forward alloc FAIL\n"); return 0; }

    /* fp32 forward vs fp64 mirror logits (assembly self-consistency) */
    double *lg64=malloc((size_t)T*pc.vocab_size*sizeof(double));
    double loss64=model_loss_fp64(m,tok,tgt,T,-1,lg64);
    float maxlog=0.0f, reflog=0.0f;
    for (int t=0;t<T;t++) for (int j=0;j<pc.vocab_size;j++){
        float a=acts->logits[(size_t)t*pc.vocab_size+j];
        float d=fabsf(a-(float)lg64[(size_t)t*pc.vocab_size+j]);
        if (d>maxlog) maxlog=d;
        if (fabsf(a)>reflog) reflog=fabsf(a);
    }
    printf("      fp32-vs-fp64 logits: rel=%.2e (max|d|=%.2e max|logit|=%.2e)\n",
           maxlog/(reflog+1e-9f), maxlog, reflog);

    /* sanity: the assembly forward must match foundry's rx_forward on identical
     * weights (full-causal attention == topk>=candidates). */
    int fwd_ok = 1;
#ifdef RX_LINK_ENGINE
    {
        rx_model *rm = rx_model_from_ptmodel(m);
        rx_state *st = rx_state_new(rm);
        float maxd=0.0f, maxref=0.0f;
        for (int t=0;t<T;t++){
            tm_tensor *lg = rx_forward(rm, (uint32_t)tok[t], st);
            const float *ld = tm_data_f32(lg);
            for (int j=0;j<pc.vocab_size;j++){
                float a=acts->logits[(size_t)t*pc.vocab_size+j];
                float d=fabsf(a-ld[j]); if(d>maxd)maxd=d;
                if(fabsf(a)>maxref)maxref=fabsf(a);
            }
            tm_free(lg);
        }
        rx_state_free(st); rx_model_free(rm);
        float rel=maxd/(maxref+1e-9f);
        printf("      assembly-vs-rx_forward logits: rel=%.2e (max|d|=%.2e)\n", rel, maxd);
        fwd_ok = rel < 5e-3f;
        if (!fwd_ok) printf("      rx_forward MISMATCH\n");
    }
#else
    printf("      assembly-vs-rx_forward: (skipped; build with RX_LINK_ENGINE)\n");
#endif

    /* analytic backward */
    double loss=rx_pretrain_backward(m,acts,tgt,T,-1);
    int finite = isfinite(loss) && isfinite(loss64);
    printf("      loss(fp32 CE)=%.6f  loss(fp64 mirror)=%.6f  finite=%s\n",
           loss, loss64, finite?"yes":"NO");

    /* per-tensor FD check */
    float wmax=0.0f; const char *wname="";
    int np=rx_ptmodel_num_params(m);
    for (int i=0;i<np;i++){
        const char *nm; float *data,*grad; int n;
        rx_ptmodel_param(m,i,&nm,&data,&grad,&n);
        /* verbose print only for representative tensors to keep output readable */
        int rep = (strstr(nm,"emb")||strstr(nm,"head")||strstr(nm,"ln_out")||
                   !strncmp(nm,"rwkv0.",6)||!strncmp(nm,"attn0.",6));
        float e=fd_tensor(m,tok,tgt,T,-1,data,grad,n,nm,rep);
        if (e>wmax){ wmax=e; wname=nm; }
    }
    int pass = wmax<1e-3f && finite && fwd_ok;
    printf("      -> worst rel-err = %.2e (%s)   %s\n", wmax, wname, pass?"PASS":"FAIL");

    free(lg64); free(tok); free(tgt);
    rx_pt_acts_free(acts); rx_ptmodel_free(m);
    return pass;
}

static void test_model_backward(void){
    printf("  Whole-model RWKV-x forward+backward — fp64 finite differences:\n");
    int ok=1;
    /* Toy A: 1 RWKV + 1 sparse-attention (the minimal RWKV-x interleave). */
    ok &= run_case("1R+1A", (rx_pt_config){
        .n_embd=64,.n_head=2,.head_size=32,.n_rwkv_layer=1,.n_attn_layer=1,
        .vocab_size=48,.Dw=8,.Da=8,.Dg=8,.Dv=8,.ffn_dim=32,
        .chunk_size=2,.moba_top_k=999}, 1234, 7);
    /* Toy B: 2 RWKV + 1 sparse-attention — exercises v_first coupling
     * (block-0 produces, block-1 consumes) AND the interleave [R0,R1,A0]. */
    ok &= run_case("2R+1A", (rx_pt_config){
        .n_embd=48,.n_head=3,.head_size=16,.n_rwkv_layer=2,.n_attn_layer=1,
        .vocab_size=40,.Dw=6,.Da=6,.Dg=6,.Dv=6,.ffn_dim=24,
        .chunk_size=2,.moba_top_k=999}, 777, 6);
    ASSERT_TRUE(ok);
}

int main(void){
    printf("\nWhole-model RWKV-x gradient check (from-scratch trainer)\n");
    RUN_TEST(model_backward);
    TEST_REPORT();
    return __ft_failed>0 ? 1 : 0;
}