/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * rx_train.h — RWKV-X LoRA training for attention blocks
 *
 * Implements manual backward pass through sparse attention blocks,
 * training loop with AdamW, and data pipeline for chat history.
 *
 * Prefix: rx_ (for RWKV-X training)
 */

#ifndef RWKVX_TRAIN_H
#define RWKVX_TRAIN_H

#include <stdint.h>
#include <stdbool.h>
#include "rwkvx/rwkvx.h"
#include "lora/lora.h"
#include "optim/optim.h"
#include "data/data.h"

#ifdef __cplusplus
extern "C" {
#endif

/* ================================================================
 * Training example
 * ================================================================ */

typedef struct {
    uint32_t *input_ids;   /* [seq_len] */
    uint32_t *target_ids;  /* [seq_len] */
    int seq_len;
} rx_train_example;

/* Free a single training example */
void rx_train_example_free(rx_train_example *ex);

/* Free array of training examples */
void rx_train_examples_free(rx_train_example *examples, int n);

/* Tokenize a text dataset into training examples.
 * Each sample must have a "text" field.
 * Returns malloc'd array, sets *n_examples_out. */
rx_train_example* rx_tokenize_dataset(
    fd_dataset *ds, rx_tokenizer *tok,
    int max_seq_len, int *n_examples_out
);

/* ================================================================
 * Saved activations (for backward pass through attention blocks)
 * ================================================================ */

typedef struct {
    float *ln_out;         /* [C] LayerNorm output (input to Q/K/V) */
    float *q_vec;          /* [C] Q projection output */
    float *y_data;         /* [C] attention output (before output projection) */
    float *attn_weights;   /* [n_head * total_tokens] softmax weights per head */
    int *head_totals;      /* [n_head] number of total tokens per head */
    int cache_len;         /* number of cache tokens at this step */
    float *x_input;        /* [C] block input (for residual backward) */
    float *x_after_attn;   /* [C] x after attention residual (input to FFN LN) */
    float *ffn_ln_out;     /* [C] LayerNorm output before FFN */
    float *shifted;        /* [C] FFN input after token shift */
    float *fk_pre_relu;    /* [ffn_dim] FFN key before relu^2 */
    float *fk_post_relu;   /* [ffn_dim] FFN key after relu^2 */
} rx_attn_saved;

/* Free saved activations */
void rx_attn_saved_free(rx_attn_saved *saved);

/* ================================================================
 * Training config
 * ================================================================ */

typedef struct {
    /* Model */
    bool quantize;             /* Use INT8 base weights */

    /* LoRA */
    int lora_rank;             /* LoRA rank (default: 32) */
    float lora_alpha;          /* LoRA alpha (default: 64) */

    /* Training */
    float lr;                  /* Learning rate (default: 2e-5) */
    float weight_decay;        /* AdamW weight decay (default: 0.01) */
    int epochs;                /* Number of epochs (default: 3) */
    int max_seq_len;           /* Max sequence length (default: 512) */
    int warmup_steps;          /* LR warmup steps (default: 100) */
    int save_every;            /* Checkpoint every N steps (default: 1000) */
    int log_every;             /* Log every N steps (default: 10) */
    float grad_clip;           /* Max gradient norm (default: 1.0) */

    /* Paths */
    const char *model_path;
    const char *vocab_path;
    const char *data_path;     /* JSONL or sqlite path */
    const char *output_dir;    /* Checkpoint output directory */
} rx_train_config;

/* Return default training config */
rx_train_config rx_train_config_defaults(void);

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

/* LayerNorm backward: given d_output, original input x, weight
 * Computes d_x (gradient to propagate through) */
void rx_layer_norm_backward(
    const float *d_output,  /* [C] gradient from above */
    const float *x,         /* [C] original input to LayerNorm */
    const float *weight,    /* [C] LN weight (gamma) */
    float *d_x,             /* [C] gradient to propagate (output) */
    int c
);

/* ================================================================
 * Attention block training forward + backward
 * ================================================================ */

/* Training forward — same as attn_block_forward_fast but saves activations.
 * lora[6]: Q, K, V, output, ffn_key, ffn_value (NULL entries = no LoRA). */
void rx_attn_forward_train(
    rx_model *model,
    int attn_idx,               /* which attention block index */
    float *x_data,              /* [C] input */
    rx_state *state,
    fl_layer **lora,            /* [6] LoRA layers (NULL = no LoRA) */
    float *x_out,               /* [C] output */
    rx_attn_saved *saved        /* activations saved for backward */
);

/* Backward pass: given d_output, compute LoRA gradients.
 * Accumulates into lora[i]->lora_A->grad and lora[i]->lora_B->grad.
 * Needs state for KV cache access during attention backward. */
void rx_attn_backward(
    rx_model *model,
    rx_state *state,
    int attn_idx,
    rx_attn_saved *saved,
    fl_layer **lora,            /* [6] LoRA layers */
    float *d_output,            /* [C] gradient from above */
    float *d_input,             /* [C] gradient to propagate down (output) */
    int n_head, int head_size
);

/* ================================================================
 * Full model training pass
 * ================================================================ */

typedef struct {
    int step;
    float loss;
    float lr;
    int tokens_seen;
    double elapsed_sec;
} rx_train_step_result;

/* Single training step: forward + loss + backward + optimizer step.
 * Returns step result with loss and timing info. */
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
);

/* Full training loop over all examples for all epochs. */
void rx_train_loop(
    rx_model *model,
    rx_train_config *config,
    rx_train_example *examples,
    int n_examples
);

/* Create LoRA adapter for attention blocks.
 * Creates 6 LoRA layers per attention block: Q, K, V, output, ffn_key, ffn_value.
 * Returns adapter with all layers. */
fl_adapter* rx_create_attn_lora(rx_model *model, int rank, float alpha);

/* Collect all trainable parameters from adapter into array.
 * Returns malloc'd array of tm_tensor* pointers (A and B for each layer).
 * Sets *n_params_out. */
tm_tensor** rx_collect_params(fl_adapter *adapter, int *n_params_out);

#ifdef __cplusplus
}
#endif

#endif /* RWKVX_TRAIN_H */