/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_optim_cuda.cu — parity harness for the GPU optimizer-step kernels
 * (cuda_kernels/optim.cu) vs the CPU optimizer, on GPU 0.  Phase 3, Chunk A.
 *
 *   SR4   : ternary_optim_step (ternary_optim/ternary_optim.c) — BIT-EXACT.
 *           Require identical int8 q integers (rel-err 0 / exact equality).
 *   AdamW : fo_adamw / adam_step (optim/optim.c, linked via a minimal tensor
 *           shim) — inf-norm max|gpu-cpu|/max|cpu| < 1e-3 (expect ~1e-6).
 *
 * Covers multiple shapes and several (seed, counter) / step values.  The CPU
 * oracle code is LINKED (not reimplemented): ternary_optim.o and optim.o +
 * tensor_shim.o.
 */
#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>

#include "cuda_kernels/optim.cu"

/* ---- CPU oracle: SR4 (self-contained C lib) ---- */
#include "ternary_optim/ternary_optim.h"

/* ---- CPU oracle: AdamW (real optim.c + tensor shim) ---- */
#include "tensor/tensor.h"
#include "optim/optim.h"
extern "C" tm_tensor *shim_tensor_f32(const float *src, int64_t n);

/* ---- deterministic fixture RNG (SplitMix64) ---- */
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 float sm_uf(void){ return (float)((sm_next()>>11)*(1.0/9007199254740992.0)); }
static float sm_sym(float s){ return (2.0f*sm_uf()-1.0f)*s; }
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*sizeof(float));
    for(size_t i=0;i<n;i++)p[i]=sm_sym(s); return p; }

static int g_fail=0;

/* =====================================================================
 * SR4 — bit-exact q comparison
 * ===================================================================== */
static void sr4_case(const char*name,int rows,int cols,uint64_t seed,
                     int64_t idx_base,int batch_size,uint64_t counter,
                     float target_step,int use_dither,float grad_scale){
    int numel=rows*cols;
    ternary_optim_config cfg; ternary_optim_config_default(&cfg);
    cfg.seed=seed; cfg.idx_base=idx_base; cfg.use_dither=use_dither;
    ternary_optim *opt=ternary_optim_create(rows,cols,&cfg);
    if(!opt){printf("  %s CREATE ERR\n",name);g_fail=1;return;}

    /* snapshot the identical initial q for the GPU side */
    const int8_t *q0=ternary_optim_q(opt);
    int8_t *q_gpu=(int8_t*)malloc((size_t)numel);
    memcpy(q_gpu,q0,(size_t)numel);

    sm_state=seed*0x1234567ULL+counter+7;
    float *grad=rnd((size_t)numel,grad_scale);

    /* CPU step (the oracle) */
    ternary_optim_step(opt,grad,batch_size,counter,target_step);
    const int8_t *q_cpu=ternary_optim_q(opt);

    /* GPU step from the same initial q + same inputs */
    int rc=sr4_step_cuda(q_gpu,grad,numel,batch_size,counter,target_step,
                         cfg.seed,cfg.idx_base,cfg.q_clamp,cfg.use_dither);
    if(rc){printf("  %s GPU ERR rc=%d\n",name,rc);g_fail=1;free(grad);free(q_gpu);ternary_optim_free(opt);return;}

    int mism=0,first=-1;
    for(int i=0;i<numel;i++) if(q_cpu[i]!=q_gpu[i]){ if(first<0)first=i; mism++; }
    const char*st=mism?"FAIL":"PASS"; if(mism)g_fail=1;
    printf("  %-22s rows=%d cols=%d bs=%d ctr=%llu dith=%d ts=%.4g : mism=%d/%d %s",
        name,rows,cols,batch_size,(unsigned long long)counter,use_dither,target_step,mism,numel,st);
    if(mism)printf("  (first@%d cpu=%d gpu=%d)",first,q_cpu[first],q_gpu[first]);
    printf("\n");

    free(grad);free(q_gpu);ternary_optim_free(opt);
}

/* =====================================================================
 * AdamW — inf-norm parity over several steps (exercises bias correction)
 * ===================================================================== */
static void adamw_case(const char*name,int64_t n,float lr,float beta1,float beta2,
                       float eps,float wd,int nsteps,uint64_t seed){
    sm_state=seed;
    float *init=rnd((size_t)n,0.5f);

    /* CPU oracle: real fo_adamw over nsteps */
    tm_tensor *p=shim_tensor_f32(init,n);
    p->grad=shim_tensor_f32(NULL,n);
    p->requires_grad=true;
    tm_tensor *params[1]={p};
    fo_adam_config ac=fo_adam_default(lr);
    ac.beta1=beta1; ac.beta2=beta2; ac.eps=eps; ac.weight_decay=wd;
    fo_optimizer *opt=fo_adamw_new(params,1,&ac);

    /* GPU: mirror with host-resident m/v, one launch per step */
    float *data_g=(float*)malloc((size_t)n*sizeof(float));
    memcpy(data_g,init,(size_t)n*sizeof(float));
    float *m_g=(float*)calloc((size_t)n,sizeof(float));
    float *v_g=(float*)calloc((size_t)n,sizeof(float));

    for(int s=0;s<nsteps;s++){
        float *g=rnd((size_t)n,0.3f);
        memcpy(tm_data_f32(p->grad),g,(size_t)n*sizeof(float));  /* same grad both sides */
        fo_step(opt);
        int rc=adamw_step_cuda(data_g,g,m_g,v_g,n,s+1,lr,beta1,beta2,eps,wd,1);
        if(rc){printf("  %s GPU ERR rc=%d\n",name,rc);g_fail=1;free(g);break;}
        free(g);
    }

    const float *data_c=tm_data_f32(p);
    double maxabs=0.0,scale=0.0;
    for(int64_t i=0;i<n;i++){ double d=fabs((double)data_g[i]-(double)data_c[i]);
        if(d>maxabs)maxabs=d; double c=fabs((double)data_c[i]); if(c>scale)scale=c; }
    double rel=maxabs/(scale+1e-30); const char*st=(rel<1e-3)?"PASS":"FAIL"; if(rel>=1e-3)g_fail=1;
    printf("  %-22s n=%lld steps=%d wd=%.3g : rel=%.3e maxabs=%.3e [%s]\n",
        name,(long long)n,nsteps,wd,rel,maxabs,st);

    free(init);free(data_g);free(m_g);free(v_g);
    fo_free(opt); tm_release(p->grad); tm_release(p);
}

int main(void){
    int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
    printf("=== GPU optimizer-step parity vs CPU optimizer (GPU %d: %s cc %d.%d) ===\n\n",
        dev,pr.name,pr.major,pr.minor);

    printf("--- SR4 (require BIT-EXACT q) ---\n");
    sr4_case("sr4_small",        4,   8,  42,   0, 1,   0, 0.003f, 1, 0.01f);
    sr4_case("sr4_small_ctr",    4,   8,  42,   0, 1,   5, 0.003f, 1, 0.01f);
    sr4_case("sr4_stratum2",     4,   8,  42,   0, 1, 513, 0.003f, 1, 0.01f);   /* counter>>8 != 0 */
    sr4_case("sr4_nodither",     4,   8,  42,   0, 1,   5, 0.003f, 0, 0.01f);
    sr4_case("sr4_idxbase",     16,  32, 123, 4096, 1,  17, 0.003f, 1, 0.02f);
    sr4_case("sr4_bs4",         16,  32,  42,   0, 4,   3, 0.003f, 1, 0.05f);
    sr4_case("sr4_bigstep",     32,  64,   7,   0, 1,   9, 0.05f,  1, 0.001f);  /* large prop -> saturate */
    sr4_case("sr4_mid",        128, 256,  99, 8192, 1,  40, 0.003f, 1, 0.02f);
    sr4_case("sr4_real768",    768, 768,  55,   0, 1, 128, 0.003f, 1, 0.015f);
    sr4_case("sr4_real_ctr300",768, 768,  55,   0, 1, 300, 0.003f, 1, 0.015f);

    printf("\n--- AdamW (inf-norm < 1e-3) ---\n");
    adamw_case("adamw_small",     64,   1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f,  1, 11);
    adamw_case("adamw_t5",        64,   1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f,  5, 12);
    adamw_case("adamw_t20",     4096,   1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, 20, 13);
    adamw_case("adamw_wd",      4096,   1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f,10, 14);
    adamw_case("adamw_beta",    4096,   3e-4f, 0.95f,0.98f,  1e-6f, 0.0f, 10, 15);
    adamw_case("adamw_big",   131072,   1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f,  8, 16);

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