/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_block_bwd.c — Gradient-check for the FULL RWKV-7 BLOCK backward
 * (rx_block_backward), the counterpart of rwkvx.c `rwkv7_block_forward_fast`:
 *   x -> ln1 -> time-mix -> +x -> ln2 -> channel-mix -> +x1
 *
 * Method (same discipline as test_timemix_bwd.c / docs §5-§6):
 *   - A DOUBLE-PRECISION mirror of the whole block computes the FD reference
 *     loss L = sum(cc * out).  fp32 forward noise would swamp the small
 *     a/b/decay/LN grads and false-fail, so the reference is fp64.
 *   - The analytic backward runs the normal fp32 path (rx_block_backward).
 *   - Central finite differences, PER-TENSOR normalised infinity-norm metric
 *     (max|g-fd| / max|fd|).  Pass: < 1e-3 for EVERY tensor and the block dx.
 *
 * Focus of this chunk: the channel-mix params (ffn_x_k, ffn_key_w,
 * ffn_value_w), the two LayerNorms (ln1/ln2 gamma/beta) and the whole-block
 * dx (residual grad threading).  The time-mix + WKV7 params are re-checked
 * for free since the block exercises them end-to-end.
 */
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"

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

typedef struct {
    int B, T, H, N, C, Dw, Da, Dg, Dv, layer_id, F;  /* F = ffn_dim */
} cfg;

/* Every learnable tensor in a block, fp32. */
typedef struct {
    /* time-mix */
    float *x_r,*x_k,*x_v,*x_w,*x_a,*x_g;
    float *receptance_w,*key_w,*value_w,*output_w;
    float *w0,*w1,*w2, *a0,*a1,*a2, *g1,*g2, *v0,*v1,*v2;
    float *k_k,*k_a,*r_k,*ln_x_w,*ln_x_b;
    /* channel-mix */
    float *ffn_x_k,*ffn_key_w,*ffn_value_w;
    /* block LNs */
    float *ln1_w,*ln1_b,*ln2_w,*ln2_b;
} BP;

static float frand(void){ return (float)rand()/(float)RAND_MAX - 0.5f; }
static float *rbuf(int n, float scale){
    float *p=malloc((size_t)n*sizeof(float));
    for (int i=0;i<n;i++) p[i]=scale*frand();
    return p;
}

static void bp_alloc(BP *m, const cfg *cf){
    int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv,F=cf->F;
    m->x_r=rbuf(C,1.0f); m->x_k=rbuf(C,1.0f); m->x_v=rbuf(C,1.0f);
    m->x_w=rbuf(C,1.0f); m->x_a=rbuf(C,1.0f); m->x_g=rbuf(C,1.0f);
    m->receptance_w=rbuf(C*C,0.3f); m->key_w=rbuf(C*C,0.3f);
    m->value_w=rbuf(C*C,0.3f);      m->output_w=rbuf(C*C,0.3f);
    m->w0=rbuf(C,0.5f); m->w1=rbuf(C*Dw,0.3f); m->w2=rbuf(Dw*C,0.3f);
    m->a0=rbuf(C,0.5f); m->a1=rbuf(C*Da,0.3f); m->a2=rbuf(Da*C,0.3f);
    m->g1=rbuf(C*Dg,0.3f); m->g2=rbuf(Dg*C,0.3f);
    if (cf->layer_id>0){ m->v0=rbuf(C,0.5f); m->v1=rbuf(C*Dv,0.3f); m->v2=rbuf(Dv*C,0.3f); }
    else { m->v0=m->v1=m->v2=NULL; }
    m->k_k=rbuf(C,0.5f); m->k_a=rbuf(C,0.5f); m->r_k=rbuf(C,0.5f);
    m->ln_x_w=malloc((size_t)C*sizeof(float)); m->ln_x_b=rbuf(C,0.2f);
    for (int i=0;i<C;i++) m->ln_x_w[i]=1.0f+0.2f*frand();
    /* channel-mix */
    m->ffn_x_k=rbuf(C,1.0f);
    m->ffn_key_w=rbuf(C*F,0.3f);
    m->ffn_value_w=rbuf(F*C,0.3f);
    /* block LNs: gamma near 1, beta small */
    m->ln1_w=malloc((size_t)C*sizeof(float)); m->ln1_b=rbuf(C,0.2f);
    m->ln2_w=malloc((size_t)C*sizeof(float)); m->ln2_b=rbuf(C,0.2f);
    for (int i=0;i<C;i++){ m->ln1_w[i]=1.0f+0.2f*frand(); m->ln2_w[i]=1.0f+0.2f*frand(); }
}
static void bp_free(BP *m){
    float **a=(float**)m; for (size_t i=0;i<sizeof(BP)/sizeof(float*);i++) free(a[i]);
}
static void bp_to_params(const BP *m, const cfg *cf, rx_block_params *P){
    rx_tm_params *tm=&P->tm;
    tm->c=cf->C; tm->n_head=cf->H; tm->head_size=cf->N;
    tm->Dw=cf->Dw; tm->Da=cf->Da; tm->Dg=cf->Dg; tm->Dv=cf->Dv; tm->layer_id=cf->layer_id;
    tm->x_r=m->x_r;tm->x_k=m->x_k;tm->x_v=m->x_v;tm->x_w=m->x_w;tm->x_a=m->x_a;tm->x_g=m->x_g;
    tm->receptance_w=m->receptance_w;tm->key_w=m->key_w;tm->value_w=m->value_w;tm->output_w=m->output_w;
    tm->w0=m->w0;tm->w1=m->w1;tm->w2=m->w2;tm->a0=m->a0;tm->a1=m->a1;tm->a2=m->a2;
    tm->g1=m->g1;tm->g2=m->g2;tm->v0=m->v0;tm->v1=m->v1;tm->v2=m->v2;
    tm->k_k=m->k_k;tm->k_a=m->k_a;tm->r_k=m->r_k;tm->ln_x_w=m->ln_x_w;tm->ln_x_b=m->ln_x_b;
    P->cm.c=cf->C; P->cm.ffn_dim=cf->F;
    P->cm.ffn_x_k=m->ffn_x_k; P->cm.ffn_key_w=m->ffn_key_w; P->cm.ffn_value_w=m->ffn_value_w;
    P->ln1_w=m->ln1_w;P->ln1_b=m->ln1_b;P->ln2_w=m->ln2_w;P->ln2_b=m->ln2_b;
}

/* ---------- fp64 reference forward + scalar loss ---------- */
static double dsig(double z){ return 1.0/(1.0+exp(-z)); }

/* fp64 LayerNorm over C (eps=1e-5). */
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];
}

static double loss_fp64(const BP *m, const cfg *cf, const float *x,
    const float *tmxp, const float *cmxp, const float *v_first, const float *cc)
{
    int B=cf->B,T=cf->T,H=cf->H,N=cf->N,C=cf->C;
    int Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv,F=cf->F;
    double L=0.0;
    size_t tot=(size_t)B*T*C;
    double *h1=malloc(tot*sizeof(double));   /* ln1(x) */
    double *tm=malloc(tot*sizeof(double));   /* time-mix out */
    double *x1=malloc(tot*sizeof(double));
    double *h2=malloc(tot*sizeof(double));   /* ln2(x1) */
    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;
    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)),*tg=malloc(C*sizeof(double));
    double *xd=malloc(C*sizeof(double));
    double *fk=malloc((size_t)F*sizeof(double));

    /* h1 = ln1(x), per token */
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        for(int i=0;i<C;i++) xd[i]=(double)x[base+i];
        ln_fp64(xd,m->ln1_w,m->ln1_b,&h1[base],C);
    }

    /* time-mix over h1 -> tm (sequential state per batch) */
    for (int b=0;b<B;b++){
        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)b*T+t)*C;
            const double *xx=&h1[base];
            const double *prev=(t==0)?NULL:&h1[base-C];
            double tmxp_v;
            #define SHIFT(dst,lerp) for(int i=0;i<C;i++){ \
                double pv = prev?prev[i] : (tmxp?(double)tmxp[(size_t)b*C+i]:0.0); \
                dst[i]=xx[i]+(pv-xx[i])*(double)lerp[i]; } (void)tmxp_v
            SHIFT(sh,m->x_r); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->receptance_w[i*C+j]; r[j]=s;}
            SHIFT(sh,m->x_k); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->key_w[i*C+j]; k[j]=s;}
            SHIFT(shv,m->x_v);
            for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*m->value_w[i*C+j]; v[j]=s;}
            SHIFT(sh,m->x_w);
            for(int d=0;d<Dw;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->w2[d*C+j];
                double sg=dsig((double)m->w0[j]+s); w[j]=exp(-0.606531*sg);}
            SHIFT(sh,m->x_a);
            for(int d=0;d<Da;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->a2[d*C+j];
                a[j]=dsig((double)m->a0[j]+s);}
            SHIFT(sh,m->x_g);
            for(int d=0;d<Dg;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->g2[d*C+j]; g[j]=s;}
            for(int i=0;i<C;i++) kk[i]=k[i]*(double)m->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)m->k_a[i]);
            if (cf->layer_id==0){ for(int i=0;i<C;i++) vf[i]=v[i]; }
            else {
                for(int d=0;d<Dv;d++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*m->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]*m->v2[d*C+j];
                    double bl=dsig((double)m->v0[j]+s);
                    vf[j]=v[j]+((double)v_first[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)m->ln_x_w[off+i]+(double)m->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)m->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++) tg[i]=ygn[i]*g[i];
            for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=tg[i]*m->output_w[i*C+j]; tm[base+j]=s;}
            #undef SHIFT
        }
    }

    /* x1 = x + tm ; h2 = ln2(x1) */
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        for(int i=0;i<C;i++) x1[base+i]=(double)x[base+i]+tm[base+i];
        ln_fp64(&x1[base],m->ln2_w,m->ln2_b,&h2[base],C);
    }

    /* channel-mix over h2 -> cm ; out = x1 + cm ; L += cc*out */
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*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]:(cmxp?(double)cmxp[(size_t)b*C+i]:0.0);
            sh[i]=xx[i]+(pv-xx[i])*(double)m->ffn_x_k[i];
        }
        for(int mm=0;mm<F;mm++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->ffn_value_w[mm*C+j];
            double out=x1[base+j]+s;
            L+=(double)cc[base+j]*out;}
    }

    free(h1);free(tm);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(tg);free(xd);free(fk);
    return L;
}

/* FD-check an array (nv elems) vs analytic grad g — per-tensor rel-err. */
static float fd_array(float *arr, const float *g, int nv,
    const BP *m, const cfg *cf, const float *x, const float *tmxp,
    const float *cmxp, const float *vfst, const float *cc,
    const char *nm, int verbose)
{
    float maxdiff=0.0f, maxfd=0.0f;
    for (int e=0;e<nv;e++){
        float o=arr[e]; float h=1e-3f*(1.0f+fabsf(o));
        arr[e]=o+h; double Lp=loss_fp64(m,cf,x,tmxp,cmxp,vfst,cc);
        arr[e]=o-h; double Lm=loss_fp64(m,cf,x,tmxp,cmxp,vfst,cc);
        arr[e]=o;
        float fd=(float)((Lp-Lm)/(2.0*(double)h));
        float d=fabsf(g[e]-fd);
        if (d>maxdiff) maxdiff=d;
        if (fabsf(fd)>maxfd) maxfd=fabsf(fd);
    }
    float re=maxdiff/(maxfd+1e-9f);
    if (verbose) printf("      %-14s rel=%.2e   (max|g-fd|=%.2e max|fd|=%.2e)\n",
                        nm,re,maxdiff,maxfd);
    return re;
}

static int run_case(const char *name, cfg cf, unsigned seed){
    srand(seed);
    cf.C = cf.H*cf.N;
    int C=cf.C, B=cf.B, T=cf.T;
    size_t tot=(size_t)B*T*C;

    BP m; bp_alloc(&m,&cf);
    rx_block_params P; bp_to_params(&m,&cf,&P);

    float *x=rbuf((int)tot,1.0f);
    float *vfst=rbuf((int)tot,0.6f);   /* external v_first (layer>0) */
    float *cc=rbuf((int)tot,1.0f);
    float *dy=malloc(tot*sizeof(float));
    for (size_t i=0;i<tot;i++) dy[i]=cc[i];  /* L = sum cc*out => dy = cc */

    /* analytic grads */
    BP gm; memset(&gm,0,sizeof gm);
    #define GA(f,n) gm.f=calloc((size_t)(n),sizeof(float))
    GA(x_r,C);GA(x_k,C);GA(x_v,C);GA(x_w,C);GA(x_a,C);GA(x_g,C);
    GA(receptance_w,C*C);GA(key_w,C*C);GA(value_w,C*C);GA(output_w,C*C);
    GA(w0,C);GA(w1,C*cf.Dw);GA(w2,cf.Dw*C);
    GA(a0,C);GA(a1,C*cf.Da);GA(a2,cf.Da*C);
    GA(g1,C*cf.Dg);GA(g2,cf.Dg*C);
    if (cf.layer_id>0){ GA(v0,C);GA(v1,C*cf.Dv);GA(v2,cf.Dv*C); }
    GA(k_k,C);GA(k_a,C);GA(r_k,C);GA(ln_x_w,C);GA(ln_x_b,C);
    GA(ffn_x_k,C);GA(ffn_key_w,C*cf.F);GA(ffn_value_w,cf.F*C);
    GA(ln1_w,C);GA(ln1_b,C);GA(ln2_w,C);GA(ln2_b,C);
    #undef GA
    rx_block_grads G; memset(&G,0,sizeof G);
    rx_tm_grads *tg=&G.tm;
    tg->x_r=gm.x_r;tg->x_k=gm.x_k;tg->x_v=gm.x_v;tg->x_w=gm.x_w;tg->x_a=gm.x_a;tg->x_g=gm.x_g;
    tg->receptance_w=gm.receptance_w;tg->key_w=gm.key_w;tg->value_w=gm.value_w;tg->output_w=gm.output_w;
    tg->w0=gm.w0;tg->w1=gm.w1;tg->w2=gm.w2;tg->a0=gm.a0;tg->a1=gm.a1;tg->a2=gm.a2;
    tg->g1=gm.g1;tg->g2=gm.g2;tg->v0=gm.v0;tg->v1=gm.v1;tg->v2=gm.v2;
    tg->k_k=gm.k_k;tg->k_a=gm.k_a;tg->r_k=gm.r_k;tg->ln_x_w=gm.ln_x_w;tg->ln_x_b=gm.ln_x_b;
    G.cm.ffn_x_k=gm.ffn_x_k;G.cm.ffn_key_w=gm.ffn_key_w;G.cm.ffn_value_w=gm.ffn_value_w;
    G.ln1_w=gm.ln1_w;G.ln1_b=gm.ln1_b;G.ln2_w=gm.ln2_w;G.ln2_b=gm.ln2_b;

    float *dx=calloc(tot,sizeof(float));
    float *dvf=cf.layer_id>0?calloc(tot,sizeof(float)):NULL;
    rx_block_backward(&P,x,NULL,NULL,vfst,dy,&G,dx,NULL,NULL,dvf,B,T);

    printf("    [%s] B=%d T=%d H=%d N=%d C=%d F=%d layer=%d\n",
           name,B,T,cf.H,cf.N,C,cf.F,cf.layer_id);
    float wmax=0.0f;
    #define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,NULL,NULL,vfst,cc,#f,1); \
        if(e>wmax)wmax=e; }while(0)
    /* channel-mix + block LNs (the new pieces) */
    CK(ffn_x_k,C);CK(ffn_key_w,C*cf.F);CK(ffn_value_w,cf.F*C);
    CK(ln1_w,C);CK(ln1_b,C);CK(ln2_w,C);CK(ln2_b,C);
    /* time-mix params (re-checked end-to-end through the block) */
    CK(x_r,C);CK(x_k,C);CK(x_v,C);CK(x_w,C);CK(x_a,C);CK(x_g,C);
    CK(receptance_w,C*C);CK(key_w,C*C);CK(value_w,C*C);CK(output_w,C*C);
    CK(w0,C);CK(w1,C*cf.Dw);CK(w2,cf.Dw*C);
    CK(a0,C);CK(a1,C*cf.Da);CK(a2,cf.Da*C);
    CK(g1,C*cf.Dg);CK(g2,cf.Dg*C);
    if (cf.layer_id>0){ CK(v0,C);CK(v1,C*cf.Dv);CK(v2,cf.Dv*C); }
    CK(k_k,C);CK(k_a,C);CK(r_k,C);CK(ln_x_w,C);CK(ln_x_b,C);
    #undef CK
    { float e=fd_array(x,dx,(int)tot,&m,&cf,x,NULL,NULL,vfst,cc,"dx(block in)",1);
      if(e>wmax)wmax=e; }
    printf("      -> worst rel-err = %.2e  (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");

    bp_free(&m); bp_free(&gm);
    free(x);free(vfst);free(cc);free(dy);free(dx);free(dvf);
    return wmax<1e-3f;
}

static void test_block_backward(void){
    printf("  RWKV-7 full block backward — fp64 finite differences:\n");
    int ok=1;
    ok &= run_case("layer0 base", (cfg){.B=1,.T=5,.H=2,.N=4,.Dw=3,.Da=3,.Dg=3,.Dv=0,.layer_id=0,.F=12}, 11);
    ok &= run_case("layer1 vres", (cfg){.B=1,.T=6,.H=2,.N=4,.Dw=3,.Da=3,.Dg=3,.Dv=3,.layer_id=1,.F=16}, 23);
    ok &= run_case("batch heads", (cfg){.B=2,.T=4,.H=3,.N=6,.Dw=4,.Da=3,.Dg=5,.Dv=0,.layer_id=0,.F=20}, 37);
    ASSERT_TRUE(ok);
}

int main(void){
    printf("\nRWKV-7 full block backward gradient check\n");
    RUN_TEST(block_backward);
    TEST_REPORT();
    return __ft_failed>0 ? 1 : 0;
}