// ============================================================================
// cactus_backprop.cpp — see cactus_backprop.h for design notes.
// ============================================================================

#include "cactus_backprop.h"

#include <cstring>
#include <cmath>

bool cactus_bp_forward(const CactusBpLayer* layers, int n_layers,
                       const float* input, int input_dim,
                       float activations_out[][CACTUS_BP_MAX_DIM]) {
    if (!layers || n_layers <= 0 || n_layers > CACTUS_BP_MAX_LAYERS) return false;
    if (input_dim <= 0 || input_dim > CACTUS_BP_MAX_DIM) return false;
    if (layers[0].in_dim != input_dim) return false;

    memset(activations_out[0], 0, sizeof(float) * CACTUS_BP_MAX_DIM);
    memcpy(activations_out[0], input, sizeof(float) * (size_t)input_dim);

    for (int i = 0; i < n_layers; i++) {
        const CactusBpLayer& L = layers[i];
        const float* x = activations_out[i];
        float* y = activations_out[i + 1];
        memset(y, 0, sizeof(float) * CACTUS_BP_MAX_DIM);

        if (L.type == CBP_LINEAR_FP32) {
            if (L.in_dim > CACTUS_BP_MAX_DIM || L.out_dim > CACTUS_BP_MAX_DIM) return false;
            const float* W = (const float*)L.W;
            for (int m = 0; m < L.out_dim; m++) y[m] = L.b[m];
            for (int k = 0; k < L.in_dim; k++) {
                float xk = x[k];
                for (int m = 0; m < L.out_dim; m++) {
                    y[m] += xk * W[k * L.out_dim + m];
                }
            }
        } else if (L.type == CBP_LINEAR_INT8) {
            if (L.in_dim > CACTUS_BP_MAX_DIM || L.out_dim > CACTUS_BP_MAX_DIM) return false;
            const int8_t* W = (const int8_t*)L.W;
            for (int m = 0; m < L.out_dim; m++) y[m] = L.b[m];
            for (int k = 0; k < L.in_dim; k++) {
                float xk = x[k];
                for (int m = 0; m < L.out_dim; m++) {
                    y[m] += xk * ((float)W[k * L.out_dim + m] * L.w_scale);
                }
            }
        } else if (L.type == CBP_RELU) {
            for (int j = 0; j < L.in_dim; j++) y[j] = x[j] > 0.0f ? x[j] : 0.0f;
        } else if (L.type == CBP_SOFTMAX) {
            float mx = x[0];
            for (int j = 1; j < L.in_dim; j++) if (x[j] > mx) mx = x[j];
            float sum = 0.0f;
            for (int j = 0; j < L.in_dim; j++) {
                y[j] = expf(x[j] - mx);
                sum += y[j];
            }
            for (int j = 0; j < L.in_dim; j++) y[j] /= sum;
        } else {
            return false;
        }
    }
    return true;
}

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[]) {
    if (!layers || n_layers <= 0) return -1.0f;

    const CactusBpLayer& last = layers[n_layers - 1];
    bool ends_in_softmax = (last.type == CBP_SOFTMAX);

    // dL/dz at the input to the final layer. If the network ends in
    // softmax (the only case this module supports -- see header), use the
    // fused softmax+cross-entropy shortcut: dL/dz = softmax_out -
    // onehot(target). This is the gradient with respect to the SOFTMAX
    // LAYER'S INPUT, so backprop below starts at layer n_layers-2 (the
    // softmax layer itself is fully absorbed into this one step, matching
    // ends_in_softmax == true meaning "one fewer explicit backward layer").
    float grad[CACTUS_BP_MAX_LAYERS + 1][CACTUS_BP_MAX_DIM];
    int start_layer;   // last layer index whose OWN backward() still needs to run
    float loss;

    if (ends_in_softmax) {
        const float* probs = activations[n_layers];   // post-softmax output
        int dim = last.in_dim;   // softmax in_dim == out_dim
        float p_target = probs[target_class] > 1e-12f ? probs[target_class] : 1e-12f;
        loss = -logf(p_target);
        for (int j = 0; j < dim; j++) {
            grad[n_layers - 1][j] = probs[j] - (j == target_class ? 1.0f : 0.0f);
        }
        start_layer = n_layers - 2;
    } else {
        // No fused shortcut available for a non-softmax-terminated network
        // -- not supported by this module (see header scope note).
        return -1.0f;
    }

    // Standard layer-by-layer backprop for everything before the (already
    // absorbed) final softmax.
    for (int i = start_layer; i >= 0; i--) {
        const CactusBpLayer& L = layers[i];
        const float* x = activations[i];       // this layer's input
        const float* dLdy = grad[i + 1];        // gradient w.r.t. this layer's output

        if (L.type == CBP_LINEAR_FP32 || L.type == CBP_LINEAR_INT8) {
            if (grad_W_out[i]) {
                for (int k = 0; k < L.in_dim; k++) {
                    for (int m = 0; m < L.out_dim; m++) {
                        grad_W_out[i][k * L.out_dim + m] = x[k] * dLdy[m];
                    }
                }
            }
            if (grad_b_out[i]) {
                for (int m = 0; m < L.out_dim; m++) grad_b_out[i][m] = dLdy[m];
            }
            // dL/dx[k] = sum_m dL/dy[m] * W[k,m]
            for (int k = 0; k < L.in_dim; k++) {
                float acc = 0.0f;
                if (L.type == CBP_LINEAR_FP32) {
                    const float* W = (const float*)L.W;
                    for (int m = 0; m < L.out_dim; m++) acc += dLdy[m] * W[k * L.out_dim + m];
                } else {
                    const int8_t* W = (const int8_t*)L.W;
                    for (int m = 0; m < L.out_dim; m++) acc += dLdy[m] * ((float)W[k * L.out_dim + m] * L.w_scale);
                }
                grad[i][k] = acc;
            }
        } else if (L.type == CBP_RELU) {
            // Pre-activation for a ReLU layer is simply its own input x.
            for (int j = 0; j < L.in_dim; j++) {
                grad[i][j] = (x[j] > 0.0f) ? dLdy[j] : 0.0f;
            }
        } else {
            // A softmax layer can only legally appear last (checked
            // above); anything else reaching here is an unsupported
            // architecture for this module.
            return -1.0f;
        }
    }

    return loss;
}

void cactus_bp_apply_gradient(CactusBpLayer* layer,
                              const float* grad_W, const float* grad_b,
                              float learning_rate) {
    if (!layer || !grad_W || !grad_b) return;
    int n_w = layer->in_dim * layer->out_dim;

    if (layer->type == CBP_LINEAR_FP32) {
        float* W = (float*)layer->W;
        for (int i = 0; i < n_w; i++) W[i] -= learning_rate * grad_W[i];
    } else if (layer->type == CBP_LINEAR_INT8) {
        int8_t* W = (int8_t*)layer->W;
        float scale = layer->w_scale;
        if (scale == 0.0f) return;   // can't dequant/requant against a zero scale
        for (int i = 0; i < n_w; i++) {
            float effective = (float)W[i] * scale - learning_rate * grad_W[i];
            float q = effective / scale;
            int qi = (int)(q >= 0.0f ? q + 0.5f : q - 0.5f);   // round-half-away-from-zero
            if (qi > 127) qi = 127;
            if (qi < -127) qi = -127;
            W[i] = (int8_t)qi;
        }
    } else {
        return;   // non-linear layers have no weights to update
    }

    for (int i = 0; i < layer->out_dim; i++) {
        layer->b[i] -= learning_rate * grad_b[i];
    }
}
