/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_timemix_bwd.c — Gradient-check for the RWKV-7 TIME-MIX BLOCK backward
 * (rx_timemix_backward), the counterpart of rwkvx.c `rwkv7_time_mix_fast`.
 *
 * Method (docs/wkv7_backward_reference.md §5/§6, "use an fp64 reference"):
 *   - A DOUBLE-PRECISION mirror of the time-mix forward computes the FD
 *     reference loss L = sum(cc * out).  fp32 forward noise would swamp the
 *     small a/b/decay-related grads, so the reference forward is fp64.
 *   - The analytic backward runs in the normal fp32 path (rx_timemix_backward).
 *   - Central finite differences per-parameter and for the block input.
 *     Pass: max relative error < 1e-3 for EVERY parameter grad and the input.
 *
 * Prints per-parameter rel-err so a miss is localised.  Several shapes are
 * exercised (T, heads, head-dim, batch) plus a layer>0 case (value-residual
 * gating path) and a non-trivial T so token-shift + the WKV recurrence run.
 */
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"

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

/* ---------- config + parameter bag ---------- */
typedef struct {
    int B, T, H, N, C, Dw, Da, Dg, Dv, layer_id;
} cfg;

/* Every learnable time-mix tensor, fp32.  Sizes derived from cfg. */
typedef struct {
    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;
} TMP;

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 tmp_alloc(TMP *m, const cfg *cf){
    int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv;
    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);
    /* group-norm affine: weight near 1, bias small */
    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();
}
static void tmp_free(TMP *m){
    float **a=(float**)m; for (size_t i=0;i<sizeof(TMP)/sizeof(float*);i++) free(a[i]);
}
static void tmp_to_params(const TMP *m, const cfg *cf, rx_tm_params *P){
    P->c=cf->C; P->n_head=cf->H; P->head_size=cf->N;
    P->Dw=cf->Dw; P->Da=cf->Da; P->Dg=cf->Dg; P->Dv=cf->Dv; P->layer_id=cf->layer_id;
    P->x_r=m->x_r;P->x_k=m->x_k;P->x_v=m->x_v;P->x_w=m->x_w;P->x_a=m->x_a;P->x_g=m->x_g;
    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->w0=m->w0;P->w1=m->w1;P->w2=m->w2;P->a0=m->a0;P->a1=m->a1;P->a2=m->a2;
    P->g1=m->g1;P->g2=m->g2;P->v0=m->v0;P->v1=m->v1;P->v2=m->v2;
    P->k_k=m->k_k;P->k_a=m->k_a;P->r_k=m->r_k;P->ln_x_w=m->ln_x_w;P->ln_x_b=m->ln_x_b;
}

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

static double loss_fp64(const TMP *m, const cfg *cf, const float *x,
    const float *x_prev_init, 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;
    double L=0.0;
    double *S=calloc((size_t)H*N*N,sizeof(double));
    double *sh=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));
    double *lo=malloc((size_t)(Dw>Da?(Dw>Dg?Dw:Dg):(Da>Dg?Da:Dg))*sizeof(double)+64);
    double *lo2=malloc((size_t)(Dv>0?Dv:1)*sizeof(double)+64);
    double *y=malloc(C*sizeof(double)),*ygn=malloc(C*sizeof(double)),*tg=malloc(C*sizeof(double));
    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 float *xx=&x[base];
            const float *prev=(t==0)?(x_prev_init?&x_prev_init[(size_t)b*C]:NULL)
                                    :&x[(((size_t)b*T)+(t-1))*C];
            #define SHIFT(dst,lerp) for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0; \
                dst[i]=(double)xx[i]+(pv-(double)xx[i])*(double)lerp[i]; }
            /* R/K/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;}
            double shv[4096]; 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;}
            /* decay */
            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);}
            /* gate a */
            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);}
            /* gate g */
            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;}
            /* key norm (per head) */
            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;}
            /* key mod */
            for(int i=0;i<C;i++) kmod[i]=k[i]*(1.0+(a[i]-1.0)*(double)m->k_a[i]);
            /* value residual */
            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;}
            }
            /* WKV7 recurrence (double) */
            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;
                }
            }
            /* group norm */
            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];}
            /* bonus */
            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];}
            /* gate + output */
            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];
                L+=(double)cc[base+j]*s;}
            #undef SHIFT
        }
    }
    free(S);free(sh);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);
    return L;
}

/* FD-check a parameter/input array (nv elements) vs analytic grad g. */
static float fd_array(float *arr, const float *g, int nv,
    const TMP *m, const cfg *cf, const float *x, const float *xpi,
    const float *vfst, const float *cc, const char *nm, int verbose)
{
    /* Per-tensor relative error (standard gradcheck metric): the max absolute
     * analytic-vs-FD mismatch, normalised by the tensor's gradient scale.  A
     * per-ELEMENT relative error is meaningless for the many near-zero entries
     * of the C×C weight matrices (a ~1e-7 absolute mismatch over a ~1e-5 entry
     * reads as ~1e-2 while g and fd still agree to 3-4 sig figs). */
    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,xpi,vfst,cc);
        arr[e]=o-h; double Lm=loss_fp64(m,cf,x,xpi,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);
    int C=cf.C, B=cf.B, T=cf.T;
    cf.C = C = cf.H*cf.N;
    size_t tot=(size_t)B*T*C;

    TMP m; tmp_alloc(&m,&cf);
    rx_tm_params P; tmp_to_params(&m,&cf,&P);

    float *x=rbuf((int)tot,1.0f);
    float *xpi=rbuf(B*C,1.0f);
    float *vfst=rbuf((int)tot,0.6f);       /* external v_first (layer>0) */
    float *cc=rbuf((int)tot,1.0f);          /* dL/dout weights */
    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 */
    TMP 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);
    #undef GA
    rx_tm_grads G;
    G.x_r=gm.x_r;G.x_k=gm.x_k;G.x_v=gm.x_v;G.x_w=gm.x_w;G.x_a=gm.x_a;G.x_g=gm.x_g;
    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.w0=gm.w0;G.w1=gm.w1;G.w2=gm.w2;G.a0=gm.a0;G.a1=gm.a1;G.a2=gm.a2;
    G.g1=gm.g1;G.g2=gm.g2;G.v0=gm.v0;G.v1=gm.v1;G.v2=gm.v2;
    G.k_k=gm.k_k;G.k_a=gm.k_a;G.r_k=gm.r_k;G.ln_x_w=gm.ln_x_w;G.ln_x_b=gm.ln_x_b;

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

    printf("    [%s] B=%d T=%d H=%d N=%d C=%d layer=%d\n",
           name,B,T,cf.H,cf.N,C,cf.layer_id);
    float wmax=0.0f;
    #define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,xpi,vfst,cc,#f,1); \
        if(e>wmax)wmax=e; }while(0)
    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
    /* block input */
    { float e=fd_array(x,dx,(int)tot,&m,&cf,x,xpi,vfst,cc,"dx(input)",1);
      if(e>wmax)wmax=e; }
    printf("      -> worst rel-err = %.2e  (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");

    /* cleanup */
    tmp_free(&m); tmp_free(&gm);
    free(x);free(xpi);free(vfst);free(cc);free(dy);free(dx); free(dvf);
    return wmax<1e-3f;
}

static void test_timemix_backward(void){
    printf("  Time-mix block backward — fp64 finite differences:\n");
    int ok=1;
    /* layer 0, small multi-head, non-trivial T */
    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}, 11);
    /* layer>0: exercises value-residual gating path */
    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}, 23);
    /* batch>1, more heads, bigger head-dim */
    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}, 37);
    ASSERT_TRUE(ok);
}

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