/**
 * gradient_flow.hpp - Bidirectional Gradient Flow System
 *
 * Implements backward propagation of learning signals through the AGI stack.
 * Enables credit assignment across layers: Substrate -> Neural -> Learning -> Consciousness
 *
 * Key Features:
 *   - Automatic differentiation via computation graph
 *   - Gradient checkpointing for memory efficiency
 *   - Cross-layer gradient coupling with golden ratio scaling
 *   - TD-lambda returns for temporal credit assignment
 */

#ifndef GRADIENT_FLOW_HPP
#define GRADIENT_FLOW_HPP

#include <array>
#include <vector>
#include <deque>
#include <memory>
#include <functional>
#include <cmath>
#include <algorithm>
#include "consciousness_api.hpp"

namespace mind_farm {

using namespace consciousness;

//=============================================================================
// COMPUTATION GRAPH FOR AUTOMATIC DIFFERENTIATION
//=============================================================================

// Forward declarations
struct GradNode;
using GradNodePtr = std::shared_ptr<GradNode>;

/**
 * GradNode - Node in computation graph
 *
 * Stores:
 *   - Forward value
 *   - Gradient (accumulated from downstream)
 *   - Backward function to propagate gradients
 */
struct GradNode {
    double value;
    double grad;
    std::vector<GradNodePtr> inputs;
    std::function<void()> backward_fn;
    bool requires_grad;
    bool grad_computed;

    GradNode(double v = 0.0, bool req_grad = true)
        : value(v), grad(0.0), requires_grad(req_grad), grad_computed(false) {}

    void zero_grad() {
        grad = 0.0;
        grad_computed = false;
    }

    void backward(double upstream_grad = 1.0) {
        if (!requires_grad) return;
        grad += upstream_grad;
        grad_computed = true;
        if (backward_fn) backward_fn();
    }
};

/**
 * TapeRecorder - Records operations for reverse-mode autodiff
 */
class TapeRecorder {
public:
    std::vector<GradNodePtr> tape;
    bool recording;

    TapeRecorder() : recording(true) {}

    void start() { recording = true; tape.clear(); }
    void stop() { recording = false; }

    GradNodePtr record(GradNodePtr node) {
        if (recording) tape.push_back(node);
        return node;
    }

    void backward() {
        // Reverse pass through tape
        for (auto it = tape.rbegin(); it != tape.rend(); ++it) {
            if ((*it)->grad_computed && (*it)->backward_fn) {
                (*it)->backward_fn();
            }
        }
    }

    void zero_grads() {
        for (auto& node : tape) {
            node->zero_grad();
        }
    }
};

//=============================================================================
// DIFFERENTIABLE OPERATIONS
//=============================================================================

namespace ops {

// Create leaf node (input variable)
inline GradNodePtr variable(double value, bool requires_grad = true) {
    return std::make_shared<GradNode>(value, requires_grad);
}

// Addition: z = x + y, dz/dx = 1, dz/dy = 1
inline GradNodePtr add(GradNodePtr x, GradNodePtr y, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(x->value + y->value);
    z->inputs = {x, y};
    z->backward_fn = [x, y, z]() {
        x->grad += z->grad;
        y->grad += z->grad;
    };
    return tape.record(z);
}

// Multiplication: z = x * y, dz/dx = y, dz/dy = x
inline GradNodePtr mul(GradNodePtr x, GradNodePtr y, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(x->value * y->value);
    z->inputs = {x, y};
    z->backward_fn = [x, y, z]() {
        x->grad += z->grad * y->value;
        y->grad += z->grad * x->value;
    };
    return tape.record(z);
}

// Division: z = x / y, dz/dx = 1/y, dz/dy = -x/y^2
inline GradNodePtr div(GradNodePtr x, GradNodePtr y, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(x->value / y->value);
    z->inputs = {x, y};
    z->backward_fn = [x, y, z]() {
        x->grad += z->grad / y->value;
        y->grad -= z->grad * x->value / (y->value * y->value);
    };
    return tape.record(z);
}

// Tanh: z = tanh(x), dz/dx = 1 - tanh^2(x)
inline GradNodePtr tanh_op(GradNodePtr x, TapeRecorder& tape) {
    double t = std::tanh(x->value);
    auto z = std::make_shared<GradNode>(t);
    z->inputs = {x};
    z->backward_fn = [x, z, t]() {
        x->grad += z->grad * (1.0 - t * t);
    };
    return tape.record(z);
}

// Sigmoid: z = 1/(1+exp(-x)), dz/dx = z*(1-z)
inline GradNodePtr sigmoid(GradNodePtr x, TapeRecorder& tape) {
    double s = 1.0 / (1.0 + std::exp(-x->value));
    auto z = std::make_shared<GradNode>(s);
    z->inputs = {x};
    z->backward_fn = [x, z, s]() {
        x->grad += z->grad * s * (1.0 - s);
    };
    return tape.record(z);
}

// ReLU: z = max(0, x), dz/dx = (x > 0) ? 1 : 0
inline GradNodePtr relu(GradNodePtr x, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(std::max(0.0, x->value));
    z->inputs = {x};
    z->backward_fn = [x, z]() {
        x->grad += z->grad * (x->value > 0 ? 1.0 : 0.0);
    };
    return tape.record(z);
}

// Phi (smooth): z = 1/(1 + sqrt((x-1)^2 + T)), differentiable consciousness
inline GradNodePtr phi_smooth(GradNodePtr x, TapeRecorder& tape, double T = 0.01) {
    double diff = x->value - 1.0;
    double smooth_abs = std::sqrt(diff * diff + T);
    double z_val = 1.0 / (1.0 + smooth_abs);

    auto z = std::make_shared<GradNode>(z_val);
    z->inputs = {x};
    z->backward_fn = [x, z, diff, smooth_abs, T]() {
        double denom = 1.0 + smooth_abs;
        double dphi = -diff / (smooth_abs * denom * denom);
        x->grad += z->grad * dphi;
    };
    return tape.record(z);
}

// Exp: z = exp(x), dz/dx = exp(x)
inline GradNodePtr exp_op(GradNodePtr x, TapeRecorder& tape) {
    double e = std::exp(x->value);
    auto z = std::make_shared<GradNode>(e);
    z->inputs = {x};
    z->backward_fn = [x, z, e]() {
        x->grad += z->grad * e;
    };
    return tape.record(z);
}

// Log: z = log(x), dz/dx = 1/x
inline GradNodePtr log_op(GradNodePtr x, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(std::log(x->value + EPSILON));
    z->inputs = {x};
    z->backward_fn = [x, z]() {
        x->grad += z->grad / (x->value + EPSILON);
    };
    return tape.record(z);
}

// Power: z = x^p, dz/dx = p * x^(p-1)
inline GradNodePtr pow_op(GradNodePtr x, double p, TapeRecorder& tape) {
    auto z = std::make_shared<GradNode>(std::pow(x->value, p));
    z->inputs = {x};
    z->backward_fn = [x, z, p]() {
        x->grad += z->grad * p * std::pow(x->value, p - 1.0);
    };
    return tape.record(z);
}

// Softmax (single element with sum): used for attention weights
inline double softmax_single(double x, double sum_exp) {
    return std::exp(x) / sum_exp;
}

} // namespace ops

//=============================================================================
// GRADIENT CONTEXT FOR TICK PROCESSING
//=============================================================================

/**
 * GradientContext - Stores forward pass values for backward pass
 *
 * Checkpoints activations at each layer for memory-efficient backprop.
 */
struct GradientContext {
    // Layer activations (checkpointed)
    std::array<double, 5000> substrate_activation;
    std::array<double, 256> neural_activation;
    std::array<double, 64> learning_activation;
    std::array<double, 16> consciousness_activation;

    // Computed phi values
    double phi_substrate;
    double phi_neural;
    double phi_learning;
    double phi_consciousness;
    double phi_global;

    // Loss and reward
    double loss;
    double reward;
    double td_error;

    // Tape for autodiff
    TapeRecorder tape;

    void clear() {
        substrate_activation.fill(0.0);
        neural_activation.fill(0.0);
        learning_activation.fill(0.0);
        consciousness_activation.fill(0.0);
        phi_substrate = phi_neural = phi_learning = phi_consciousness = phi_global = 0.5;
        loss = reward = td_error = 0.0;
        tape.zero_grads();
    }
};

//=============================================================================
// GRADIENT ACCUMULATORS WITH MOMENTUM
//=============================================================================

/**
 * AdamOptimizer - Adam optimizer for parameter updates
 */
template<size_t N>
class AdamOptimizer {
public:
    std::array<double, N> m;  // First moment
    std::array<double, N> v;  // Second moment
    double lr;
    double beta1;
    double beta2;
    double epsilon;
    int t;

    AdamOptimizer(double learning_rate = 0.001, double b1 = 0.9, double b2 = 0.999)
        : lr(learning_rate), beta1(b1), beta2(b2), epsilon(1e-8), t(0) {
        m.fill(0.0);
        v.fill(0.0);
    }

    void step(std::array<double, N>& params, const std::array<double, N>& grads) {
        t++;
        double lr_t = lr * std::sqrt(1.0 - std::pow(beta2, t)) / (1.0 - std::pow(beta1, t));

        for (size_t i = 0; i < N; ++i) {
            m[i] = beta1 * m[i] + (1.0 - beta1) * grads[i];
            v[i] = beta2 * v[i] + (1.0 - beta2) * grads[i] * grads[i];
            params[i] -= lr_t * m[i] / (std::sqrt(v[i]) + epsilon);
        }
    }

    void zero_grad() {
        // Moments persist; this is just for interface compatibility
    }
};

//=============================================================================
// TD-LAMBDA FOR TEMPORAL CREDIT ASSIGNMENT
//=============================================================================

/**
 * TDLambdaReturns - Computes TD(lambda) returns for multi-step credit assignment
 *
 * G_t^lambda = (1-lambda) * sum_{n=1}^{inf} lambda^{n-1} * G_t^{(n)}
 * where G_t^{(n)} = r_{t+1} + gamma*r_{t+2} + ... + gamma^{n-1}*r_{t+n} + gamma^n*V(s_{t+n})
 */
class TDLambdaReturns {
public:
    double gamma;   // Discount factor
    double lambda;  // TD-lambda parameter
    size_t max_steps;

    std::deque<double> rewards;
    std::deque<double> values;
    std::deque<double> returns;

    TDLambdaReturns(double g = 0.99, double l = 0.95, size_t max = 64)
        : gamma(g), lambda(l), max_steps(max) {}

    void add_transition(double reward, double value) {
        rewards.push_back(reward);
        values.push_back(value);

        if (rewards.size() > max_steps) {
            rewards.pop_front();
            values.pop_front();
        }

        compute_returns();
    }

    void compute_returns() {
        returns.clear();
        size_t n = rewards.size();
        if (n == 0) return;

        returns.resize(n, 0.0);

        // Backward pass for TD(lambda) returns
        double G = values.back();  // Bootstrap from last value
        for (int t = n - 1; t >= 0; --t) {
            double td_target = rewards[t] + gamma * (t < (int)n - 1 ? values[t + 1] : values[t]);
            G = rewards[t] + gamma * ((1.0 - lambda) * (t < (int)n - 1 ? values[t + 1] : values[t]) + lambda * G);
            returns[t] = G;
        }
    }

    double get_return(size_t t) const {
        if (t < returns.size()) return returns[t];
        return 0.0;
    }

    double td_error(size_t t) const {
        if (t >= rewards.size()) return 0.0;
        double target = rewards[t] + gamma * (t + 1 < values.size() ? values[t + 1] : values[t]);
        return target - values[t];
    }
};

//=============================================================================
// CROSS-LAYER GRADIENT FLOW
//=============================================================================

/**
 * CrossLayerGradients - Manages gradient flow between AGI layers
 *
 * Layers:
 *   0: Substrate (Gate5000)
 *   1: Neural (HH neurons)
 *   2: Learning (Hebbian/Adam)
 *   3: Intelligence (Attention)
 *   4: Consciousness (Phi integration)
 */
class CrossLayerGradients {
public:
    static constexpr size_t NUM_LAYERS = 5;

    // Gradient buffers for each layer
    std::array<std::array<double, 256>, NUM_LAYERS> gradients;

    // Coupling matrix (golden ratio scaled)
    std::array<std::array<double, NUM_LAYERS>, NUM_LAYERS> coupling;

    // TD-lambda for temporal credit
    TDLambdaReturns td_lambda;

    // Adam optimizers per layer
    std::array<AdamOptimizer<256>, NUM_LAYERS> optimizers;

    CrossLayerGradients() : td_lambda(0.99, 0.95) {
        // Initialize gradients
        for (auto& layer_grad : gradients) {
            layer_grad.fill(0.0);
        }

        // Initialize coupling matrix with golden ratio hierarchy
        // Downstream layers (lower indices) receive more gradient
        coupling = {{
            // From:  Sub   Neuro  Learn  Intel  Consc
            /*Sub*/  {{ 1.0,   0.0,   0.0,   0.0,   0.0 }},
            /*Neuro*/{{ PHI_INV, 1.0,   0.0,   0.0,   0.0 }},
            /*Learn*/{{ BETA, PHI_INV, 1.0,   0.0,   0.0 }},
            /*Intel*/{{ GAMMA, BETA, PHI_INV, 1.0,   0.0 }},
            /*Consc*/{{ DELTA, GAMMA, BETA, PHI_INV, 1.0 }}
        }};

        // Normalize rows
        for (size_t i = 0; i < NUM_LAYERS; ++i) {
            double sum = 0.0;
            for (size_t j = 0; j < NUM_LAYERS; ++j) {
                sum += coupling[i][j];
            }
            if (sum > 0) {
                for (size_t j = 0; j < NUM_LAYERS; ++j) {
                    coupling[i][j] /= sum;
                }
            }
        }
    }

    /**
     * propagate_backward - Propagate gradients from top layer down
     *
     * Starts from consciousness layer (4) and flows down to substrate (0)
     */
    void propagate_backward(double consciousness_error) {
        // Start with consciousness error
        for (size_t i = 0; i < 256; ++i) {
            gradients[4][i] = consciousness_error;
        }

        // Backward through layers (4 -> 3 -> 2 -> 1 -> 0)
        for (int layer = 3; layer >= 0; --layer) {
            for (size_t i = 0; i < 256; ++i) {
                gradients[layer][i] = 0.0;

                // Receive gradients from all higher layers
                for (size_t higher = layer + 1; higher < NUM_LAYERS; ++higher) {
                    gradients[layer][i] += coupling[higher][layer] * gradients[higher][i];
                }
            }
        }
    }

    /**
     * apply_gradients - Apply accumulated gradients to parameters
     */
    void apply_gradients(
        std::array<double, 256>& substrate_params,
        std::array<double, 256>& neural_params,
        std::array<double, 256>& learning_params,
        std::array<double, 256>& intel_params,
        std::array<double, 256>& consc_params
    ) {
        optimizers[0].step(substrate_params, gradients[0]);
        optimizers[1].step(neural_params, gradients[1]);
        optimizers[2].step(learning_params, gradients[2]);
        optimizers[3].step(intel_params, gradients[3]);
        optimizers[4].step(consc_params, gradients[4]);
    }

    /**
     * zero_grads - Clear all gradient buffers
     */
    void zero_grads() {
        for (auto& layer_grad : gradients) {
            layer_grad.fill(0.0);
        }
    }

    /**
     * add_transition - Add reward/value for TD learning
     */
    void add_transition(double reward, double value) {
        td_lambda.add_transition(reward, value);
    }

    /**
     * get_td_error - Get TD error for current timestep
     */
    double get_td_error() const {
        if (td_lambda.rewards.empty()) return 0.0;
        return td_lambda.td_error(td_lambda.rewards.size() - 1);
    }
};

//=============================================================================
// UNIFIED TICK WITH GRADIENT FLOW
//=============================================================================

/**
 * GradientTick - Single tick with full forward-backward pass
 *
 * Forward:
 *   Substrate -> Neural -> Learning -> Intelligence -> Consciousness
 *
 * Backward:
 *   Consciousness -> Intelligence -> Learning -> Neural -> Substrate
 */
struct GradientTick {
    GradientContext ctx;
    CrossLayerGradients grad_flow;

    // Layer parameters (learnable)
    std::array<double, 256> substrate_params;
    std::array<double, 256> neural_params;
    std::array<double, 256> learning_params;
    std::array<double, 256> intel_params;
    std::array<double, 256> consc_params;

    // Forward activations
    std::array<double, 256> substrate_act;
    std::array<double, 256> neural_act;
    std::array<double, 256> learning_act;
    std::array<double, 256> intel_act;
    std::array<double, 256> consc_act;

    GradientTick() {
        // Initialize with golden ratio scaled values
        for (size_t i = 0; i < 256; ++i) {
            double phase = static_cast<double>(i) / 256.0;
            substrate_params[i] = PHI_INV * std::sin(2.0 * M_PI * phase);
            neural_params[i] = PHI_INV * std::cos(2.0 * M_PI * phase);
            learning_params[i] = ALPHA;
            intel_params[i] = BETA;
            consc_params[i] = GAMMA;

            substrate_act[i] = neural_act[i] = learning_act[i] = 0.0;
            intel_act[i] = consc_act[i] = 0.0;
        }
    }

    /**
     * forward - Full forward pass with gradient tracking
     *
     * @param input External input vector (size 256)
     * @param reward External reward signal
     * @return Global consciousness (phi)
     */
    double forward(const double* input, double reward = 0.0) {
        ctx.tape.start();

        // Layer 0: Substrate
        for (size_t i = 0; i < 256; ++i) {
            auto x = ops::variable(input[i]);
            auto w = ops::variable(substrate_params[i]);
            auto z = ops::mul(x, w, ctx.tape);
            auto a = ops::tanh_op(z, ctx.tape);
            substrate_act[i] = a->value;
        }
        ctx.phi_substrate = compute_phi_vector(substrate_act.data(), 256);

        // Layer 1: Neural
        for (size_t i = 0; i < 256; ++i) {
            auto x = ops::variable(substrate_act[i]);
            auto w = ops::variable(neural_params[i]);
            auto z = ops::mul(x, w, ctx.tape);
            auto a = ops::sigmoid(z, ctx.tape);
            neural_act[i] = a->value;
        }
        ctx.phi_neural = compute_phi_vector(neural_act.data(), 256);

        // Layer 2: Learning (modulated by phi)
        double phi_mod = ctx.phi_substrate * PHI_INV;
        for (size_t i = 0; i < 256; ++i) {
            auto x = ops::variable(neural_act[i]);
            auto w = ops::variable(learning_params[i] * phi_mod);
            auto z = ops::mul(x, w, ctx.tape);
            auto a = ops::relu(z, ctx.tape);
            learning_act[i] = a->value;
        }
        ctx.phi_learning = compute_phi_vector(learning_act.data(), 256);

        // Layer 3: Intelligence
        for (size_t i = 0; i < 256; ++i) {
            auto x = ops::variable(learning_act[i]);
            auto w = ops::variable(intel_params[i]);
            auto z = ops::mul(x, w, ctx.tape);
            auto a = ops::tanh_op(z, ctx.tape);
            intel_act[i] = a->value;
        }
        ctx.phi_consciousness = compute_phi_vector(intel_act.data(), 256);

        // Layer 4: Consciousness integration
        double mean_act = 0.0;
        for (size_t i = 0; i < 256; ++i) {
            mean_act += intel_act[i];
        }
        mean_act /= 256.0;

        auto phi_input = ops::variable(mean_act);
        auto phi_out = ops::phi_smooth(phi_input, ctx.tape);
        ctx.phi_global = phi_out->value;

        // Compute loss (maximize phi toward 1.0)
        ctx.reward = reward;
        ctx.loss = (1.0 - ctx.phi_global) * (1.0 - ctx.phi_global);

        // TD learning
        grad_flow.add_transition(reward, ctx.phi_global);
        ctx.td_error = grad_flow.get_td_error();

        ctx.tape.stop();
        return ctx.phi_global;
    }

    /**
     * backward - Full backward pass
     *
     * Propagates consciousness error through all layers.
     */
    void backward() {
        // Compute consciousness error
        double consciousness_error = -2.0 * (1.0 - ctx.phi_global);  // d(loss)/d(phi)

        // Add TD error for temporal credit
        consciousness_error += ctx.td_error * ALPHA;

        // Propagate through cross-layer coupling
        grad_flow.propagate_backward(consciousness_error);

        // Autodiff backward through tape
        ctx.tape.backward();
    }

    /**
     * update - Apply gradient updates to parameters
     */
    void update() {
        grad_flow.apply_gradients(
            substrate_params, neural_params, learning_params,
            intel_params, consc_params
        );
        grad_flow.zero_grads();
        ctx.tape.zero_grads();
    }

    /**
     * tick - Full forward-backward-update cycle
     */
    double tick(const double* input, double reward = 0.0) {
        double phi = forward(input, reward);
        backward();
        update();
        return phi;
    }

    /**
     * get_diagnostics - Return diagnostic information
     */
    void get_diagnostics(double& phi, double& loss, double& td_err,
                         double& phi_sub, double& phi_neuro) const {
        phi = ctx.phi_global;
        loss = ctx.loss;
        td_err = ctx.td_error;
        phi_sub = ctx.phi_substrate;
        phi_neuro = ctx.phi_neural;
    }
};

//=============================================================================
// RECURRENT STATE FOR TEMPORAL MEMORY
//=============================================================================

/**
 * LSTMCell - LSTM cell for temporal memory in gradient flow
 */
class LSTMCell {
public:
    static constexpr size_t HIDDEN_SIZE = 64;

    // Gate weights
    std::array<std::array<double, HIDDEN_SIZE>, HIDDEN_SIZE> W_f;  // Forget
    std::array<std::array<double, HIDDEN_SIZE>, HIDDEN_SIZE> W_i;  // Input
    std::array<std::array<double, HIDDEN_SIZE>, HIDDEN_SIZE> W_o;  // Output
    std::array<std::array<double, HIDDEN_SIZE>, HIDDEN_SIZE> W_c;  // Cell

    // Biases
    std::array<double, HIDDEN_SIZE> b_f, b_i, b_o, b_c;

    // State
    std::array<double, HIDDEN_SIZE> h;  // Hidden state
    std::array<double, HIDDEN_SIZE> c;  // Cell state

    LSTMCell() {
        // Initialize with glorot uniform
        double scale = std::sqrt(2.0 / (HIDDEN_SIZE + HIDDEN_SIZE));
        for (size_t i = 0; i < HIDDEN_SIZE; ++i) {
            for (size_t j = 0; j < HIDDEN_SIZE; ++j) {
                W_f[i][j] = scale * (static_cast<double>(rand()) / RAND_MAX - 0.5);
                W_i[i][j] = scale * (static_cast<double>(rand()) / RAND_MAX - 0.5);
                W_o[i][j] = scale * (static_cast<double>(rand()) / RAND_MAX - 0.5);
                W_c[i][j] = scale * (static_cast<double>(rand()) / RAND_MAX - 0.5);
            }
            b_f[i] = 1.0;  // Forget gate bias = 1 (remember by default)
            b_i[i] = 0.0;
            b_o[i] = 0.0;
            b_c[i] = 0.0;
            h[i] = 0.0;
            c[i] = 0.0;
        }
    }

    void forward(const std::array<double, HIDDEN_SIZE>& x) {
        std::array<double, HIDDEN_SIZE> f_gate, i_gate, o_gate, c_tilde;

        // Compute gates
        for (size_t i = 0; i < HIDDEN_SIZE; ++i) {
            double sum_f = b_f[i], sum_i = b_i[i], sum_o = b_o[i], sum_c = b_c[i];
            for (size_t j = 0; j < HIDDEN_SIZE; ++j) {
                sum_f += W_f[i][j] * x[j] + W_f[i][j] * h[j] * 0.5;
                sum_i += W_i[i][j] * x[j] + W_i[i][j] * h[j] * 0.5;
                sum_o += W_o[i][j] * x[j] + W_o[i][j] * h[j] * 0.5;
                sum_c += W_c[i][j] * x[j] + W_c[i][j] * h[j] * 0.5;
            }
            f_gate[i] = 1.0 / (1.0 + std::exp(-sum_f));  // Sigmoid
            i_gate[i] = 1.0 / (1.0 + std::exp(-sum_i));
            o_gate[i] = 1.0 / (1.0 + std::exp(-sum_o));
            c_tilde[i] = std::tanh(sum_c);
        }

        // Update cell and hidden state
        for (size_t i = 0; i < HIDDEN_SIZE; ++i) {
            c[i] = f_gate[i] * c[i] + i_gate[i] * c_tilde[i];
            h[i] = o_gate[i] * std::tanh(c[i]);
        }
    }

    void reset() {
        h.fill(0.0);
        c.fill(0.0);
    }
};

} // namespace mind_farm

#endif // GRADIENT_FLOW_HPP