#ifndef CACTUS_BACKPROP_H
#define CACTUS_BACKPROP_H

// ============================================================================
// cactus_backprop.{h,cpp} — pure, native-testable backpropagation for the
// cactus_fullmode MLP architecture (lib/cactus_fullmode).
//
// Added 2026-07-30 as the first real piece of "federated learning v2" (the
// on-device gradient-push side talon_federated.h documents as deferred:
// "requires backprop over int8 activations ... deferred until model_manager
// exposes a gradient callback"). This module is that missing math.
//
// Deliberately a SEPARATE module from cactus_fullmode.cpp, not an addition
// to it: cactus_fullmode.cpp unconditionally includes <Arduino.h> and
// <esp_heap_caps.h> (PSRAM), so it cannot be compiled or tested natively
// (its own test, test_fullmode_ondevice.cpp, says so explicitly and only
// runs on real hardware). Delicate gradient math needs the same rigor this
// session already applied to the Ed25519/X25519/Reticulum work -- verified
// against ground truth, not just self-consistency -- and the only
// ground truth available here (no external ML framework reference exists
// for this exact hand-rolled architecture) is numerical gradient checking:
// perturb each weight by +-epsilon, measure the resulting change in loss
// via two forward passes, and confirm it matches the analytical gradient
// this module computes. See test/test_cactus_backprop/test_cactus_backprop.cpp.
//
// CactusBpLayer mirrors cactus_fullmode.cpp's private LayerDesc structurally
// (same fields, same meaning) but is its own type here, since LayerDesc
// itself is private to that Arduino-only file. A thin bridge function in
// cactus_fullmode.cpp (ARDUINO-side only) converts a loaded slot's real
// LayerDesc array to this pure representation (the W/b pointers are
// reused as-is, not copied -- they already point into the slot's PSRAM
// pool) and calls into this module.
//
// SCOPE, STATED HONESTLY: this is the gradient-computation primitive only.
// It does NOT decide what counts as a valid on-device training example
// (talon_federated.h's talon_collect_local_anchors() needs "access to a
// labeled local detection window" -- there is no ground-truth label source
// on a T1 sensor today; inventing one here would be a product/ML-policy
// decision, not an engineering fill-in-the-blank, so that function is left
// exactly as honestly-stubbed as it already was). Nor does it decide a
// learning rate or handle int8 scale drift beyond the conservative,
// documented policy in cactus_bp_apply_gradient() below.
// ============================================================================

#include <cstdint>
#include <cstddef>

enum CactusBpLayerType : uint8_t {
    CBP_LINEAR_FP32 = 0,
    CBP_LINEAR_INT8 = 1,
    CBP_RELU        = 2,
    CBP_SOFTMAX     = 3,
};

// Mirrors cactus_fullmode.cpp's private LayerDesc (type, in_dim, out_dim,
// w_scale, W, b) field-for-field -- see the design note above for why this
// is a separate type rather than a shared one.
struct CactusBpLayer {
    CactusBpLayerType type;
    int     in_dim;
    int     out_dim;
    float   w_scale;   // int8 layers only; ignored for fp32/relu/softmax
    void*   W;         // float* (fp32) or int8_t* (int8); linear layers only, row-major (in_dim, out_dim)
    float*  b;         // linear layers only, length out_dim
};

#define CACTUS_BP_MAX_LAYERS 16
// Generous headroom over the shipped threat_classifier's largest dimension
// (128 in this session's audit) -- bump if a wider architecture ships.
#define CACTUS_BP_MAX_DIM    256

// ---------------------------------------------------------------------------
// Forward pass with full activation caching (unlike cactus_fm_infer(),
// which ping-pongs through 2 scratch buffers and destroys every
// intermediate result except the current pair -- backward() cannot recover
// what forward() doesn't keep).
//
// activations_out[0]   = network input (copied in, length input_dim)
// activations_out[i+1] = layer i's OUTPUT (post-activation for relu/
//                        softmax, post-linear for linear layers)
// A ReLU layer's PRE-activation is simply its own input, i.e.
// activations_out[i] for layer i -- no separate storage needed for that;
// backward() reads it directly from this same array.
// activations_out must have (n_layers+1) rows of CACTUS_BP_MAX_DIM floats,
// caller-allocated. Returns false on a dimension mismatch or
// n_layers > CACTUS_BP_MAX_LAYERS.
// ---------------------------------------------------------------------------
bool cactus_bp_forward(const CactusBpLayer* layers, int n_layers,
                       const float* input, int input_dim,
                       float activations_out[][CACTUS_BP_MAX_DIM]);

// ---------------------------------------------------------------------------
// Backprop from a one-hot classification target, assuming the LAST layer
// is softmax and the loss is categorical cross-entropy. Uses the standard
// fused softmax+cross-entropy gradient identity dL/dz = softmax_out -
// onehot(target) at the softmax layer's input, rather than a separate
// softmax Jacobian step -- mathematically equivalent, numerically simpler,
// standard practice. Requires the activations array from a prior
// cactus_bp_forward() call on the SAME input.
//
// grad_W_out[i] / grad_b_out[i] receive the gradient for linear layer i
// (nullptr for non-linear layers) -- caller allocates arrays sized
// in_dim*out_dim and out_dim respectively for each linear layer.
// Returns the cross-entropy loss (for monitoring/convergence checks; not
// used internally by the fused-gradient shortcut above).
// ---------------------------------------------------------------------------
float cactus_bp_backward(const CactusBpLayer* layers, int n_layers,
                         const float activations[][CACTUS_BP_MAX_DIM],
                         int target_class,
                         float* grad_W_out[], float* grad_b_out[]);

// ---------------------------------------------------------------------------
// Applies a computed gradient to one linear layer's weights/biases IN
// PLACE, via plain gradient descent: W -= learning_rate * grad_W (and same
// for b). For fp32 layers this is a direct update. For int8 layers, the
// update is computed in the dequantized (effective fp32) domain, then
// re-quantized: new_effective = (W[k,m]*w_scale) - lr*grad_W[k,m];
// new_int8 = round(new_effective / w_scale), clamped to [-127, 127].
//
// DOCUMENTED LIMITATION, not silently glossed over: w_scale itself is
// never adjusted. If updates repeatedly push a weight's effective value
// outside what the ORIGINAL w_scale can represent at int8 resolution, the
// clamp will saturate that weight and further updates in that direction
// become no-ops -- there is no rescaling/recalibration here. A real
// production federated-learning loop would need to either widen w_scale
// periodically (recorded and shipped alongside the delta) or cap the
// per-round learning rate low enough that this practically never
// triggers; deciding which is a policy call outside this function's scope.
// ---------------------------------------------------------------------------
void cactus_bp_apply_gradient(CactusBpLayer* layer,
                              const float* grad_W, const float* grad_b,
                              float learning_rate);

#endif // CACTUS_BACKPROP_H
