/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_attn_bwd.c — Gradient-check for the SPARSE (MoBA-style) ATTENTION BLOCK
 * backward with BASE-weight gradients (rx_attn_backward_full), the counterpart
 * of rwkvx.c `attn_block_forward_fast`:
 *   h1 = ln1(x); q,k,v = h1@{Wq,Wk,Wv};
 *   att = MoBA-sparse causal softmax-attn(q,k,v); att_out = att@Wo;
 *   x1 = x + att_out; h2 = ln2(x1); cm = channelmix(h2); out = x1 + cm
 *
 * Method (same discipline as test_block_bwd.c / docs §5-§6):
 *   - A DOUBLE-PRECISION mirror of the whole attention block computes the FD
 *     reference loss L = sum(cc * out).  fp32 forward noise would swamp the
 *     small LN grads and false-fail; the reference is fp64.  The double mirror
 *     accumulates the softmax + LN reductions in double (cancellation-prone).
 *   - The analytic backward runs the normal fp32 path (rx_attn_backward_full).
 *   - Central finite differences, PER-TENSOR normalised infinity-norm metric
 *     (max|g-fd| / max|fd|).  Pass: < 1e-3 for EVERY tensor and the block dx.
 *
 * TOP-K SELECTION (non-differentiable): the mirror recomputes the MoBA chunk
 * selection at the perturbed weights exactly as the kernel does at the base
 * weights.  The analytic backward treats the base selection as FIXED (keys of
 * unselected chunks carry weight 0 and vanish from the reductions).  These
 * agree ONLY while a perturbation does not flip the selection, so the FD test
 * is run in a selection-STABLE regime:
 *   - cases "allsel*" set moba_top_k >= #candidate chunks => all chunks always
 *     selected => plain causal attention => selection can never flip.
 *   - case "sparse" genuinely drops chunks (moba_top_k < #candidates); with a
 *     well-separated random init the gate ordering is stable under the tiny
 *     central-difference step, so it still passes.  This is the case that
 *     exercises the true sparse (masked) backward.
 */
#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, F;        /* C=H*N, F=ffn_dim */
    int chunk_size, moba_top_k;  /* MoBA config */
} cfg;

/* Every learnable tensor in the attention block, fp32. */
typedef struct {
    float *ln1_w,*ln1_b;
    float *receptance_w,*key_w,*value_w,*output_w;
    float *ln2_w,*ln2_b;
    float *ffn_x_k,*ffn_key_w,*ffn_value_w;
} AP;

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 ap_alloc(AP *m, const cfg *cf){
    int C=cf->C,F=cf->F;
    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(); }
    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->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);
}
static void ap_free(AP *m){
    float **a=(float**)m; for (size_t i=0;i<sizeof(AP)/sizeof(float*);i++) free(a[i]);
}
static void ap_to_params(const AP *m, const cfg *cf, rx_attn_params *P){
    P->c=cf->C; P->n_head=cf->H; P->head_size=cf->N;
    P->chunk_size=cf->chunk_size; P->moba_top_k=cf->moba_top_k;
    P->ln1_w=m->ln1_w; P->ln1_b=m->ln1_b;
    P->receptance_w=m->receptance_w; P->key_w=m->key_w;
    P->value_w=m->value_w; P->output_w=m->output_w;
    P->ln2_w=m->ln2_w; P->ln2_b=m->ln2_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;
}

/* ---------- fp64 reference forward + scalar loss ---------- */
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 AP *m, const cfg *cf, const float *x,
    const float *cmxp, const float *cc)
{
    int B=cf->B,T=cf->T,H=cf->H,N=cf->N,C=cf->C,F=cf->F;
    int cs=cf->chunk_size<1?1:cf->chunk_size;
    double scale=1.0/sqrt((double)N);
    size_t tot=(size_t)B*T*C;
    double L=0.0;
    double *h1=malloc(tot*sizeof(double));
    double *q =malloc(tot*sizeof(double));
    double *k =malloc(tot*sizeof(double));
    double *v =malloc(tot*sizeof(double));
    double *y =malloc(tot*sizeof(double));
    double *x1=malloc(tot*sizeof(double));
    double *h2=malloc(tot*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));

    /* h1 = ln1(x); q,k,v = h1 @ {Wq,Wk,Wv} */
    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);
        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*m->receptance_w[i*C+j]; sk+=hv*m->key_w[i*C+j]; sv+=hv*m->value_w[i*C+j]; }
            q[base+j]=sq; k[base+j]=sk; v[base+j]=sv; }
    }

    /* MoBA-sparse causal attention */
    for (int b=0;b<B;b++) for (int h=0;h<H;h++){
        int off=h*N;
        for (int tq=0;tq<T;tq++){
            size_t qb=((size_t)b*T+tq)*C+off;
            int local_chunk=tq/cs, ncand=local_chunk;
            unsigned char csel[T>0?T:1];
            int keepall=(cf->moba_top_k<=0 || cf->moba_top_k>=ncand);
            if (keepall){ for(int cc=0;cc<ncand;cc++) csel[cc]=1; }
            else {
                double gate[T>0?T:1];
                for(int c2=0;c2<ncand;c2++){ double g=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)b*T+s)*C+off+j];
                        mkj/=(double)cs; g+=q[qb+j]*mkj; }
                    gate[c2]=g; csel[c2]=0; }
                for(int it=0;it<cf->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; }
            }
            double sc[T>0?T:1]; 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)b*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)b*T+s)*C+off+j]; }
                y[qb+j]=acc; }
        }
    }

    /* att_out = y @ Wo; x1 = x + att_out; 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 j=0;j<C;j++){ double s=0; for(int i=0;i<C;i++) s+=y[base+i]*m->output_w[i*C+j];
            x1[base+j]=(double)x[base+j]+s; }
        ln_fp64(&x1[base],m->ln2_w,m->ln2_b,&h2[base],C);
    }

    /* channel-mix over h2; 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(q);free(k);free(v);free(y);free(x1);free(h2);
    free(xd);free(sh);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 AP *m, const cfg *cf, const float *x, const float *cmxp,
    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,cmxp,cc);
        arr[e]=o-h; double Lm=loss_fp64(m,cf,x,cmxp,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;

    AP m; ap_alloc(&m,&cf);
    rx_attn_params P; ap_to_params(&m,&cf,&P);

    float *x=rbuf((int)tot,1.0f);
    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 */
    AP gm; memset(&gm,0,sizeof gm);
    #define GA(f,n) gm.f=calloc((size_t)(n),sizeof(float))
    GA(ln1_w,C);GA(ln1_b,C);GA(ln2_w,C);GA(ln2_b,C);
    GA(receptance_w,C*C);GA(key_w,C*C);GA(value_w,C*C);GA(output_w,C*C);
    GA(ffn_x_k,C);GA(ffn_key_w,C*cf.F);GA(ffn_value_w,cf.F*C);
    #undef GA
    rx_attn_grads G; memset(&G,0,sizeof G);
    G.ln1_w=gm.ln1_w;G.ln1_b=gm.ln1_b;G.ln2_w=gm.ln2_w;G.ln2_b=gm.ln2_b;
    G.receptance_w=gm.receptance_w;G.key_w=gm.key_w;G.value_w=gm.value_w;G.output_w=gm.output_w;
    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;

    float *dx=calloc(tot,sizeof(float));
    rx_attn_backward_full(&P,x,NULL,dy,&G,dx,NULL,B,T);

    printf("    [%s] B=%d T=%d H=%d N=%d C=%d F=%d chunk=%d topk=%d\n",
           name,B,T,cf.H,cf.N,C,cf.F,cf.chunk_size,cf.moba_top_k);
    float wmax=0.0f;
    #define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,NULL,cc,#f,1); \
        if(e>wmax)wmax=e; }while(0)
    CK(receptance_w,C*C);CK(key_w,C*C);CK(value_w,C*C);CK(output_w,C*C);
    CK(ln1_w,C);CK(ln1_b,C);CK(ln2_w,C);CK(ln2_b,C);
    CK(ffn_x_k,C);CK(ffn_key_w,C*cf.F);CK(ffn_value_w,cf.F*C);
    #undef CK
    { float e=fd_array(x,dx,(int)tot,&m,&cf,x,NULL,cc,"dx(block in)",1);
      if(e>wmax)wmax=e; }
    printf("      -> worst rel-err = %.2e  (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");

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

static void test_attn_backward(void){
    printf("  Sparse-attention block backward (base-weight grads) — fp64 FD:\n");
    int ok=1;
    /* all-selected (moba_top_k >= #candidate chunks) => full causal, selection
     * can never flip under the FD step. */
    ok &= run_case("allsel s1", (cfg){.B=1,.T=6,.H=2,.N=4,.F=12,
                                      .chunk_size=2,.moba_top_k=99}, 11);
    ok &= run_case("allsel bh", (cfg){.B=2,.T=5,.H=3,.N=6,.F=20,
                                      .chunk_size=2,.moba_top_k=99}, 23);
    ok &= run_case("allsel cs3",(cfg){.B=1,.T=7,.H=2,.N=4,.F=16,
                                      .chunk_size=3,.moba_top_k=0 /*<=0=all*/}, 37);
    /* genuinely sparse: drops chunks; passes while gate ordering is stable. */
    ok &= run_case("sparse   ", (cfg){.B=1,.T=8,.H=2,.N=4,.F=12,
                                      .chunk_size=2,.moba_top_k=1}, 41);
    ASSERT_TRUE(ok);
}

int main(void){
    printf("\nSparse-attention block backward gradient check\n");
    RUN_TEST(attn_backward);
    TEST_REPORT();
    return __ft_failed>0 ? 1 : 0;
}