/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#ifndef FOUNDRY_FT_LAYER_OPS_H
#define FOUNDRY_FT_LAYER_OPS_H

#include "ternarize.h"

#ifdef __cplusplus
extern "C" {
#endif

/*
 * Layer forward/backward dispatch for ternarizer.
 *
 * Each layer applies ternary weights via {add, nop, sub} integer ops
 * scaled by per-group float scales. Forward computes activations,
 * backward accumulates gradient statistics.
 *
 * LLaMA-family layer structure:
 *   RMSNorm -> Q/K/V proj -> RoPE -> GQA attention -> O proj -> residual
 *   -> RMSNorm -> gate/up proj -> SiLU -> down proj -> residual
 */

/* Forward pass through one transformer layer with ternary weights.
 * input:  [batch, n_embd]  (pre-residual activations)
 * output: [batch, n_embd]  (post-residual activations)
 * Adds residual internally. */
void ft_layer_forward(const ft_layer_state *state, const ft_config *cfg,
                      const float *input, float *output, int batch_size);

/* Backward pass through one transformer layer.
 * Accumulates gradient statistics (grad_sum, grad_sq_sum, hessian_diag)
 * into the provided ft_grad_stats structure.
 * d_output: [batch, n_embd]  gradient w.r.t. layer output
 * d_input:  [batch, n_embd]  gradient w.r.t. layer input (accumulated) */
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);

/* ---- Helper ops ---- */

/* RMSNorm: out[i] = weight[i] * x[i] / rms(x), where rms = sqrt(mean(x^2) + eps) */
void ft_rmsnorm(const float *x, const float *weight, float *out,
                int n, float eps);

/* RMSNorm backward: d_x given d_out, x, weight, eps */
void ft_rmsnorm_backward(const float *x, const float *weight,
                         const float *d_out, float *d_x,
                         int n, float eps);

/* Ternary matvec: out[row] = scale * sum_k(ternary[row,k] * x[k])
 * Uses sign+mask bit arrays + per-group scales.
 * state_offset: offset into the layer state's sign/mask arrays. */
void ft_ternary_matvec(const ft_layer_state *state, int64_t state_offset,
                       int rows, int cols,
                       const float *x, float *out);

/* Ternary matvec backward (d_input): d_x[k] += scale * sum_row(ternary[row,k] * d_out[row]) */
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);

/* Accumulate weight gradient stats for a ternary matmul.
 * For each weight [row,col]: grad += x[col] * d_out[row]
 * Accumulates into stats->grad_sum, grad_sq_sum, hessian_diag. */
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);

/* SiLU activation: out[i] = x[i] * sigmoid(x[i]) */
void ft_silu(const float *x, float *out, int n);

/* SiLU backward: d_x[i] = d_out[i] * (sigmoid(x[i]) + x[i] * sigmoid(x[i]) * (1 - sigmoid(x[i]))) */
void ft_silu_backward(const float *x, const float *d_out, float *d_x, int n);

#ifdef __cplusplus
}
#endif

#endif /* FOUNDRY_FT_LAYER_OPS_H */