// ============================================================================
// test/test_cactus_backprop/test_cactus_backprop.cpp
//
// Unity tests for src/cactus_backprop.{h,cpp} -- see that header for the
// full design rationale.
//
// THE CORE VERIFICATION STRATEGY: numerical gradient checking. There is no
// external reference implementation to check this exact hand-rolled
// architecture against (unlike the Ed25519/X25519/Reticulum work earlier
// this session, which was verified byte-for-byte against the real
// `cryptography`/`RNS` Python libraries) -- so instead, every analytical
// gradient this module computes is cross-checked against a numerical
// (finite-difference) estimate of the SAME gradient: perturb one weight by
// +-epsilon, run two real forward passes, measure how much the loss
// actually changed, and confirm that matches what cactus_bp_backward()
// claims. This is mathematically self-verifying and independent of any
// bug that might exist identically in both a hand-derivation and a naive
// implementation -- if forward() and backward() disagree about the
// gradient of the very function forward() computes, this WILL catch it.
//
// Run with:
//   pio test -e native -f test_cactus_backprop
// ============================================================================

#include <unity.h>
#include <cstring>
#include <cmath>
#include <cstdio>
#include "../../src/cactus_backprop.cpp"

void setUp(void) {}
void tearDown(void) {}

// ---------------------------------------------------------------------------
// A small, fixed, non-degenerate synthetic network: 4 -> linear(4,3,fp32)
// -> relu -> linear(3,3,fp32) -> softmax. Weights are hand-picked non-zero,
// non-symmetric values so no gradient direction is accidentally zero.
// ---------------------------------------------------------------------------
static float g_W1[4 * 3] = {
     0.20f, -0.10f,  0.05f,
    -0.30f,  0.15f,  0.25f,
     0.10f,  0.35f, -0.20f,
    -0.05f, -0.15f,  0.30f,
};
static float g_b1[3] = { 0.01f, -0.02f, 0.03f };

static float g_W2[3 * 3] = {
     0.40f, -0.25f,  0.10f,
    -0.10f,  0.30f, -0.15f,
     0.20f, -0.05f,  0.25f,
};
static float g_b2[3] = { -0.01f, 0.02f, 0.00f };

static void make_fp32_net(CactusBpLayer layers[3]) {
    layers[0] = { CBP_LINEAR_FP32, 4, 3, 0.0f, g_W1, g_b1 };
    layers[1] = { CBP_RELU,        3, 3, 0.0f, nullptr, nullptr };
    layers[2] = { CBP_LINEAR_FP32, 3, 3, 0.0f, g_W2, g_b2 };
    // NOTE: this 3-layer array does NOT end in softmax -- tests that need
    // a full network add a 4th CBP_SOFTMAX layer themselves, since a bare
    // linear/relu prefix is also useful on its own (e.g. for probing
    // intermediate activations in a future test).
}

static const float TEST_INPUT[4] = { 0.5f, -0.3f, 0.8f, 0.1f };

// ---------------------------------------------------------------------------
// Forward pass sanity
// ---------------------------------------------------------------------------
void test_forward_softmax_output_sums_to_one(void) {
    CactusBpLayer layers[4];
    make_fp32_net(layers);
    layers[3] = { CBP_SOFTMAX, 3, 3, 0.0f, nullptr, nullptr };

    float act[5][CACTUS_BP_MAX_DIM];
    TEST_ASSERT_TRUE(cactus_bp_forward(layers, 4, TEST_INPUT, 4, act));

    float sum = 0.0f;
    for (int i = 0; i < 3; i++) {
        TEST_ASSERT_TRUE(act[4][i] >= 0.0f);
        sum += act[4][i];
    }
    TEST_ASSERT_FLOAT_WITHIN(1e-5f, 1.0f, sum);
}

void test_forward_relu_zeros_negative_inputs(void) {
    CactusBpLayer layers[2];
    layers[0] = { CBP_LINEAR_FP32, 4, 3, 0.0f, g_W1, g_b1 };
    layers[1] = { CBP_RELU, 3, 3, 0.0f, nullptr, nullptr };

    float act[3][CACTUS_BP_MAX_DIM];
    TEST_ASSERT_TRUE(cactus_bp_forward(layers, 2, TEST_INPUT, 4, act));
    for (int i = 0; i < 3; i++) {
        TEST_ASSERT_TRUE(act[2][i] >= 0.0f);
        if (act[1][i] <= 0.0f) TEST_ASSERT_EQUAL_FLOAT(0.0f, act[2][i]);
        else TEST_ASSERT_EQUAL_FLOAT(act[1][i], act[2][i]);
    }
}

void test_forward_rejects_dimension_mismatch(void) {
    CactusBpLayer layers[1];
    layers[0] = { CBP_LINEAR_FP32, 4, 3, 0.0f, g_W1, g_b1 };
    float act[2][CACTUS_BP_MAX_DIM];
    const float wrong_input[3] = {1,2,3};
    TEST_ASSERT_FALSE(cactus_bp_forward(layers, 1, wrong_input, 3, act));
}

// ---------------------------------------------------------------------------
// Numerical gradient check -- fp32 layers, both W and b, for every element
// ---------------------------------------------------------------------------
static float loss_for_weights(CactusBpLayer layers[4], int target_class) {
    float act[5][CACTUS_BP_MAX_DIM];
    cactus_bp_forward(layers, 4, TEST_INPUT, 4, act);
    const float* probs = act[4];
    float p = probs[target_class] > 1e-12f ? probs[target_class] : 1e-12f;
    return -logf(p);
}

void test_numerical_gradient_check_fp32_layer1_weights(void) {
    CactusBpLayer layers[4];
    make_fp32_net(layers);
    layers[3] = { CBP_SOFTMAX, 3, 3, 0.0f, nullptr, nullptr };
    const int target = 1;

    float act[5][CACTUS_BP_MAX_DIM];
    TEST_ASSERT_TRUE(cactus_bp_forward(layers, 4, TEST_INPUT, 4, act));

    float grad_W1[4 * 3], grad_b1[3], grad_W2[3 * 3], grad_b2[3];
    float* gW[4] = { grad_W1, nullptr, grad_W2, nullptr };
    float* gb[4] = { grad_b1, nullptr, grad_b2, nullptr };
    float loss = cactus_bp_backward(layers, 4, act, target, gW, gb);
    TEST_ASSERT_TRUE(loss > 0.0f);

    const float eps = 1e-3f;
    for (int i = 0; i < 4 * 3; i++) {
        float orig = g_W1[i];
        g_W1[i] = orig + eps;
        float loss_plus = loss_for_weights(layers, target);
        g_W1[i] = orig - eps;
        float loss_minus = loss_for_weights(layers, target);
        g_W1[i] = orig;

        float numerical = (loss_plus - loss_minus) / (2.0f * eps);
        char msg[64];
        snprintf(msg, sizeof(msg), "W1[%d]: analytical=%f numerical=%f", i, grad_W1[i], numerical);
        TEST_ASSERT_FLOAT_WITHIN_MESSAGE(2e-2f, numerical, grad_W1[i], msg);
    }
}

void test_numerical_gradient_check_fp32_layer1_bias(void) {
    CactusBpLayer layers[4];
    make_fp32_net(layers);
    layers[3] = { CBP_SOFTMAX, 3, 3, 0.0f, nullptr, nullptr };
    const int target = 2;

    float act[5][CACTUS_BP_MAX_DIM];
    cactus_bp_forward(layers, 4, TEST_INPUT, 4, act);
    float grad_W1[4 * 3], grad_b1[3], grad_W2[3 * 3], grad_b2[3];
    float* gW[4] = { grad_W1, nullptr, grad_W2, nullptr };
    float* gb[4] = { grad_b1, nullptr, grad_b2, nullptr };
    cactus_bp_backward(layers, 4, act, target, gW, gb);

    const float eps = 1e-3f;
    for (int i = 0; i < 3; i++) {
        float orig = g_b1[i];
        g_b1[i] = orig + eps;
        float loss_plus = loss_for_weights(layers, target);
        g_b1[i] = orig - eps;
        float loss_minus = loss_for_weights(layers, target);
        g_b1[i] = orig;

        float numerical = (loss_plus - loss_minus) / (2.0f * eps);
        char msg[64];
        snprintf(msg, sizeof(msg), "b1[%d]: analytical=%f numerical=%f", i, grad_b1[i], numerical);
        TEST_ASSERT_FLOAT_WITHIN_MESSAGE(2e-2f, numerical, grad_b1[i], msg);
    }
}

void test_numerical_gradient_check_fp32_layer2_weights_and_bias(void) {
    CactusBpLayer layers[4];
    make_fp32_net(layers);
    layers[3] = { CBP_SOFTMAX, 3, 3, 0.0f, nullptr, nullptr };
    const int target = 0;

    float act[5][CACTUS_BP_MAX_DIM];
    cactus_bp_forward(layers, 4, TEST_INPUT, 4, act);
    float grad_W1[4 * 3], grad_b1[3], grad_W2[3 * 3], grad_b2[3];
    float* gW[4] = { grad_W1, nullptr, grad_W2, nullptr };
    float* gb[4] = { grad_b1, nullptr, grad_b2, nullptr };
    cactus_bp_backward(layers, 4, act, target, gW, gb);

    const float eps = 1e-3f;
    for (int i = 0; i < 3 * 3; i++) {
        float orig = g_W2[i];
        g_W2[i] = orig + eps;
        float loss_plus = loss_for_weights(layers, target);
        g_W2[i] = orig - eps;
        float loss_minus = loss_for_weights(layers, target);
        g_W2[i] = orig;
        float numerical = (loss_plus - loss_minus) / (2.0f * eps);
        TEST_ASSERT_FLOAT_WITHIN(2e-2f, numerical, grad_W2[i]);
    }
    for (int i = 0; i < 3; i++) {
        float orig = g_b2[i];
        g_b2[i] = orig + eps;
        float loss_plus = loss_for_weights(layers, target);
        g_b2[i] = orig - eps;
        float loss_minus = loss_for_weights(layers, target);
        g_b2[i] = orig;
        float numerical = (loss_plus - loss_minus) / (2.0f * eps);
        TEST_ASSERT_FLOAT_WITHIN(2e-2f, numerical, grad_b2[i]);
    }
}

// ---------------------------------------------------------------------------
// Numerical gradient check -- int8 layer. Perturbation must be a multiple
// of w_scale, since that's the only granularity a real int8 weight update
// can actually take -- this also exercises the dequantized forward-pass
// math (w[k,m]*w_scale) end-to-end, not just the fp32 path.
// ---------------------------------------------------------------------------
static int8_t g_W1_int8[4 * 3];
static const float INT8_W_SCALE = 0.05f;

static void make_int8_net(CactusBpLayer layers[4]) {
    for (int i = 0; i < 4 * 3; i++) {
        int q = (int)(g_W1[i] / INT8_W_SCALE);
        if (q > 127) q = 127; if (q < -127) q = -127;
        g_W1_int8[i] = (int8_t)q;
    }
    layers[0] = { CBP_LINEAR_INT8, 4, 3, INT8_W_SCALE, g_W1_int8, g_b1 };
    layers[1] = { CBP_RELU,        3, 3, 0.0f, nullptr, nullptr };
    layers[2] = { CBP_LINEAR_FP32, 3, 3, 0.0f, g_W2, g_b2 };
    layers[3] = { CBP_SOFTMAX,     3, 3, 0.0f, nullptr, nullptr };
}

void test_numerical_gradient_check_int8_layer(void) {
    CactusBpLayer layers[4];
    make_int8_net(layers);
    const int target = 1;

    float act[5][CACTUS_BP_MAX_DIM];
    TEST_ASSERT_TRUE(cactus_bp_forward(layers, 4, TEST_INPUT, 4, act));

    float grad_W1[4 * 3], grad_b1[3], grad_W2[3 * 3], grad_b2[3];
    float* gW[4] = { grad_W1, nullptr, grad_W2, nullptr };
    float* gb[4] = { grad_b1, nullptr, grad_b2, nullptr };
    cactus_bp_backward(layers, 4, act, target, gW, gb);

    for (int i = 0; i < 4 * 3; i++) {
        int8_t orig = g_W1_int8[i];
        if (orig >= 126 || orig <= -126) continue;   // skip near-saturation elements

        g_W1_int8[i] = (int8_t)(orig + 1);
        float loss_plus = loss_for_weights(layers, target);
        g_W1_int8[i] = (int8_t)(orig - 1);
        float loss_minus = loss_for_weights(layers, target);
        g_W1_int8[i] = orig;

        // d(loss)/d(int8 unit) = d(loss)/d(effective_weight) * w_scale
        float numerical_per_int8_step = (loss_plus - loss_minus) / 2.0f;
        float analytical_per_int8_step = grad_W1[i] * INT8_W_SCALE;
        char msg[80];
        snprintf(msg, sizeof(msg), "int8 W1[%d]: analytical=%f numerical=%f", i,
                analytical_per_int8_step, numerical_per_int8_step);
        TEST_ASSERT_FLOAT_WITHIN_MESSAGE(5e-3f, numerical_per_int8_step, analytical_per_int8_step, msg);
    }
}

// ---------------------------------------------------------------------------
// cactus_bp_apply_gradient -- direct arithmetic checks (fp32 and int8)
// ---------------------------------------------------------------------------
void test_apply_gradient_fp32_moves_weight_downhill(void) {
    float W[2] = { 1.0f, -1.0f };
    float b[1] = { 0.5f };
    CactusBpLayer layer = { CBP_LINEAR_FP32, 2, 1, 0.0f, W, b };
    float grad_W[2] = { 0.2f, -0.3f };
    float grad_b[1] = { 0.1f };

    cactus_bp_apply_gradient(&layer, grad_W, grad_b, 0.1f);

    TEST_ASSERT_FLOAT_WITHIN(1e-6f, 1.0f - 0.1f * 0.2f, W[0]);
    TEST_ASSERT_FLOAT_WITHIN(1e-6f, -1.0f - 0.1f * -0.3f, W[1]);
    TEST_ASSERT_FLOAT_WITHIN(1e-6f, 0.5f - 0.1f * 0.1f, b[0]);
}

void test_apply_gradient_int8_requantizes_correctly(void) {
    int8_t W[1] = { 40 };   // effective value = 40 * 0.05 = 2.0
    float b[1] = { 0.0f };
    CactusBpLayer layer = { CBP_LINEAR_INT8, 1, 1, 0.05f, W, b };
    // A gradient + learning rate chosen to move the effective value by
    // exactly one int8 step (0.05): new_effective = 2.0 - lr*grad = 1.95
    // -> 1.95/0.05 = 39.0 -> rounds to 39.
    float grad_W[1] = { 1.0f };
    float grad_b[1] = { 0.0f };
    cactus_bp_apply_gradient(&layer, grad_W, grad_b, 0.05f);
    TEST_ASSERT_EQUAL_INT8(39, W[0]);
}

void test_apply_gradient_int8_clamps_at_saturation(void) {
    int8_t W[1] = { 127 };
    float b[1] = { 0.0f };
    CactusBpLayer layer = { CBP_LINEAR_INT8, 1, 1, 0.05f, W, b };
    float grad_W[1] = { -100.0f };   // pushes strongly positive, past +127
    float grad_b[1] = { 0.0f };
    cactus_bp_apply_gradient(&layer, grad_W, grad_b, 1.0f);
    TEST_ASSERT_EQUAL_INT8(127, W[0]);
}

int main(int, char**) {
    UNITY_BEGIN();

    RUN_TEST(test_forward_softmax_output_sums_to_one);
    RUN_TEST(test_forward_relu_zeros_negative_inputs);
    RUN_TEST(test_forward_rejects_dimension_mismatch);

    RUN_TEST(test_numerical_gradient_check_fp32_layer1_weights);
    RUN_TEST(test_numerical_gradient_check_fp32_layer1_bias);
    RUN_TEST(test_numerical_gradient_check_fp32_layer2_weights_and_bias);
    RUN_TEST(test_numerical_gradient_check_int8_layer);

    RUN_TEST(test_apply_gradient_fp32_moves_weight_downhill);
    RUN_TEST(test_apply_gradient_int8_requantizes_correctly);
    RUN_TEST(test_apply_gradient_int8_clamps_at_saturation);

    return UNITY_END();
}
