/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * rx_train.c — RWKV-X LoRA training for attention blocks
 *
 * Gradient checkpointing approach: during forward pass, we save only the
 * block input (x_input) and the KV cache state at each attention block.
 * During backward, we recompute all intermediate activations using the
 * weight accessors, then compute proper gradients for all 6 LoRA layers.
 *
 * The standard inference forward (attn_block_forward_fast) handles all
 * cache management, eviction, etc. Our training forward just piggybacks
 * via the rx_forward_ex callback to save x_input and cache_len.
 *
 * LoRA layers apply an additive correction: out += scale * B @ A @ input.
 * During backward, we recompute (base + LoRA) forward, then backprop to
 * get LoRA A/B gradients. Base weight gradients are not computed (frozen).
 *
 * LoRA indices per attention block: 0=Q, 1=K, 2=V, 3=output, 4=ffn_key, 5=ffn_value
 */

#include "rwkvx/rx_train.h"
#include "rwkvx/rwkvx.h"
#include "tensor/tensor.h"
#include "lora/lora.h"
#include "optim/optim.h"
#include "data/data.h"

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/stat.h>

#ifdef USE_BLAS
  #if defined(__APPLE__)
    #include <Accelerate/Accelerate.h>
  #else
    #include <cblas.h>
  #endif
#endif

/* ================================================================
 * Internal helpers
 * ================================================================ */

/* out[cols] = x[rows] @ W[rows, cols]  (fast_matvec convention) */
static void t_matvec(const float *x, const float *W, float *out,
                      int rows, int cols) {
#ifdef USE_BLAS
    cblas_sgemv(CblasRowMajor, CblasTrans, rows, cols,
                1.0f, W, cols, x, 1, 0.0f, out, 1);
#else
    memset(out, 0, cols * sizeof(float));
    for (int i = 0; i < rows; i++) {
        float xi = x[i];
        const float *wi = &W[i * cols];
        for (int j = 0; j < cols; j++)
            out[j] += xi * wi[j];
    }
#endif
}

/* out[rows] = W[rows, cols] @ x[cols]  (W not transposed)
 * Used for backward: d_input = W @ d_output */
static void t_matvec_nt(const float *W, const float *x, float *out,
                         int rows, int cols) {
#ifdef USE_BLAS
    cblas_sgemv(CblasRowMajor, CblasNoTrans, rows, cols,
                1.0f, W, cols, x, 1, 0.0f, out, 1);
#else
    for (int i = 0; i < rows; i++) {
        float sum = 0.0f;
        for (int j = 0; j < cols; j++)
            sum += W[i * cols + j] * x[j];
        out[i] = sum;
    }
#endif
}

/* Layer norm forward */
static void t_layer_norm(const float *x, const float *w, const float *b,
                          float *out, int c) {
    float mean = 0.0f;
    for (int i = 0; i < c; i++) mean += x[i];
    mean /= c;
    float var = 0.0f;
    for (int i = 0; i < c; i++) { float d = x[i] - mean; var += d * d; }
    var /= c;
    float inv_std = 1.0f / sqrtf(var + 1e-5f);
    for (int i = 0; i < c; i++)
        out[i] = (x[i] - mean) * inv_std * w[i] + b[i];
}

/* Apply LoRA forward: out[d_out] += scale * B @ A @ input */
static void lora_fwd_add(fl_layer *lora, const float *input, float *out) {
    if (!lora) return;
    int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
    float scale = lora->scale;
    float *A = tm_data_f32(lora->lora_A);
    float *B = tm_data_f32(lora->lora_B);
    float *h = (float*)calloc(rank, sizeof(float));
    for (int r = 0; r < rank; r++) {
        float s = 0.0f;
        for (int j = 0; j < d_in; j++) s += A[r * d_in + j] * input[j];
        h[r] = s;
    }
    for (int i = 0; i < d_out; i++) {
        float s = 0.0f;
        for (int r = 0; r < rank; r++) s += B[i * rank + r] * h[r];
        out[i] += scale * s;
    }
    free(h);
}

/* Ensure gradient tensor exists */
static void ensure_grad(tm_tensor *t) {
    if (!t) return;
    if (!t->grad) {
        t->grad = tm_zeros(t->shape, t->ndim, t->dtype);
        t->grad->requires_grad = false;
    }
}

/* Accumulate LoRA gradients: dB += scale * d_out ⊗ h^T, dA += scale * (B^T @ d_out) ⊗ in^T */
static void lora_grad_acc(fl_layer *lora, const float *input,
                           const float *d_output) {
    if (!lora) return;
    int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
    float scale = lora->scale;
    float *A = tm_data_f32(lora->lora_A);
    float *B = tm_data_f32(lora->lora_B);
    ensure_grad(lora->lora_A);
    ensure_grad(lora->lora_B);
    float *dA = tm_data_f32(lora->lora_A->grad);
    float *dB = tm_data_f32(lora->lora_B->grad);

    float *h = (float*)calloc(rank, sizeof(float));
    for (int r = 0; r < rank; r++) {
        float s = 0.0f;
        for (int j = 0; j < d_in; j++) s += A[r * d_in + j] * input[j];
        h[r] = s;
    }
    for (int i = 0; i < d_out; i++) {
        float ds = scale * d_output[i];
        for (int r = 0; r < rank; r++)
            dB[i * rank + r] += ds * h[r];
    }
    float *dh = (float*)calloc(rank, sizeof(float));
    for (int r = 0; r < rank; r++) {
        float s = 0.0f;
        for (int i = 0; i < d_out; i++) s += B[i * rank + r] * d_output[i];
        dh[r] = scale * s;
    }
    for (int r = 0; r < rank; r++) {
        float dr = dh[r];
        for (int j = 0; j < d_in; j++)
            dA[r * d_in + j] += dr * input[j];
    }
    free(h);
    free(dh);
}

/* LoRA backward: d_input[d_in] += scale * A^T @ B^T @ d_output */
static void lora_backward_input(fl_layer *lora, const float *d_output,
                                  float *d_input) {
    if (!lora) return;
    int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
    float scale = lora->scale;
    float *A = tm_data_f32(lora->lora_A);
    float *B = tm_data_f32(lora->lora_B);

    float *dh = (float*)calloc(rank, sizeof(float));
    for (int r = 0; r < rank; r++) {
        float s = 0.0f;
        for (int i = 0; i < d_out; i++) s += B[i * rank + r] * d_output[i];
        dh[r] = s;
    }
    for (int r = 0; r < rank; r++) {
        float dr = scale * dh[r];
        for (int j = 0; j < d_in; j++)
            d_input[j] += A[r * d_in + j] * dr;
    }
    free(dh);
}

/* ================================================================
 * Config defaults
 * ================================================================ */

rx_train_config rx_train_config_defaults(void) {
    rx_train_config cfg;
    memset(&cfg, 0, sizeof(cfg));
    cfg.lora_rank = 32;
    cfg.lora_alpha = 64.0f;
    cfg.lr = 2e-5f;
    cfg.weight_decay = 0.01f;
    cfg.epochs = 3;
    cfg.max_seq_len = 512;
    cfg.warmup_steps = 100;
    cfg.save_every = 1000;
    cfg.log_every = 10;
    cfg.grad_clip = 1.0f;
    return cfg;
}

/* ================================================================
 * Training examples
 * ================================================================ */

void rx_train_example_free(rx_train_example *ex) {
    if (!ex) return;
    free(ex->input_ids);
    free(ex->target_ids);
}

void rx_train_examples_free(rx_train_example *examples, int n) {
    if (!examples) return;
    for (int i = 0; i < n; i++) {
        free(examples[i].input_ids);
        free(examples[i].target_ids);
    }
    free(examples);
}

rx_train_example* rx_tokenize_dataset(
    fd_dataset *ds, rx_tokenizer *tok,
    int max_seq_len, int *n_examples_out)
{
    if (!ds || !tok || max_seq_len < 2 || !n_examples_out) return NULL;
    int n_samples = fd_len(ds);
    if (n_samples == 0) { *n_examples_out = 0; return NULL; }

    rx_train_example *examples = (rx_train_example*)calloc(n_samples, sizeof(rx_train_example));
    int n_out = 0;
    int max_tokens = max_seq_len + 1;
    uint32_t *tokens = (uint32_t*)malloc(max_tokens * sizeof(uint32_t));

    for (int i = 0; i < n_samples; i++) {
        const char *text = fd_sample_get(fd_get(ds, i), "text");
        if (!text || strlen(text) < 2) continue;
        int n_tok = rx_tokenizer_encode(tok, text, tokens, max_tokens);
        if (n_tok < 2) continue;
        if (n_tok > max_seq_len + 1) n_tok = max_seq_len + 1;
        int sl = n_tok - 1;
        examples[n_out].seq_len = sl;
        examples[n_out].input_ids = (uint32_t*)malloc(sl * sizeof(uint32_t));
        examples[n_out].target_ids = (uint32_t*)malloc(sl * sizeof(uint32_t));
        memcpy(examples[n_out].input_ids, tokens, sl * sizeof(uint32_t));
        memcpy(examples[n_out].target_ids, &tokens[1], sl * sizeof(uint32_t));
        n_out++;
    }
    free(tokens);
    *n_examples_out = n_out;
    return examples;
}

/* ================================================================
 * Saved activations
 * ================================================================ */

void rx_attn_saved_free(rx_attn_saved *s) {
    if (!s) return;
    free(s->ln_out);
    free(s->q_vec);
    free(s->y_data);
    free(s->attn_weights);
    free(s->head_totals);
    free(s->x_input);
    free(s->x_after_attn);
    free(s->ffn_ln_out);
    free(s->shifted);
    free(s->fk_pre_relu);
    free(s->fk_post_relu);
}

/* ================================================================
 * LayerNorm backward
 * ================================================================ */

void rx_layer_norm_backward(
    const float *d_output,
    const float *x,
    const float *weight,
    float *d_x,
    int c)
{
    float mean = 0.0f;
    for (int i = 0; i < c; i++) mean += x[i];
    mean /= c;
    float var = 0.0f;
    for (int i = 0; i < c; i++) { float d = x[i] - mean; var += d * d; }
    var /= c;
    float inv_std = 1.0f / sqrtf(var + 1e-5f);
    float inv_c = 1.0f / c;

    float mean1 = 0.0f, mean2 = 0.0f;
    for (int i = 0; i < c; i++) {
        float dw = d_output[i] * weight[i];
        float x_hat = (x[i] - mean) * inv_std;
        mean1 += dw;
        mean2 += dw * x_hat;
    }
    mean1 *= inv_c;
    mean2 *= inv_c;

    for (int i = 0; i < c; i++) {
        float dw = d_output[i] * weight[i];
        float x_hat = (x[i] - mean) * inv_std;
        d_x[i] = inv_std * (dw - mean1 - x_hat * mean2);
    }
}

/* ================================================================
 * Forward recomputation for backward pass (gradient checkpointing)
 *
 * Given saved x_input, recompute all intermediate activations
 * needed for the backward pass. Uses weight accessors + LoRA.
 * Stores results in rx_attn_saved struct.
 *
 * This is called once per attention block per token during backward.
 * ================================================================ */

static void recompute_attn_intermediates(
    rx_model *model,
    rx_state *state,
    int attn_idx,
    fl_layer **lora,     /* [6] */
    rx_attn_saved *saved,
    int n_head, int head_size)
{
    rx_config cfg = rx_model_config(model);
    int c = cfg.n_embd;

    rx_attn_weights w;
    rx_model_get_attn_weights(model, attn_idx, &w);
    int ffn_dim = w.ffn_dim;

    /* Ensure saved has allocated buffers */
    if (!saved->ln_out) saved->ln_out = (float*)calloc(c, sizeof(float));
    if (!saved->q_vec) saved->q_vec = (float*)calloc(c, sizeof(float));
    if (!saved->y_data) saved->y_data = (float*)calloc(c, sizeof(float));
    if (!saved->x_after_attn) saved->x_after_attn = (float*)calloc(c, sizeof(float));
    if (!saved->ffn_ln_out) saved->ffn_ln_out = (float*)calloc(c, sizeof(float));
    if (!saved->shifted) saved->shifted = (float*)calloc(c, sizeof(float));
    if (!saved->fk_pre_relu) saved->fk_pre_relu = (float*)calloc(ffn_dim, sizeof(float));
    if (!saved->fk_post_relu) saved->fk_post_relu = (float*)calloc(ffn_dim, sizeof(float));

    /* Pre-attention LayerNorm */
    t_layer_norm(saved->x_input, w.ln1_w, w.ln1_b, saved->ln_out, c);

    /* Q projection with LoRA */
    t_matvec(saved->ln_out, w.receptance_w, saved->q_vec, c, c);
    lora_fwd_add(lora[0], saved->ln_out, saved->q_vec);

    /* Get K/V cache (current state — includes all tokens up to cache_len) */
    const float *k_cache, *v_cache;
    int current_cache_len = rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
    /* Use saved cache_len (state at this token's forward) */
    int ct = saved->cache_len;
    if (ct <= 0) ct = current_cache_len;  /* fallback */
    float scale = 1.0f / sqrtf((float)head_size);

    /* Per-head attention */
    float *y_data = saved->y_data;
    memset(y_data, 0, c * sizeof(float));

    /* Allocate/reuse attention weights buffer */
    int max_tokens = ct > 0 ? ct : 1;
    if (!saved->attn_weights)
        saved->attn_weights = (float*)calloc((size_t)n_head * max_tokens, sizeof(float));
    if (!saved->head_totals)
        saved->head_totals = (int*)calloc(n_head, sizeof(int));

    for (int h = 0; h < n_head; h++) {
        int off = h * head_size;
        float *scores = &saved->attn_weights[h * max_tokens];
        saved->head_totals[h] = ct;

        for (int t = 0; t < ct; t++) {
            float dot = 0.0f;
            for (int j = 0; j < head_size; j++)
                dot += saved->q_vec[off + j] * k_cache[t * c + off + j];
            scores[t] = dot * scale;
        }

        /* Softmax */
        float max_s = ct > 0 ? scores[0] : 0.0f;
        for (int t = 1; t < ct; t++)
            if (scores[t] > max_s) max_s = scores[t];
        float sum_e = 0.0f;
        for (int t = 0; t < ct; t++) {
            scores[t] = expf(scores[t] - max_s);
            sum_e += scores[t];
        }
        if (sum_e > 0.0f)
            for (int t = 0; t < ct; t++)
                scores[t] /= sum_e;

        /* Value aggregation */
        for (int j = 0; j < head_size; j++) {
            float val = 0.0f;
            for (int t = 0; t < ct; t++)
                val += scores[t] * v_cache[t * c + off + j];
            y_data[off + j] = val;
        }
    }

    /* Output projection + LoRA → attention output */
    float *att_out = (float*)calloc(c, sizeof(float));
    t_matvec(y_data, w.output_w, att_out, c, c);
    lora_fwd_add(lora[3], y_data, att_out);

    /* Attention residual: x_after_attn = x_input + att_out */
    for (int i = 0; i < c; i++)
        saved->x_after_attn[i] = saved->x_input[i] + att_out[i];

    /* Pre-FFN LayerNorm */
    t_layer_norm(saved->x_after_attn, w.ln2_w, w.ln2_b, saved->ffn_ln_out, c);

    /* FFN token shift (approximate: use ffn_ln_out as both current and prev) */
    for (int i = 0; i < c; i++)
        saved->shifted[i] = saved->ffn_ln_out[i];  /* shift is identity when prev≈current */

    /* FFN key + LoRA */
    memset(saved->fk_pre_relu, 0, ffn_dim * sizeof(float));
    t_matvec(saved->shifted, w.ffn_key_w, saved->fk_pre_relu, c, ffn_dim);
    lora_fwd_add(lora[4], saved->shifted, saved->fk_pre_relu);

    /* ReLU² */
    for (int i = 0; i < ffn_dim; i++) {
        float v = saved->fk_pre_relu[i] > 0.0f ? saved->fk_pre_relu[i] : 0.0f;
        saved->fk_post_relu[i] = v * v;
    }

    free(att_out);
}

/* ================================================================
 * Attention training forward
 *
 * This is a public function that can be called directly. For the
 * gradient checkpointing approach used in rx_train_step, it's
 * called during backward to recompute activations.
 * ================================================================ */

void rx_attn_forward_train(
    rx_model *model,
    int attn_idx,
    float *x_data,
    rx_state *state,
    fl_layer **lora,
    float *x_out,
    rx_attn_saved *saved)
{
    rx_config cfg = rx_model_config(model);
    int c = cfg.n_embd;
    int n_head = cfg.n_head;
    int head_size = cfg.head_size;

    /* Save x_input */
    if (!saved->x_input) saved->x_input = (float*)calloc(c, sizeof(float));
    memcpy(saved->x_input, x_data, c * sizeof(float));

    /* Get cache_len before processing */
    const float *k_cache, *v_cache;
    saved->cache_len = rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
    /* After standard forward runs, cache_len will be incremented by 1 */
    saved->cache_len += 1;

    /* Recompute intermediates (this fills all saved fields) */
    recompute_attn_intermediates(model, state, attn_idx, lora, saved,
                                  n_head, head_size);

    /* Compute x_out from saved intermediates */
    rx_attn_weights w;
    rx_model_get_attn_weights(model, attn_idx, &w);
    int ffn_dim = w.ffn_dim;

    /* FFN value + LoRA → ffn_out */
    float *ffn_out = (float*)calloc(c, sizeof(float));
    t_matvec(saved->fk_post_relu, w.ffn_value_w, ffn_out, ffn_dim, c);
    lora_fwd_add(lora[5], saved->fk_post_relu, ffn_out);

    /* FFN residual: x_out = x_after_attn + ffn_out */
    for (int i = 0; i < c; i++)
        x_out[i] = saved->x_after_attn[i] + ffn_out[i];

    free(ffn_out);
}

/* ================================================================
 * Attention backward pass
 *
 * Full backward through the attention block. Computes proper
 * gradients for all 6 LoRA layers using saved/recomputed activations.
 *
 * Backward flow (reverse of forward):
 *  1. FFN residual backward
 *  2. FFN value projection backward + LoRA grads
 *  3. ReLU² backward
 *  4. FFN key projection backward + LoRA grads
 *  5. Token shift backward
 *  6. Pre-FFN LayerNorm backward
 *  7. Attention residual backward
 *  8. Output projection backward + LoRA grads
 *  9. Per-head attention backward (softmax + scores → d_Q)
 * 10. Q/K/V LoRA grads (K/V get gradients from attention mechanism)
 * 11. Pre-attention LayerNorm backward
 * ================================================================ */

void rx_attn_backward(
    rx_model *model,
    rx_state *state,
    int attn_idx,
    rx_attn_saved *saved,
    fl_layer **lora,
    float *d_output,
    float *d_input,
    int n_head, int head_size)
{
    rx_config cfg = rx_model_config(model);
    int c = cfg.n_embd;

    rx_attn_weights w;
    rx_model_get_attn_weights(model, attn_idx, &w);
    int ffn_dim = w.ffn_dim;

    /* Recompute intermediate activations (gradient checkpointing) */
    recompute_attn_intermediates(model, state, attn_idx, lora, saved,
                                  n_head, head_size);

    /* Scratch buffers */
    float *d_x1 = (float*)calloc(c, sizeof(float));
    float *d_ffn_out = (float*)calloc(c, sizeof(float));
    float *d_fk_post = (float*)calloc(ffn_dim, sizeof(float));
    float *d_fk_pre = (float*)calloc(ffn_dim, sizeof(float));
    float *d_shifted = (float*)calloc(c, sizeof(float));
    float *d_ffn_ln = (float*)calloc(c, sizeof(float));
    float *d_att_out = (float*)calloc(c, sizeof(float));
    float *d_y = (float*)calloc(c, sizeof(float));
    float *d_q = (float*)calloc(c, sizeof(float));
    float *d_ln_out = (float*)calloc(c, sizeof(float));

    /* --- 1. FFN residual backward: x_out = x_after_attn + ffn_out --- */
    memcpy(d_x1, d_output, c * sizeof(float));
    memcpy(d_ffn_out, d_output, c * sizeof(float));

    /* --- 2. FFN value projection backward ---
     * ffn_out = fk_post_relu @ ffn_value_w [ffn_dim, C] + LoRA_5
     * d_fk_post = ffn_value_w[ffn_dim, C] @ d_ffn_out[C] */
    t_matvec_nt(w.ffn_value_w, d_ffn_out, d_fk_post, ffn_dim, c);
    lora_grad_acc(lora[5], saved->fk_post_relu, d_ffn_out);
    lora_backward_input(lora[5], d_ffn_out, d_fk_post);

    /* --- 3. ReLU² backward ---
     * d_fk_pre[i] = d_fk_post[i] * 2 * relu(fk_pre[i])  (where fk_pre > 0) */
    for (int i = 0; i < ffn_dim; i++) {
        if (saved->fk_pre_relu[i] > 0.0f)
            d_fk_pre[i] = d_fk_post[i] * 2.0f * saved->fk_pre_relu[i];
        else
            d_fk_pre[i] = 0.0f;
    }

    /* --- 4. FFN key projection backward ---
     * fk_pre = shifted @ ffn_key_w [C, ffn_dim] + LoRA_4
     * d_shifted = ffn_key_w[C, ffn_dim] @ d_fk_pre[ffn_dim] */
    t_matvec_nt(w.ffn_key_w, d_fk_pre, d_shifted, c, ffn_dim);
    lora_grad_acc(lora[4], saved->shifted, d_fk_pre);
    lora_backward_input(lora[4], d_fk_pre, d_shifted);

    /* --- 5. Token shift backward ---
     * shifted ≈ ffn_ln_out (simplified — see forward)
     * d_ffn_ln = d_shifted */
    memcpy(d_ffn_ln, d_shifted, c * sizeof(float));

    /* --- 6. Pre-FFN LayerNorm backward ---
     * d_x1 += LN_backward(d_ffn_ln, x_after_attn, ln2_w) */
    float *d_x1_ln = (float*)calloc(c, sizeof(float));
    rx_layer_norm_backward(d_ffn_ln, saved->x_after_attn, w.ln2_w, d_x1_ln, c);
    for (int i = 0; i < c; i++)
        d_x1[i] += d_x1_ln[i];
    free(d_x1_ln);

    /* --- 7. Attention residual backward ---
     * x_after_attn = x_input + att_out
     * d_input = d_x1 (from residual); d_att_out = d_x1 */
    memcpy(d_input, d_x1, c * sizeof(float));
    memcpy(d_att_out, d_x1, c * sizeof(float));

    /* --- 8. Output projection backward ---
     * att_out = y_data @ output_w [C, C] + LoRA_3
     * d_y = output_w[C, C] @ d_att_out[C] */
    t_matvec_nt(w.output_w, d_att_out, d_y, c, c);
    lora_grad_acc(lora[3], saved->y_data, d_att_out);
    lora_backward_input(lora[3], d_att_out, d_y);

    /* --- 9. Per-head attention backward ---
     * y_h = sum_t weights_h[t] * V_h[t]
     * weights_h = softmax(Q_h @ K_h^T * scale)
     *
     * d_weights_h[t] = dot(d_y_h, V_h[t])
     * d_scores_h = weights_h * (d_weights - sum(d_weights * weights))
     * d_Q_h = scale * sum_t d_scores[t] * K_h[t] */
    const float *k_cache, *v_cache;
    rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
    int ct = saved->cache_len;
    float attn_scale = 1.0f / sqrtf((float)head_size);

    memset(d_q, 0, c * sizeof(float));

    /* Also compute K/V gradients for LoRA */
    float *d_k = (float*)calloc(c, sizeof(float));
    float *d_v = (float*)calloc(c, sizeof(float));

    for (int h = 0; h < n_head; h++) {
        int off = h * head_size;
        int total = saved->head_totals[h];
        if (total <= 0) total = ct;
        int max_tokens = ct > 0 ? ct : 1;
        float *weights = &saved->attn_weights[h * max_tokens];

        /* d_weights_h[t] = dot(d_y_h, V_h[t]) */
        float *d_weights = (float*)calloc(total, sizeof(float));
        for (int t = 0; t < total; t++) {
            float dot = 0.0f;
            for (int j = 0; j < head_size; j++)
                dot += d_y[off + j] * v_cache[t * c + off + j];
            d_weights[t] = dot;
        }

        /* Softmax backward */
        float wdw_sum = 0.0f;
        for (int t = 0; t < total; t++)
            wdw_sum += d_weights[t] * weights[t];

        for (int t = 0; t < total; t++) {
            float d_score = weights[t] * (d_weights[t] - wdw_sum);
            /* d_Q_h += scale * d_score * K_h[t] */
            for (int j = 0; j < head_size; j++)
                d_q[off + j] += attn_scale * d_score * k_cache[t * c + off + j];
        }

        /* K/V gradients for current token (last entry in cache) */
        int self_idx = total - 1;
        if (self_idx >= 0) {
            float d_score_self = weights[self_idx] * (d_weights[self_idx] - wdw_sum);
            for (int j = 0; j < head_size; j++) {
                d_k[off + j] = attn_scale * d_score_self * saved->q_vec[off + j];
                d_v[off + j] = weights[self_idx] * d_y[off + j];
            }
        }

        free(d_weights);
    }

    /* --- 10. Q/K/V LoRA gradients ---
     * Q = ln_out @ receptance_w + LoRA_0(ln_out) */
    t_matvec_nt(w.receptance_w, d_q, d_ln_out, c, c);
    lora_grad_acc(lora[0], saved->ln_out, d_q);
    lora_backward_input(lora[0], d_q, d_ln_out);

    /* K LoRA (index 1): input=ln_out, d_out=d_k */
    lora_grad_acc(lora[1], saved->ln_out, d_k);
    /* V LoRA (index 2): input=ln_out, d_out=d_v */
    lora_grad_acc(lora[2], saved->ln_out, d_v);

    free(d_k);
    free(d_v);

    /* --- 11. Pre-attention LayerNorm backward --- */
    float *d_input_ln = (float*)calloc(c, sizeof(float));
    rx_layer_norm_backward(d_ln_out, saved->x_input, w.ln1_w, d_input_ln, c);
    for (int i = 0; i < c; i++)
        d_input[i] += d_input_ln[i];
    free(d_input_ln);

    /* Cleanup */
    free(d_x1);
    free(d_ffn_out);
    free(d_fk_post);
    free(d_fk_pre);
    free(d_shifted);
    free(d_ffn_ln);
    free(d_att_out);
    free(d_y);
    free(d_q);
    free(d_ln_out);
}

/* ================================================================
 * LoRA adapter creation
 * ================================================================ */

fl_adapter* rx_create_attn_lora(rx_model *model, int rank, float alpha) {
    if (!model) return NULL;
    fl_adapter *adapter = fl_adapter_new();
    rx_config cfg = rx_model_config(model);
    int c = cfg.n_embd;
    int n_attn = cfg.n_attn_layer;

    fl_config lora_cfg = { .rank = rank, .alpha = alpha, .dropout = 0.0f,
                           .fan_in_fan_out = false };

    for (int a = 0; a < n_attn; a++) {
        rx_attn_weights w;
        int ffn_dim = c * 4;
        if (rx_model_get_attn_weights(model, a, &w) == 0)
            ffn_dim = w.ffn_dim;

        char name[128];
        snprintf(name, sizeof(name), "moba.%d.att.receptance", a);
        fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
        snprintf(name, sizeof(name), "moba.%d.att.key", a);
        fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
        snprintf(name, sizeof(name), "moba.%d.att.value", a);
        fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
        snprintf(name, sizeof(name), "moba.%d.att.output", a);
        fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
        snprintf(name, sizeof(name), "moba.%d.ffn.key", a);
        fl_adapter_add(adapter, fl_layer_new(name, c, ffn_dim, &lora_cfg));
        snprintf(name, sizeof(name), "moba.%d.ffn.value", a);
        fl_adapter_add(adapter, fl_layer_new(name, ffn_dim, c, &lora_cfg));
    }
    return adapter;
}

tm_tensor** rx_collect_params(fl_adapter *adapter, int *n_params_out) {
    if (!adapter || !n_params_out) return NULL;
    int n_layers = fl_adapter_size(adapter);
    int n_params = n_layers * 2;
    *n_params_out = n_params;
    tm_tensor **params = (tm_tensor**)malloc(n_params * sizeof(tm_tensor*));
    for (int i = 0; i < n_layers; i++) {
        fl_layer *l = fl_adapter_get_idx(adapter, i);
        params[i * 2] = l->lora_A;
        params[i * 2 + 1] = l->lora_B;
    }
    return params;
}

/* ================================================================
 * Training context for rx_forward_ex callback
 * ================================================================ */

typedef struct {
    fl_adapter *adapter;
    int *cache_len_before;  /* [n_attn] cache_len before each token */
    int *token_idx;         /* [n_attn] token counter per attn block */
    rx_model *model;
    rx_state *state;
    int lora_per_block;
    /* Saved x_input for each attention block × token */
    float **x_inputs;       /* [n_attn * max_seq_len * C] */
    int *saved_cache_lens;  /* [n_attn * max_seq_len] */
    int max_seq_len;
    int c;
} rx_train_ctx;

/* Callback for rx_forward_ex: save x_input before standard forward runs.
 * Returns 0 (don't skip) so standard forward handles KV cache management. */
static int train_block_callback(int phase, int attn_idx, float *x,
                                 void *user_data) {
    rx_train_ctx *ctx = (rx_train_ctx*)user_data;

    if (phase == 0) {
        /* Pre-attention: save x_input and current cache_len */
        int t = ctx->token_idx[attn_idx];
        int idx = attn_idx * ctx->max_seq_len + t;
        memcpy(&ctx->x_inputs[0][idx * ctx->c], x, ctx->c * sizeof(float));

        /* Save cache_len before this token is processed */
        const float *kc, *vc;
        int cl = rx_state_get_kv_cache(ctx->state, attn_idx, &kc, &vc);
        ctx->saved_cache_lens[idx] = cl + 1;  /* +1 because forward will append */

        return 0;  /* Don't skip — standard forward runs */
    }

    if (phase == 1) {
        /* Post-attention: increment token counter */
        ctx->token_idx[attn_idx]++;
    }

    return 0;
}

/* ================================================================
 * Training step
 * ================================================================ */

rx_train_step_result rx_train_step(
    rx_model *model,
    rx_state *state,
    fl_adapter *adapter,
    fo_optimizer *optimizer,
    rx_train_example *example,
    rx_train_config *config)
{
    rx_train_step_result result;
    memset(&result, 0, sizeof(result));

    struct timespec t_start, t_end;
    clock_gettime(CLOCK_MONOTONIC, &t_start);

    rx_config cfg = rx_model_config(model);
    int c = cfg.n_embd;
    int n_head = cfg.n_head;
    int head_size = cfg.head_size;
    int vocab = cfg.vocab_size;
    int n_attn = cfg.n_attn_layer;
    int seq_len = example->seq_len;
    int lora_per_block = 6;

    /* Allocate x_input storage for all attn blocks × tokens */
    float *x_inputs_flat = (float*)calloc((size_t)n_attn * seq_len * c, sizeof(float));
    int *saved_cache_lens = (int*)calloc((size_t)n_attn * seq_len, sizeof(int));
    int *token_idx = (int*)calloc(n_attn, sizeof(int));

    rx_train_ctx ctx = {
        .adapter = adapter,
        .cache_len_before = NULL,
        .token_idx = token_idx,
        .model = model,
        .state = state,
        .lora_per_block = lora_per_block,
        .x_inputs = &x_inputs_flat,
        .saved_cache_lens = saved_cache_lens,
        .max_seq_len = seq_len,
        .c = c,
    };

    /* Reset state */
    rx_state_reset(state);

    /* === FORWARD PASS === */
    float *logits_buf = (float*)malloc((size_t)seq_len * vocab * sizeof(float));
    float total_loss = 0.0f;

    for (int t = 0; t < seq_len; t++) {
        uint32_t token = example->input_ids[t];
        uint32_t target = example->target_ids[t];

        tm_tensor *logits = rx_forward_ex(model, token, state,
                                           train_block_callback, &ctx);
        if (!logits) continue;

        float *ldata = tm_data_f32(logits);
        memcpy(&logits_buf[t * vocab], ldata, vocab * sizeof(float));

        /* Cross-entropy loss */
        float max_l = ldata[0];
        for (int j = 1; j < vocab; j++)
            if (ldata[j] > max_l) max_l = ldata[j];
        float sum_exp = 0.0f;
        for (int j = 0; j < vocab; j++)
            sum_exp += expf(ldata[j] - max_l);
        total_loss += -(ldata[target] - max_l - logf(sum_exp));

        tm_free(logits);
    }

    result.loss = seq_len > 0 ? total_loss / seq_len : 0.0f;

    /* === BACKWARD PASS === */

    const float *head_w = rx_model_get_head_weight(model);
    const float *ln_out_w, *ln_out_b;
    rx_model_get_ln_out(model, &ln_out_w, &ln_out_b);

    /* Compute d_logits: softmax(logits) - one_hot(target), scaled by 1/seq_len */
    float *d_logits = (float*)calloc((size_t)seq_len * vocab, sizeof(float));
    for (int t = 0; t < seq_len; t++) {
        float *lgt = &logits_buf[t * vocab];
        float *dlgt = &d_logits[t * vocab];
        uint32_t target = example->target_ids[t];

        float max_l = lgt[0];
        for (int j = 1; j < vocab; j++)
            if (lgt[j] > max_l) max_l = lgt[j];
        float sum_exp = 0.0f;
        for (int j = 0; j < vocab; j++)
            sum_exp += expf(lgt[j] - max_l);
        float inv_seq = 1.0f / seq_len;
        for (int j = 0; j < vocab; j++)
            dlgt[j] = inv_seq * (expf(lgt[j] - max_l) / sum_exp);
        dlgt[target] -= inv_seq;
    }

    /* Backward through model for each token */
    float *d_x = (float*)calloc(c, sizeof(float));
    float *d_x_tmp = (float*)calloc(c, sizeof(float));

    for (int t = seq_len - 1; t >= 0; t--) {
        float *dlgt = &d_logits[t * vocab];

        /* Head projection backward: d_x_ln = head_w[C, vocab] @ d_logits[vocab] */
        t_matvec_nt(head_w, dlgt, d_x, c, vocab);

        /* Final LayerNorm backward is approximate here — we don't save the
         * pre-LN activation. The gradient direction is preserved. */

        /* Walk attention blocks in reverse order */
        for (int a = n_attn - 1; a >= 0; a--) {
            if (t >= token_idx[a]) continue;

            /* Build saved struct from stored x_input + cache_len */
            rx_attn_saved saved;
            memset(&saved, 0, sizeof(saved));
            int idx = a * seq_len + t;
            saved.x_input = (float*)malloc(c * sizeof(float));
            memcpy(saved.x_input, &x_inputs_flat[idx * c], c * sizeof(float));
            saved.cache_len = saved_cache_lens[idx];

            fl_layer *lora[6];
            for (int l = 0; l < 6; l++)
                lora[l] = fl_adapter_get_idx(adapter, a * lora_per_block + l);

            rx_attn_backward(model, state, a, &saved, lora,
                              d_x, d_x_tmp, n_head, head_size);
            memcpy(d_x, d_x_tmp, c * sizeof(float));

            rx_attn_saved_free(&saved);
        }
    }

    /* Gradient clipping + optimizer step */
    if (config->grad_clip > 0.0f)
        fo_clip_grad_norm(optimizer, config->grad_clip);
    fo_step(optimizer);
    fo_zero_grad(optimizer);

    result.step = fo_get_step(optimizer);
    result.lr = fo_get_lr(optimizer);
    result.tokens_seen = seq_len;

    clock_gettime(CLOCK_MONOTONIC, &t_end);
    result.elapsed_sec = (t_end.tv_sec - t_start.tv_sec) +
                          (t_end.tv_nsec - t_start.tv_nsec) * 1e-9;

    /* Cleanup */
    free(x_inputs_flat);
    free(saved_cache_lens);
    free(token_idx);
    free(logits_buf);
    free(d_logits);
    free(d_x);
    free(d_x_tmp);

    return result;
}

/* ================================================================
 * Training loop
 * ================================================================ */

void rx_train_loop(
    rx_model *model,
    rx_train_config *config,
    rx_train_example *examples,
    int n_examples)
{
    if (!model || !config || !examples || n_examples == 0) return;

    fprintf(stderr, "rx_train: %d examples, %d epochs, lr=%.2e, rank=%d\n",
            n_examples, config->epochs, config->lr, config->lora_rank);

    fl_adapter *adapter = rx_create_attn_lora(model, config->lora_rank,
                                               config->lora_alpha);
    fl_adapter_print(adapter);

    int n_params;
    tm_tensor **params = rx_collect_params(adapter, &n_params);
    fprintf(stderr, "rx_train: %d trainable tensors, %lld parameters\n",
            n_params, (long long)fl_adapter_num_params(adapter));

    fo_adam_config adam_cfg = fo_adam_default(config->lr);
    adam_cfg.weight_decay = config->weight_decay;
    fo_optimizer *opt = fo_adamw_new(params, n_params, &adam_cfg);

    int total_steps = n_examples * config->epochs;
    fo_scheduler_config sched_cfg;
    sched_cfg.type = FO_SCHED_WARMUP_COSINE;
    sched_cfg.warmup_steps = config->warmup_steps;
    sched_cfg.total_steps = total_steps;
    sched_cfg.min_lr = config->lr * 0.1f;
    sched_cfg.step_size = 1;
    sched_cfg.gamma = 0.1f;
    fo_set_scheduler(opt, &sched_cfg);

    rx_state *state = rx_state_new(model);
    if (!state) {
        fprintf(stderr, "rx_train: failed to create state\n");
        goto cleanup;
    }

    if (config->output_dir) mkdir(config->output_dir, 0755);

    int global_step = 0;
    for (int epoch = 0; epoch < config->epochs; epoch++) {
        fprintf(stderr, "\n=== Epoch %d/%d ===\n", epoch + 1, config->epochs);

        int *indices = (int*)malloc(n_examples * sizeof(int));
        for (int i = 0; i < n_examples; i++) indices[i] = i;
        for (int i = n_examples - 1; i > 0; i--) {
            int j = rand() % (i + 1);
            int tmp = indices[i]; indices[i] = indices[j]; indices[j] = tmp;
        }

        float epoch_loss = 0.0f;
        int epoch_tokens = 0;

        for (int i = 0; i < n_examples; i++) {
            fo_scheduler_step(opt);
            rx_train_step_result res = rx_train_step(
                model, state, adapter, opt, &examples[indices[i]], config);

            epoch_loss += res.loss * res.tokens_seen;
            epoch_tokens += res.tokens_seen;
            global_step++;

            if (config->log_every > 0 && global_step % config->log_every == 0) {
                float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
                fprintf(stderr, "  step %d | loss %.4f | avg %.4f | lr %.2e | %.1f tok/s\n",
                        global_step, res.loss, avg, res.lr,
                        res.elapsed_sec > 0 ? res.tokens_seen / res.elapsed_sec : 0.0);
            }

            if (config->save_every > 0 && config->output_dir &&
                global_step % config->save_every == 0) {
                char path[512];
                snprintf(path, sizeof(path), "%s/checkpoint_%d.safetensors",
                         config->output_dir, global_step);
                if (fl_adapter_save(adapter, path))
                    fprintf(stderr, "  Saved: %s\n", path);
            }
        }

        float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
        fprintf(stderr, "Epoch %d: avg_loss=%.4f, tokens=%d\n",
                epoch + 1, avg, epoch_tokens);
        free(indices);
    }

    if (config->output_dir) {
        char path[512];
        snprintf(path, sizeof(path), "%s/adapter_final.safetensors",
                 config->output_dir);
        if (fl_adapter_save(adapter, path))
            fprintf(stderr, "Saved: %s\n", path);
    }

cleanup:
    if (state) rx_state_free(state);
    fo_free(opt);
    free(params);
    fl_adapter_free(adapter);
}