#include "trainer.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>

/* ─────────────────────────────────────────
   INIT / FREE
   ───────────────────────────────────────── */

#define ZERO_PAIR(m, v, r, c) \
    (m) = tensor_zeros((r),(c)); \
    (v) = tensor_zeros((r),(c));

void trainer_init(Trainer* tr, float lr) {
    tr->lr           = lr;
    tr->beta1        = 0.9f;
    tr->beta2        = 0.999f;
    tr->eps          = 1e-8f;
    tr->weight_decay = 0.01f;
    tr->step         = 0;

    ZERO_PAIR(tr->m_token_embed, tr->v_token_embed, JOY_VOCAB,  JOY_DMODEL);
    ZERO_PAIR(tr->m_pos_embed,   tr->v_pos_embed,   JOY_SEQLEN, JOY_DMODEL);

    for (int i = 0; i < JOY_NLAYERS; i++) {
        ZERO_PAIR(tr->layers[i].m_Wq,   tr->layers[i].v_Wq,   JOY_DMODEL, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_Wk,   tr->layers[i].v_Wk,   JOY_DMODEL, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_Wv,   tr->layers[i].v_Wv,   JOY_DMODEL, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_Wo,   tr->layers[i].v_Wo,   JOY_DMODEL, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_W1,   tr->layers[i].v_W1,   JOY_DMODEL, JOY_DFF);
        ZERO_PAIR(tr->layers[i].m_W2,   tr->layers[i].v_W2,   JOY_DFF,    JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_ln1_w, tr->layers[i].v_ln1_w, 1, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_ln1_b, tr->layers[i].v_ln1_b, 1, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_ln2_w, tr->layers[i].v_ln2_w, 1, JOY_DMODEL);
        ZERO_PAIR(tr->layers[i].m_ln2_b, tr->layers[i].v_ln2_b, 1, JOY_DMODEL);
    }

    ZERO_PAIR(tr->m_ln_f_w, tr->v_ln_f_w, 1, JOY_DMODEL);
    ZERO_PAIR(tr->m_ln_f_b, tr->v_ln_f_b, 1, JOY_DMODEL);
    ZERO_PAIR(tr->m_lm_head, tr->v_lm_head, JOY_DMODEL, JOY_VOCAB);

    /* gradient buffers */
    tr->g_token_embed = tensor_zeros(JOY_VOCAB,  JOY_DMODEL);
    tr->g_pos_embed   = tensor_zeros(JOY_SEQLEN, JOY_DMODEL);

    for (int i = 0; i < JOY_NLAYERS; i++) {
        tr->g_layers[i].g_Wq   = tensor_zeros(JOY_DMODEL, JOY_DMODEL);
        tr->g_layers[i].g_Wk   = tensor_zeros(JOY_DMODEL, JOY_DMODEL);
        tr->g_layers[i].g_Wv   = tensor_zeros(JOY_DMODEL, JOY_DMODEL);
        tr->g_layers[i].g_Wo   = tensor_zeros(JOY_DMODEL, JOY_DMODEL);
        tr->g_layers[i].g_W1   = tensor_zeros(JOY_DMODEL, JOY_DFF);
        tr->g_layers[i].g_W2   = tensor_zeros(JOY_DFF,    JOY_DMODEL);
        tr->g_layers[i].g_ln1_w = tensor_zeros(1, JOY_DMODEL);
        tr->g_layers[i].g_ln1_b = tensor_zeros(1, JOY_DMODEL);
        tr->g_layers[i].g_ln2_w = tensor_zeros(1, JOY_DMODEL);
        tr->g_layers[i].g_ln2_b = tensor_zeros(1, JOY_DMODEL);
    }
    tr->g_ln_f_w = tensor_zeros(1, JOY_DMODEL);
    tr->g_ln_f_b = tensor_zeros(1, JOY_DMODEL);
    tr->g_lm_head = tensor_zeros(JOY_DMODEL, JOY_VOCAB);

    printf("[TRAINER] AdamW initialized, lr=%.5f\n", lr);
}

void trainer_free(Trainer* tr) {
    tensor_free(&tr->m_token_embed); tensor_free(&tr->v_token_embed);
    tensor_free(&tr->m_pos_embed);   tensor_free(&tr->v_pos_embed);
    for (int i = 0; i < JOY_NLAYERS; i++) {
        tensor_free(&tr->layers[i].m_Wq); tensor_free(&tr->layers[i].v_Wq);
        tensor_free(&tr->layers[i].m_Wk); tensor_free(&tr->layers[i].v_Wk);
        tensor_free(&tr->layers[i].m_Wv); tensor_free(&tr->layers[i].v_Wv);
        tensor_free(&tr->layers[i].m_Wo); tensor_free(&tr->layers[i].v_Wo);
        tensor_free(&tr->layers[i].m_W1); tensor_free(&tr->layers[i].v_W1);
        tensor_free(&tr->layers[i].m_W2); tensor_free(&tr->layers[i].v_W2);
        tensor_free(&tr->layers[i].m_ln1_w); tensor_free(&tr->layers[i].v_ln1_w);
        tensor_free(&tr->layers[i].m_ln1_b); tensor_free(&tr->layers[i].v_ln1_b);
        tensor_free(&tr->layers[i].m_ln2_w); tensor_free(&tr->layers[i].v_ln2_w);
        tensor_free(&tr->layers[i].m_ln2_b); tensor_free(&tr->layers[i].v_ln2_b);
    }
    tensor_free(&tr->m_ln_f_w); tensor_free(&tr->v_ln_f_w);
    tensor_free(&tr->m_ln_f_b); tensor_free(&tr->v_ln_f_b);
    tensor_free(&tr->m_lm_head); tensor_free(&tr->v_lm_head);

    tensor_free(&tr->g_token_embed); tensor_free(&tr->g_pos_embed);
    for (int i = 0; i < JOY_NLAYERS; i++) {
        tensor_free(&tr->g_layers[i].g_Wq); tensor_free(&tr->g_layers[i].g_Wk);
        tensor_free(&tr->g_layers[i].g_Wv); tensor_free(&tr->g_layers[i].g_Wo);
        tensor_free(&tr->g_layers[i].g_W1); tensor_free(&tr->g_layers[i].g_W2);
        tensor_free(&tr->g_layers[i].g_ln1_w); tensor_free(&tr->g_layers[i].g_ln1_b);
        tensor_free(&tr->g_layers[i].g_ln2_w); tensor_free(&tr->g_layers[i].g_ln2_b);
    }
    tensor_free(&tr->g_ln_f_w); tensor_free(&tr->g_ln_f_b);
    tensor_free(&tr->g_lm_head);
}

void trainer_zero_grad(Trainer* tr) {
    int n;

    n = tr->g_token_embed.rows * tr->g_token_embed.cols;
    memset(tr->g_token_embed.data, 0, n * sizeof(float));
    n = tr->g_pos_embed.rows * tr->g_pos_embed.cols;
    memset(tr->g_pos_embed.data, 0, n * sizeof(float));

    for (int i = 0; i < JOY_NLAYERS; i++) {
        #define ZG(f) n = tr->g_layers[i].f.rows * tr->g_layers[i].f.cols; \
                      memset(tr->g_layers[i].f.data, 0, n * sizeof(float));
        ZG(g_Wq) ZG(g_Wk) ZG(g_Wv) ZG(g_Wo)
        ZG(g_W1) ZG(g_W2)
        ZG(g_ln1_w) ZG(g_ln1_b) ZG(g_ln2_w) ZG(g_ln2_b)
        #undef ZG
    }

    n = tr->g_ln_f_w.rows * tr->g_ln_f_w.cols;
    memset(tr->g_ln_f_w.data, 0, n * sizeof(float));
    memset(tr->g_ln_f_b.data, 0, n * sizeof(float));
    n = tr->g_lm_head.rows * tr->g_lm_head.cols;
    memset(tr->g_lm_head.data, 0, n * sizeof(float));
}

/* ─────────────────────────────────────────
   FORWARD + LOSS + BACKWARD
   We use a simplified gradient computation:
   - Forward pass stores activations
   - Compute cross-entropy loss on output tokens
   - Backprop gradients through lm_head
   - Update embedding gradients
   This is a teaching-signal approach: only
   compute loss on output portion (after SEP).
   ───────────────────────────────────────── */

float trainer_step(Trainer* tr,
                   Transformer* model,
                   const int* tokens,
                   int n_tokens) {
    if (n_tokens < 2) return 0.0f;
    int T = n_tokens < JOY_SEQLEN ? n_tokens : JOY_SEQLEN;

    /* forward pass */
    Tensor logits = transformer_forward(model, tokens, T);

    /* find SEP position to compute loss only on output tokens */
    int sep_pos = -1;
    for (int i = 0; i < T; i++) {
        if (tokens[i] == TOK_SEP) { sep_pos = i; break; }
    }
    if (sep_pos < 0) sep_pos = 0;

    /* cross entropy loss on output tokens (after SEP) */
    float loss = 0.0f;
    int count = 0;

    /* gradient of logits: dL/d(logits) = softmax(logits) - one_hot(target) */
    Tensor d_logits = tensor_zeros(T, JOY_VOCAB);

    for (int t = sep_pos; t < T - 1; t++) {
        int target = tokens[t + 1];
        if (target == TOK_PAD) continue;

        float* row = logits.data + t * JOY_VOCAB;
        float* drow = d_logits.data + t * JOY_VOCAB;

        /* softmax */
        float mx = row[0];
        for (int v = 1; v < JOY_VOCAB; v++) if (row[v] > mx) mx = row[v];
        float sum = 0.0f;
        float probs[JOY_VOCAB];
        for (int v = 0; v < JOY_VOCAB; v++) {
            probs[v] = expf(row[v] - mx);
            sum += probs[v];
        }
        for (int v = 0; v < JOY_VOCAB; v++) probs[v] /= sum;

        loss -= logf(probs[target] + 1e-9f);
        count++;

        /* gradient: softmax - one_hot */
        for (int v = 0; v < JOY_VOCAB; v++)
            drow[v] = probs[v];
        drow[target] -= 1.0f;
    }

    if (count > 0) {
        loss /= count;
        tensor_mul_scalar(&d_logits, 1.0f / count);
    }

    /* ── backprop through lm_head ──
       logits = x_final @ lm_head   (T×D @ D×V = T×V)
       d_lm_head += x_final.T @ d_logits
       d_x_final  = d_logits @ lm_head.T
       We approximate x_final by running a second forward pass
       (simple approach for training on phone) */

    /* gradient of lm_head via outer products at output positions */
    /* We do a simplified update: accumulate g_lm_head per token */
    /* For a real backprop we need stored activations — here we use
       a finite-difference-free approximation using the logit gradient
       directly projected back through the embedding matrix */

    /* Update lm_head gradient: g_lm_head += x.T @ d_logits
       x is the final hidden state — we approximate it as embed(token) */
    for (int t = sep_pos; t < T - 1; t++) {
        int tok_id = tokens[t];
        float* drow = d_logits.data + t * JOY_VOCAB;
        for (int d = 0; d < JOY_DMODEL; d++) {
            float x_d = model->token_embed.data[tok_id * JOY_DMODEL + d];
            for (int v = 0; v < JOY_VOCAB; v++)
                tr->g_lm_head.data[d * JOY_VOCAB + v] += x_d * drow[v];
        }

        /* gradient flows back to token embeddings */
        for (int d = 0; d < JOY_DMODEL; d++) {
            float grad = 0.0f;
            for (int v = 0; v < JOY_VOCAB; v++)
                grad += drow[v] * model->lm_head.data[d * JOY_VOCAB + v];
            tr->g_token_embed.data[tok_id * JOY_DMODEL + d] += grad;
        }

        /* positional embedding gradient */
        for (int d = 0; d < JOY_DMODEL; d++) {
            float grad = 0.0f;
            for (int v = 0; v < JOY_VOCAB; v++)
                grad += drow[v] * model->lm_head.data[d * JOY_VOCAB + v];
            if (t < JOY_SEQLEN)
                tr->g_pos_embed.data[t * JOY_DMODEL + d] += grad;
        }
    }

    tensor_free(&logits);
    tensor_free(&d_logits);
    return loss;
}

/* ─────────────────────────────────────────
   ADAMW UPDATE
   ───────────────────────────────────────── */

static void adamw_update_tensor(Tensor* w, Tensor* g,
                                  Tensor* m, Tensor* v,
                                  float lr, float b1, float b2,
                                  float eps, float wd, int step) {
    float bc1 = 1.0f - powf(b1, (float)step);
    float bc2 = 1.0f - powf(b2, (float)step);
    int n = w->rows * w->cols;

    for (int i = 0; i < n; i++) {
        float gi = g->data[i];
        m->data[i] = b1 * m->data[i] + (1.0f - b1) * gi;
        v->data[i] = b2 * v->data[i] + (1.0f - b2) * gi * gi;

        float m_hat = m->data[i] / bc1;
        float v_hat = v->data[i] / bc2;

        /* AdamW: weight decay on weights, not gradients */
        w->data[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * w->data[i]);
    }
}

void trainer_update(Trainer* tr, Transformer* model) {
    tr->step++;
    float lr = tr->lr;
    float b1 = tr->beta1, b2 = tr->beta2;
    float eps = tr->eps, wd = tr->weight_decay;
    int s = tr->step;

    adamw_update_tensor(&model->token_embed, &tr->g_token_embed,
                        &tr->m_token_embed, &tr->v_token_embed,
                        lr, b1, b2, eps, wd, s);

    adamw_update_tensor(&model->pos_embed, &tr->g_pos_embed,
                        &tr->m_pos_embed, &tr->v_pos_embed,
                        lr, b1, b2, eps, 0.0f, s); /* no wd on pos embed */

    for (int i = 0; i < JOY_NLAYERS; i++) {
        TransformerLayer* l = &model->layers[i];
        #define UPD(W, G, M, V) adamw_update_tensor(&l->W, &tr->g_layers[i].G, \
            &tr->layers[i].M, &tr->layers[i].V, lr, b1, b2, eps, wd, s);

        UPD(Wq, g_Wq, m_Wq, v_Wq)
        UPD(Wk, g_Wk, m_Wk, v_Wk)
        UPD(Wv, g_Wv, m_Wv, v_Wv)
        UPD(Wo, g_Wo, m_Wo, v_Wo)
        UPD(W1, g_W1, m_W1, v_W1)
        UPD(W2, g_W2, m_W2, v_W2)
        UPD(ln1_w, g_ln1_w, m_ln1_w, v_ln1_w)
        UPD(ln1_b, g_ln1_b, m_ln1_b, v_ln1_b)
        UPD(ln2_w, g_ln2_w, m_ln2_w, v_ln2_w)
        UPD(ln2_b, g_ln2_b, m_ln2_b, v_ln2_b)
        #undef UPD
    }

    adamw_update_tensor(&model->ln_f_w, &tr->g_ln_f_w,
                        &tr->m_ln_f_w, &tr->v_ln_f_w,
                        lr, b1, b2, eps, 0.0f, s);
    adamw_update_tensor(&model->ln_f_b, &tr->g_ln_f_b,
                        &tr->m_ln_f_b, &tr->v_ln_f_b,
                        lr, b1, b2, eps, 0.0f, s);
    adamw_update_tensor(&model->lm_head, &tr->g_lm_head,
                        &tr->m_lm_head, &tr->v_lm_head,
                        lr, b1, b2, eps, wd, s);
}

void trainer_clip_grad_norm(Trainer* tr, float max_norm) {
    float norm_sq = 0.0f;
    int n;

    #define SQ(g) n = tr->g.rows * tr->g.cols; \
        for (int i = 0; i < n; i++) norm_sq += tr->g.data[i] * tr->g.data[i];

    SQ(g_token_embed) SQ(g_pos_embed) SQ(g_lm_head)
    #undef SQ

    float norm = sqrtf(norm_sq);
    if (norm > max_norm) {
        float scale = max_norm / (norm + 1e-6f);
        tensor_mul_scalar(&tr->g_token_embed, scale);
        tensor_mul_scalar(&tr->g_pos_embed, scale);
        tensor_mul_scalar(&tr->g_lm_head, scale);
        for (int i = 0; i < JOY_NLAYERS; i++) {
            tensor_mul_scalar(&tr->g_layers[i].g_Wq, scale);
            tensor_mul_scalar(&tr->g_layers[i].g_Wk, scale);
            tensor_mul_scalar(&tr->g_layers[i].g_Wv, scale);
            tensor_mul_scalar(&tr->g_layers[i].g_Wo, scale);
            tensor_mul_scalar(&tr->g_layers[i].g_W1, scale);
            tensor_mul_scalar(&tr->g_layers[i].g_W2, scale);
        }
    }
}