/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#include "ft_layer_ops.h"

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

/* ============================================================= */
/* RMSNorm                                                       */
/* ============================================================= */

void ft_rmsnorm(const float *x, const float *weight, float *out,
                int n, float eps) {
    double sum_sq = 0.0;
    for (int i = 0; i < n; i++)
        sum_sq += (double)x[i] * x[i];
    float rms = sqrtf((float)(sum_sq / n) + eps);
    float inv_rms = 1.0f / rms;
    for (int i = 0; i < n; i++)
        out[i] = weight[i] * x[i] * inv_rms;
}

void ft_rmsnorm_backward(const float *x, const float *weight,
                         const float *d_out, float *d_x,
                         int n, float eps) {
    /* Forward recomputation */
    double sum_sq = 0.0;
    for (int i = 0; i < n; i++)
        sum_sq += (double)x[i] * x[i];
    float variance = (float)(sum_sq / n);
    float inv_rms = 1.0f / sqrtf(variance + eps);

    /* d_x[i] = weight[i] * d_out[i] * inv_rms
     *         - x[i] * inv_rms^3 * (1/n) * sum_j(weight[j] * d_out[j] * x[j]) */
    double dot = 0.0;
    for (int i = 0; i < n; i++)
        dot += (double)weight[i] * d_out[i] * x[i];
    float coeff = (float)(dot / n) * inv_rms * inv_rms * inv_rms;

    for (int i = 0; i < n; i++)
        d_x[i] += weight[i] * d_out[i] * inv_rms - x[i] * coeff;
}

/* ============================================================= */
/* SiLU                                                          */
/* ============================================================= */

static float sigmoid_f(float x) {
    return 1.0f / (1.0f + expf(-x));
}

void ft_silu(const float *x, float *out, int n) {
    for (int i = 0; i < n; i++)
        out[i] = x[i] * sigmoid_f(x[i]);
}

void ft_silu_backward(const float *x, const float *d_out, float *d_x, int n) {
    for (int i = 0; i < n; i++) {
        float s = sigmoid_f(x[i]);
        d_x[i] += d_out[i] * (s + x[i] * s * (1.0f - s));
    }
}

/* ============================================================= */
/* Ternary Matvec (using bit-packed sign+mask)                   */
/* ============================================================= */

void ft_ternary_matvec(const ft_layer_state *state, int64_t state_offset,
                       int rows, int cols,
                       const float *x, float *out) {
    /*
     * out[row] = sum_k( ternary[row,k] * scale_group(row*cols+k) * x[k] )
     *
     * ternary value: mask ? (sign ? +1 : -1) : 0
     * scale from per-group array
     */
    for (int row = 0; row < rows; row++) {
        float sum = 0.0f;
        for (int k = 0; k < cols; k++) {
            int64_t gi = state_offset + (int64_t)row * cols + k;
            int64_t byte_idx = gi / 8;
            int bit = (int)(gi % 8);

            int m = (state->mask[byte_idx] >> bit) & 1;
            if (!m) continue;

            int s = (state->sign[byte_idx] >> bit) & 1;
            float w_val = s ? 1.0f : -1.0f;

            /* Per-group scale */
            int group_idx = (int)(gi / state->group_size);
            float scale = state->scales[group_idx];

            sum += w_val * scale * x[k];
        }
        out[row] = sum;
    }
}

void ft_ternary_matvec_backward_input(const ft_layer_state *state, int64_t state_offset,
                                       int rows, int cols,
                                       const float *d_out, float *d_x) {
    /* d_x[k] += sum_row( ternary[row,k] * scale * d_out[row] ) */
    for (int row = 0; row < rows; row++) {
        if (d_out[row] == 0.0f) continue;
        for (int k = 0; k < cols; k++) {
            int64_t gi = state_offset + (int64_t)row * cols + k;
            int64_t byte_idx = gi / 8;
            int bit = (int)(gi % 8);

            int m = (state->mask[byte_idx] >> bit) & 1;
            if (!m) continue;

            int s = (state->sign[byte_idx] >> bit) & 1;
            float w_val = s ? 1.0f : -1.0f;

            int group_idx = (int)(gi / state->group_size);
            float scale = state->scales[group_idx];

            d_x[k] += w_val * scale * d_out[row];
        }
    }
}

void ft_ternary_matvec_backward_weight(const ft_layer_state *state, int64_t state_offset,
                                        int rows, int cols,
                                        const float *x, const float *d_out,
                                        ft_grad_stats *stats) {
    (void)state; /* Weight values not needed — we accumulate FP32 gradients */

    /* For each weight [row, col]:
     *   grad = x[col] * d_out[row]
     *   accumulate into grad_sum, grad_sq_sum, hessian_diag */
    for (int row = 0; row < rows; row++) {
        for (int k = 0; k < cols; k++) {
            int64_t gi = state_offset + (int64_t)row * cols + k;
            float grad = x[k] * d_out[row];

            stats->grad_sum[gi] += grad;
            stats->grad_sq_sum[gi] += grad * grad;

            /* Gauss-Newton Hessian diagonal approximation:
             * h_diag ~= x[k]^2 * d_out[row]^2 (outer product of squared terms) */
            stats->hessian_diag[gi] += x[k] * x[k] * d_out[row] * d_out[row];
        }
    }
}

/* ============================================================= */
/* Layer Forward (LLaMA-family)                                  */
/* ============================================================= */

void ft_layer_forward(const ft_layer_state *state, const ft_config *cfg,
                      const float *input, float *output, int batch_size) {
    int n = cfg->n_embd;
    int q_dim = cfg->n_heads * cfg->head_dim;
    int kv_dim = cfg->n_kv_heads * cfg->head_dim;

    for (int b = 0; b < batch_size; b++) {
        const float *x = input + (int64_t)b * n;
        float *y = output + (int64_t)b * n;

        /* --- Attention sub-layer --- */

        /* RMSNorm */
        float *x_norm = calloc(n, sizeof(float));
        ft_rmsnorm(x, state->norm_attn, x_norm, n, cfg->rms_eps);

        /* Q/K/V projections */
        float *q = calloc(q_dim, sizeof(float));
        float *k = calloc(kv_dim, sizeof(float));
        float *v = calloc(kv_dim, sizeof(float));

        ft_ternary_matvec(state, state->weights[0].offset, q_dim, n, x_norm, q);
        ft_ternary_matvec(state, state->weights[1].offset, kv_dim, n, x_norm, k);
        ft_ternary_matvec(state, state->weights[2].offset, kv_dim, n, x_norm, v);

        /* Simplified attention: for ternarizer calibration, we use a proxy.
         * Full multi-head attention with RoPE is expensive and not needed
         * for gradient estimation. Use V directly as the attention output.
         * This is a simplification — the gradient stats are still meaningful
         * because they capture the sensitivity of each weight to loss. */
        float *attn_out = calloc(q_dim, sizeof(float));
        /* Expand KV if GQA: replicate kv_heads to match n_heads */
        int repeat = cfg->n_heads / cfg->n_kv_heads;
        for (int h = 0; h < cfg->n_heads; h++) {
            int kv_h = h / repeat;
            memcpy(attn_out + h * cfg->head_dim,
                   v + kv_h * cfg->head_dim,
                   cfg->head_dim * sizeof(float));
        }

        /* O projection */
        float *o_out = calloc(n, sizeof(float));
        ft_ternary_matvec(state, state->weights[3].offset, n, q_dim, attn_out, o_out);

        /* Residual */
        float *h = calloc(n, sizeof(float));
        for (int i = 0; i < n; i++)
            h[i] = x[i] + o_out[i];

        /* --- FFN sub-layer --- */

        /* RMSNorm */
        float *h_norm = calloc(n, sizeof(float));
        ft_rmsnorm(h, state->norm_ffn, h_norm, n, cfg->rms_eps);

        /* Gate + Up projections */
        float *gate = calloc(cfg->ffn_dim, sizeof(float));
        float *up   = calloc(cfg->ffn_dim, sizeof(float));

        ft_ternary_matvec(state, state->weights[4].offset, cfg->ffn_dim, n, h_norm, gate);
        ft_ternary_matvec(state, state->weights[5].offset, cfg->ffn_dim, n, h_norm, up);

        /* SiLU(gate) * up */
        float *ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
        ft_silu(gate, ffn_hidden, cfg->ffn_dim);
        for (int i = 0; i < cfg->ffn_dim; i++)
            ffn_hidden[i] *= up[i];

        /* Down projection */
        float *down_out = calloc(n, sizeof(float));
        ft_ternary_matvec(state, state->weights[6].offset, n, cfg->ffn_dim, ffn_hidden, down_out);

        /* Residual */
        for (int i = 0; i < n; i++)
            y[i] = h[i] + down_out[i];

        free(x_norm); free(q); free(k); free(v);
        free(attn_out); free(o_out); free(h);
        free(h_norm); free(gate); free(up);
        free(ffn_hidden); free(down_out);
    }
}

/* ============================================================= */
/* Layer Backward (LLaMA-family)                                 */
/* ============================================================= */

void ft_layer_backward(const ft_layer_state *state, const ft_config *cfg,
                       const float *input, const float *d_output,
                       float *d_input, ft_grad_stats *stats, int batch_size) {
    int n = cfg->n_embd;
    int q_dim = cfg->n_heads * cfg->head_dim;
    int kv_dim = cfg->n_kv_heads * cfg->head_dim;

    for (int b = 0; b < batch_size; b++) {
        const float *x = input + (int64_t)b * n;
        const float *dy = d_output + (int64_t)b * n;

        /* Recompute forward intermediates */
        float *x_norm = calloc(n, sizeof(float));
        ft_rmsnorm(x, state->norm_attn, x_norm, n, cfg->rms_eps);

        float *v = calloc(kv_dim, sizeof(float));
        ft_ternary_matvec(state, state->weights[2].offset, kv_dim, n, x_norm, v);

        float *attn_out = calloc(q_dim, sizeof(float));
        int repeat = cfg->n_heads / cfg->n_kv_heads;
        for (int h_i = 0; h_i < cfg->n_heads; h_i++) {
            int kv_h = h_i / repeat;
            memcpy(attn_out + h_i * cfg->head_dim,
                   v + kv_h * cfg->head_dim,
                   cfg->head_dim * sizeof(float));
        }

        float *o_out = calloc(n, sizeof(float));
        ft_ternary_matvec(state, state->weights[3].offset, n, q_dim, attn_out, o_out);

        float *h = calloc(n, sizeof(float));
        for (int i = 0; i < n; i++) h[i] = x[i] + o_out[i];

        float *h_norm = calloc(n, sizeof(float));
        ft_rmsnorm(h, state->norm_ffn, h_norm, n, cfg->rms_eps);

        float *gate = calloc(cfg->ffn_dim, sizeof(float));
        float *up = calloc(cfg->ffn_dim, sizeof(float));
        ft_ternary_matvec(state, state->weights[4].offset, cfg->ffn_dim, n, h_norm, gate);
        ft_ternary_matvec(state, state->weights[5].offset, cfg->ffn_dim, n, h_norm, up);

        float *gate_silu = calloc(cfg->ffn_dim, sizeof(float));
        ft_silu(gate, gate_silu, cfg->ffn_dim);

        float *ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
        for (int i = 0; i < cfg->ffn_dim; i++)
            ffn_hidden[i] = gate_silu[i] * up[i];

        /* --- Backward through FFN sub-layer --- */

        /* dy flows to both residual paths */
        float *d_down_out = calloc(n, sizeof(float));
        memcpy(d_down_out, dy, n * sizeof(float)); /* d_residual */

        /* Backward through down_proj */
        float *d_ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
        ft_ternary_matvec_backward_input(state, state->weights[6].offset,
                                          n, cfg->ffn_dim, d_down_out, d_ffn_hidden);
        ft_ternary_matvec_backward_weight(state, state->weights[6].offset,
                                           n, cfg->ffn_dim, ffn_hidden, d_down_out, stats);

        /* Backward through SiLU(gate) * up */
        float *d_gate_silu = calloc(cfg->ffn_dim, sizeof(float));
        float *d_up = calloc(cfg->ffn_dim, sizeof(float));
        for (int i = 0; i < cfg->ffn_dim; i++) {
            d_gate_silu[i] = d_ffn_hidden[i] * up[i];
            d_up[i] = d_ffn_hidden[i] * gate_silu[i];
        }

        float *d_gate = calloc(cfg->ffn_dim, sizeof(float));
        ft_silu_backward(gate, d_gate_silu, d_gate, cfg->ffn_dim);

        /* Backward through gate/up projections */
        float *d_h_norm_ffn = calloc(n, sizeof(float));
        ft_ternary_matvec_backward_input(state, state->weights[4].offset,
                                          cfg->ffn_dim, n, d_gate, d_h_norm_ffn);
        ft_ternary_matvec_backward_weight(state, state->weights[4].offset,
                                           cfg->ffn_dim, n, h_norm, d_gate, stats);

        float *d_h_norm_ffn2 = calloc(n, sizeof(float));
        ft_ternary_matvec_backward_input(state, state->weights[5].offset,
                                          cfg->ffn_dim, n, d_up, d_h_norm_ffn2);
        ft_ternary_matvec_backward_weight(state, state->weights[5].offset,
                                           cfg->ffn_dim, n, h_norm, d_up, stats);

        /* Combine gate+up gradients */
        for (int i = 0; i < n; i++) d_h_norm_ffn[i] += d_h_norm_ffn2[i];

        /* Backward through FFN RMSNorm */
        float *d_h = calloc(n, sizeof(float));
        ft_rmsnorm_backward(h, state->norm_ffn, d_h_norm_ffn, d_h, n, cfg->rms_eps);

        /* Add FFN residual gradient */
        for (int i = 0; i < n; i++) d_h[i] += dy[i];

        /* --- Backward through attention sub-layer --- */

        float *d_o_out = calloc(n, sizeof(float));
        memcpy(d_o_out, d_h, n * sizeof(float));

        /* Backward through O projection */
        float *d_attn_out = calloc(q_dim, sizeof(float));
        ft_ternary_matvec_backward_input(state, state->weights[3].offset,
                                          n, q_dim, d_o_out, d_attn_out);
        ft_ternary_matvec_backward_weight(state, state->weights[3].offset,
                                           n, q_dim, attn_out, d_o_out, stats);

        /* Backward through simplified attention (V passthrough) */
        float *d_v = calloc(kv_dim, sizeof(float));
        for (int h_i = 0; h_i < cfg->n_heads; h_i++) {
            int kv_h = h_i / repeat;
            for (int d = 0; d < cfg->head_dim; d++)
                d_v[kv_h * cfg->head_dim + d] += d_attn_out[h_i * cfg->head_dim + d];
        }

        /* Backward through Q/K/V projections.
         * Q and K don't contribute to output in our simplified attention,
         * but we still accumulate weight gradients for scoring. */
        float *d_x_norm = calloc(n, sizeof(float));

        /* V backward */
        ft_ternary_matvec_backward_input(state, state->weights[2].offset,
                                          kv_dim, n, d_v, d_x_norm);
        ft_ternary_matvec_backward_weight(state, state->weights[2].offset,
                                           kv_dim, n, x_norm, d_v, stats);

        /* Q backward: use zero d_out for input grad but compute weight grad
         * with a small proxy signal to get meaningful statistics */
        float *d_q_proxy = calloc(q_dim, sizeof(float));
        for (int i = 0; i < q_dim; i++) d_q_proxy[i] = d_attn_out[i] * 0.1f;
        ft_ternary_matvec_backward_weight(state, state->weights[0].offset,
                                           q_dim, n, x_norm, d_q_proxy, stats);

        /* K backward: similar proxy */
        float *d_k_proxy = calloc(kv_dim, sizeof(float));
        for (int i = 0; i < kv_dim; i++) d_k_proxy[i] = d_v[i] * 0.1f;
        ft_ternary_matvec_backward_weight(state, state->weights[1].offset,
                                           kv_dim, n, x_norm, d_k_proxy, stats);

        /* Backward through attention RMSNorm */
        float *d_x_local = calloc(n, sizeof(float));
        ft_rmsnorm_backward(x, state->norm_attn, d_x_norm, d_x_local, n, cfg->rms_eps);

        /* Add attention residual gradient */
        for (int i = 0; i < n; i++)
            d_x_local[i] += d_h[i];

        /* Accumulate into d_input */
        for (int i = 0; i < n; i++)
            d_input[(int64_t)b * n + i] += d_x_local[i];

        /* Cleanup */
        free(x_norm); free(v); free(attn_out); free(o_out); free(h);
        free(h_norm); free(gate); free(up); free(gate_silu); free(ffn_hidden);
        free(d_down_out); free(d_ffn_hidden); free(d_gate_silu); free(d_up);
        free(d_gate); free(d_h_norm_ffn); free(d_h_norm_ffn2); free(d_h);
        free(d_o_out); free(d_attn_out); free(d_v); free(d_x_norm);
        free(d_q_proxy); free(d_k_proxy); free(d_x_local);
    }
}