/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_overfit.c — Training-loop smoke test for the RWKV-x from-scratch trainer,
 * with the optimizer SWAPPABLE between AdamW (all-continuous baseline) and SR4
 * (ternary projection matrices via the latent 4-bit accumulator).
 *
 * This is the FIRST proof that either optimizer drives an actual language model:
 * we overfit a single FIXED random token sequence (pure memorization, no corpus)
 * and watch the cross-entropy fall.
 *
 *   - AdamW arm : every parameter is continuous, trained by foundry optim
 *                 (fo_adamw_new + fo_step).  Expected to drive CE -> ~0
 *                 (near-perfect memorization) => proves the loop is correct.
 *   - SR4 arm   : the big linear projection matrices (attention + time-mix
 *                 receptance/key/value/output, channel-mix ffn_key/ffn_value)
 *                 are TERNARY {-1,0,+1} materialized from a per-weight int4
 *                 accumulator `q` (ternary_optim / SR4), each with a per-OUTPUT-
 *                 COLUMN fp32 scale restoring dynamic range:
 *                     W_eff[i,j] = scale[j] * ternaryW[i,j]
 *                     Z = W_eff^T x  ==  scale (*) (ternaryW^T x)     (folded)
 *                 Everything else (embedding, head, LayerNorm gamma/beta, the
 *                 RWKV-7 LoRA gate matrices, biases, lerps, AND the per-column
 *                 scales themselves) stays continuous and trains with AdamW.
 *
 * Param split is by registry NAME suffix (see is_ternary_name).
 *
 * Ternary-scale coupling (the one real subtlety):
 *   The model forward/backward only ever sees W_eff (we overwrite the model's
 *   own weight buffer with scale*ternaryW before each forward), so the analytic
 *   grad the backward produces is g = dL/dW_eff.  We then split it:
 *       grad wrt ternary weight (drives q) : gq[i,j] = g[i,j] * scale[j]
 *       grad wrt scale                     : sg[j]   = sum_i g[i,j]*ternaryW[i,j]
 *   SR4-step q with gq; AdamW-step the scale vector with sg (as a normal
 *   continuous param).
 *
 * Build: see Makefile target `test-overfit` (standalone; NOT in the default
 * `test` aggregate — it is a slow optimization loop).
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>

#include "../rx_pretrain.h"
#include "../../ternary_optim/ternary_optim.h"
#include "../../optim/optim.h"
#include "../../tensor/tensor.h"

/* ------------------------------------------------------------------ */
/* Tiny config for the overfit smoke.                                  */
/* ------------------------------------------------------------------ */
#define CFG_C        64      /* n_embd                                  */
#define CFG_HEAD     4
#define CFG_HS       16      /* head_size (HEAD*HS == C)                */
#define CFG_VOCAB    64
#define CFG_T        16      /* sequence length                         */
#define CFG_FFN      128
#define CFG_LORA     16

static rx_pt_config make_cfg(void){
    rx_pt_config pc;
    memset(&pc,0,sizeof pc);
    pc.n_embd=CFG_C; pc.n_head=CFG_HEAD; pc.head_size=CFG_HS;
    pc.n_rwkv_layer=1; pc.n_attn_layer=1;
    pc.vocab_size=CFG_VOCAB;
    pc.Dw=CFG_LORA; pc.Da=CFG_LORA; pc.Dg=CFG_LORA; pc.Dv=CFG_LORA;
    pc.ffn_dim=CFG_FFN;
    pc.chunk_size=CFG_T;          /* one chunk => full causal attention */
    pc.moba_top_k=0;              /* keep all */
    return pc;
}

/* Fixed random token sequence (input) + shifted target. length T+1. */
static void make_fixed_seq(int *tokens, int *targets, int T, int vocab){
    uint64_t s = 0xC0FFEEULL;
    int seq[CFG_T+1];
    for (int t=0;t<=T;t++){
        s = s*6364136223846793005ULL + 1442695040888963407ULL;
        seq[t] = (int)((s>>33) % (uint64_t)vocab);
    }
    for (int t=0;t<T;t++){ tokens[t]=seq[t]; targets[t]=seq[t+1]; }
}

/* ------------------------------------------------------------------ */
/* Param classification: the big projection matrices go ternary.       */
/* ------------------------------------------------------------------ */
static const char *suffix_after_dot(const char *name){
    const char *dot = strrchr(name,'.');
    return dot ? dot+1 : name;
}
static int is_ternary_name(const char *name){
    const char *s = suffix_after_dot(name);
    return  !strcmp(s,"receptance_w") || !strcmp(s,"key_w") ||
            !strcmp(s,"value_w")      || !strcmp(s,"output_w") ||
            !strcmp(s,"ffn_key_w")    || !strcmp(s,"ffn_value_w");
}
/* rows of a [in,out] row-major projection so cols(=out) = numel/rows. */
static int rows_of(const char *name, int numel){
    const char *s = suffix_after_dot(name);
    if (!strcmp(s,"ffn_key_w"))   return CFG_C;    /* [C, ffn] */
    if (!strcmp(s,"ffn_value_w")) return CFG_FFN;  /* [ffn, C] */
    (void)numel;
    return CFG_C;                                  /* [C, C] square projections */
}

/* ------------------------------------------------------------------ */
/* Wrap a model float buffer + its grad buffer as a foundry tm_tensor. */
/* ------------------------------------------------------------------ */
static tm_tensor *wrap_param(float *data, float *grad, int n){
    int64_t shape[1] = { n };
    tm_tensor *t = tm_from_data(data, shape, 1, TM_F32);
    t->grad = tm_from_data(grad, shape, 1, TM_F32);
    t->requires_grad = true;
    return t;
}

/* ------------------------------------------------------------------ */
/* Cosine LR schedule.                                                 */
/* ------------------------------------------------------------------ */
static float cosine_lr(float lr0, float lr_min, int step, int total){
    if (step >= total) return lr_min;
    float c = 0.5f*(1.0f + cosf((float)M_PI * (float)step/(float)total));
    return lr_min + (lr0-lr_min)*c;
}

/* ================================================================== */
/* AdamW arm: every param continuous.                                  */
/* ================================================================== */
static double run_adamw(const rx_pt_config *pc, const int *tokens,
                        const int *targets, int T,
                        int steps, int report_every,
                        float lr0, float lr_min, int verbose){
    rx_ptmodel *m = rx_model_new_random(pc, 1234);
    rx_pt_acts *acts = rx_pt_acts_new(m, T);
    int nP = rx_ptmodel_num_params(m);

    tm_tensor **params = calloc(nP, sizeof(tm_tensor*));
    for (int i=0;i<nP;i++){
        float *d,*g; int n; const char *nm;
        rx_ptmodel_param(m,i,&nm,&d,&g,&n);
        params[i] = wrap_param(d,g,n);
    }
    fo_adam_config cfg = fo_adam_default(lr0);
    cfg.weight_decay = 0.0f;            /* pure memorization, no decay */
    fo_optimizer *opt = fo_adamw_new(params, nP, &cfg);

    double loss = 0.0;
    if (verbose) printf("  [AdamW] lr0=%.4g -> lr_min=%.4g, %d steps\n", lr0, lr_min, steps);
    for (int step=0; step<steps; step++){
        float lr = cosine_lr(lr0, lr_min, step, steps);
        fo_set_lr(opt, lr);
        rx_ptmodel_zero_grad(m);
        rx_pretrain_forward(m, tokens, T, acts);
        loss = rx_pretrain_backward(m, acts, targets, T, -1);
        fo_step(opt);
        if (verbose && (step % report_every == 0 || step==steps-1))
            printf("    step %4d  lr %.4g  CE %.6f\n", step, lr, loss);
    }

    for (int i=0;i<nP;i++){ tm_free(params[i]); }  /* releases ->grad too */
    free(params);
    fo_free(opt);
    rx_pt_acts_free(acts);
    rx_ptmodel_free(m);
    return loss;
}

/* ================================================================== */
/* SR4 arm: ternary projections via latent int4 q + per-column scale.  */
/* ================================================================== */
typedef struct {
    float *data;            /* model weight buffer (overwritten w/ W_eff) */
    float *grad;            /* dL/dW_eff (from backward)                  */
    int rows, cols, numel;
    ternary_optim *opt;     /* owns the int4 accumulator q                */
    float *tern;            /* pure ternary {-1,0,+1} [numel]             */
    float *scale;           /* per-output-column fp32 scale [cols]        */
    float *sgrad;           /* scale grad [cols]                          */
    float *gq;              /* grad wrt ternary weight = g*scale [numel]  */
    tm_tensor *scale_t;     /* scale as a continuous AdamW param          */
} tern_mat;

/* freeze_noncore: when nonzero, ONLY the ternary q accumulators and their
 * per-column scales train — emb/head/LayerNorm/LoRA/biases are held at random
 * init.  This isolates how much of the loss drop is genuinely driven by SR4 on
 * the ternary projections vs. carried by the powerful continuous emb/head. */
static double run_sr4(const rx_pt_config *pc, const int *tokens,
                      const int *targets, int T,
                      int steps, int report_every,
                      float target_step, float cont_lr, float scale_lr,
                      int freeze_noncore, int verbose){
    rx_ptmodel *m = rx_model_new_random(pc, 1234);
    rx_pt_acts *acts = rx_pt_acts_new(m, T);
    int nP = rx_ptmodel_num_params(m);

    /* Partition params: ternary matrices vs continuous. */
    tern_mat *tmat = calloc(nP, sizeof(tern_mat));
    int n_tern = 0;
    tm_tensor **cont = calloc(2*nP, sizeof(tm_tensor*));  /* + scale vectors */
    int n_cont = 0;
    int64_t idx_base = 0;

    for (int i=0;i<nP;i++){
        float *d,*g; int n; const char *nm;
        rx_ptmodel_param(m,i,&nm,&d,&g,&n);
        if (is_ternary_name(nm)){
            tern_mat *tt = &tmat[n_tern++];
            tt->data=d; tt->grad=g; tt->numel=n;
            tt->rows = rows_of(nm,n);
            tt->cols = n / tt->rows;
            ternary_optim_config oc; ternary_optim_config_default(&oc);
            oc.seed = 42; oc.idx_base = idx_base; idx_base += n;
            tt->opt  = ternary_optim_create(tt->rows, tt->cols, &oc);
            tt->tern = malloc((size_t)n*sizeof(float));
            tt->gq   = malloc((size_t)n*sizeof(float));
            tt->scale= malloc((size_t)tt->cols*sizeof(float));
            tt->sgrad= calloc((size_t)tt->cols,sizeof(float));
            for (int c=0;c<tt->cols;c++) tt->scale[c]=1.0f;
            /* scale trains as a continuous AdamW param (grad in sgrad) */
            tt->scale_t = wrap_param(tt->scale, tt->sgrad, tt->cols);
            cont[n_cont++] = tt->scale_t;
        } else if (!freeze_noncore){
            cont[n_cont++] = wrap_param(d,g,n);
        }
    }

    /* AdamW over all continuous params (incl. per-column scales). */
    fo_adam_config cfg = fo_adam_default(cont_lr);
    cfg.weight_decay = 0.0f;
    fo_optimizer *opt = fo_adamw_new(cont, n_cont, &cfg);

    if (verbose)
        printf("  [SR4] target_step=%.4g  cont_lr=%.4g  scale_lr=%.4g  "
               "(%d ternary mats, %d continuous params)\n",
               target_step, cont_lr, scale_lr, n_tern, n_cont);

    double loss = 0.0;
    for (int step=0; step<steps; step++){
        /* 1. materialize W_eff = scale (*) ternaryW into the model buffers */
        for (int t=0;t<n_tern;t++){
            tern_mat *tt=&tmat[t];
            ternary_optim_materialize(tt->opt, tt->tern);
            for (int r=0;r<tt->rows;r++)
                for (int c=0;c<tt->cols;c++)
                    tt->data[(size_t)r*tt->cols+c] =
                        tt->scale[c]*tt->tern[(size_t)r*tt->cols+c];
        }
        /* 2. forward + backward (grads land in model buffers = dL/dW_eff) */
        rx_ptmodel_zero_grad(m);
        rx_pretrain_forward(m, tokens, T, acts);
        loss = rx_pretrain_backward(m, acts, targets, T, -1);

        /* 3. split ternary grad -> {q-grad, scale-grad} */
        for (int t=0;t<n_tern;t++){
            tern_mat *tt=&tmat[t];
            for (int c=0;c<tt->cols;c++){
                double sg=0.0;
                for (int r=0;r<tt->rows;r++){
                    size_t k=(size_t)r*tt->cols+c;
                    sg += (double)tt->grad[k]*(double)tt->tern[k];
                    tt->gq[k] = tt->grad[k]*tt->scale[c];  /* dL/d ternaryW */
                }
                tt->sgrad[c]=(float)sg;                     /* dL/d scale[c] */
            }
        }
        /* 4a. SR4-step the ternary accumulators */
        for (int t=0;t<n_tern;t++)
            ternary_optim_step(tmat[t].opt, tmat[t].gq, 1,
                               (uint64_t)step, target_step);
        /* 4b. AdamW-step the continuous params (incl. scales) */
        fo_set_lr(opt, cont_lr);
        /* scale-specific LR: temporarily not separated in foundry optim; we
         * fold scale_lr in by pre-multiplying its grad relative to cont_lr. */
        if (scale_lr != cont_lr){
            float k = scale_lr/cont_lr;
            for (int t=0;t<n_tern;t++)
                for (int c=0;c<tmat[t].cols;c++) tmat[t].sgrad[c]*=k;
        }
        fo_step(opt);

        if (verbose && (step % report_every == 0 || step==steps-1))
            printf("    step %4d  CE %.6f\n", step, loss);
    }

    /* teardown */
    for (int i=0;i<n_cont;i++){ tm_free(cont[i]); }  /* releases ->grad too */
    for (int t=0;t<n_tern;t++){
        ternary_optim_free(tmat[t].opt);
        free(tmat[t].tern); free(tmat[t].gq);
        free(tmat[t].scale); free(tmat[t].sgrad);
    }
    free(cont); free(tmat);
    fo_free(opt);
    rx_pt_acts_free(acts);
    rx_ptmodel_free(m);
    return loss;
}

/* ================================================================== */
int main(int argc, char **argv){
    const char *mode = (argc>1)? argv[1] : "both";
    rx_pt_config pc = make_cfg();
    int tokens[CFG_T], targets[CFG_T];
    make_fixed_seq(tokens, targets, CFG_T, CFG_VOCAB);

    double ln_vocab = log((double)CFG_VOCAB);
    printf("=== RWKV-x overfit smoke ===\n");
    printf("config: C=%d heads=%d hs=%d vocab=%d T=%d ffn=%d lora=%d "
           "(1 RWKV + 1 attn layer)\n",
           CFG_C,CFG_HEAD,CFG_HS,CFG_VOCAB,CFG_T,CFG_FFN,CFG_LORA);
    printf("fixed sequence (memorization target), random-init CE ~= ln(vocab) = %.4f\n\n", ln_vocab);

    int fail = 0;

    if (!strcmp(mode,"adamw") || !strcmp(mode,"both")){
        printf("--- AdamW arm (all-continuous baseline) ---\n");
        double L = run_adamw(&pc, tokens, targets, CFG_T,
                             1200, 100, 1e-2f, 1e-4f, 1);
        printf("  => AdamW final CE = %.6f  (%s memorized)\n\n",
               L, (L<0.05)?"YES —":"NOT fully");
        if (L >= 0.05) fail = 1;
    }

    if (!strcmp(mode,"sr4") || !strcmp(mode,"both")){
        printf("--- SR4 arm (ternary projections) ---\n");
        /* coarse step-size sweep */
        float cand[] = {0.003f, 0.01f, 0.03f, 0.1f, 0.3f};
        int   ncand  = (int)(sizeof(cand)/sizeof(cand[0]));
        float best_ts = cand[0]; double best_L = 1e9;
        printf("  coarse target_step sweep (400 steps each):\n");
        for (int i=0;i<ncand;i++){
            double L = run_sr4(&pc, tokens, targets, CFG_T,
                               400, 400, cand[i], 5e-3f, 1e-2f, 0, 0);
            printf("    target_step=%-6.3g  final CE %.6f\n", cand[i], L);
            if (L < best_L){ best_L=L; best_ts=cand[i]; }
        }
        printf("  best target_step = %.3g (CE %.6f); running full curve "
               "(prescribed split: emb/head/LN/LoRA continuous):\n", best_ts, best_L);
        double L = run_sr4(&pc, tokens, targets, CFG_T,
                           1500, 100, best_ts, 5e-3f, 1e-2f, 0, 1);
        printf("  => SR4 final CE = %.6f  (best target_step %.3g)\n\n", L, best_ts);

        /* Isolation probe: freeze emb/head/LN/LoRA at random init; only SR4's
         * ternary q + per-column scales may train.  Answers "is SR4 actually
         * training the projections, or is the continuous emb/head doing it?" */
        printf("  --- isolation probe: ONLY ternary q + scales train "
               "(emb/head/LN/LoRA frozen at random init) ---\n");
        double Lf = run_sr4(&pc, tokens, targets, CFG_T,
                            1500, 100, best_ts, 5e-3f, 1e-2f, 1, 1);
        printf("  => SR4-only-projections final CE = %.6f "
               "(start ~%.3f); ternary path %s learning\n\n",
               Lf, ln_vocab, (Lf < ln_vocab-0.3)?"IS":"is NOT");
    }

    if (!strcmp(mode,"both")){
        if (fail){ printf("RESULT: FAIL — AdamW did not memorize.\n"); return 1; }
        printf("RESULT: PASS — AdamW memorized; SR4 wired and reported.\n");
    }
    return 0;
}