#include "_rnn_ops.h"
#include <math.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

#ifdef _OPENMP
#include <omp.h>
#endif

void sigmoid_forward(float* x, int n) {
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int i = 0; i < n; i++) x[i] = 1.0f / (1.0f + expf(-x[i]));
}

void tanh_forward(float* x, int n) {
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int i = 0; i < n; i++) x[i] = tanhf(x[i]);
}

static void matmul_bias(float* C, const float* A, const float* B, const float* bias, int m, int n, int k) {
    #ifdef _OPENMP
    #pragma omp parallel for collapse(2)
    #endif
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            float sum = bias[j];
            for (int l = 0; l < k; l++) sum += A[i * k + l] * B[l * n + j];
            C[i * n + j] = sum;
        }
    }
}

void lstm_cell_forward(float* input, float* hidden, float* cell, float* output, const float* w_ih, const float* w_hh, const float* b_ih, const float* b_hh, int batch_size, int input_size, int hidden_size) {
    int gate_size = 4 * hidden_size;
    float* gates = (float*)malloc(batch_size * gate_size * sizeof(float));
    if (!gates) return;
    
    matmul_bias(gates, input, w_ih, b_ih, batch_size, gate_size, input_size);
    
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int g = 0; g < gate_size; g++) {
            float h_sum = b_hh[g];
            for (int h = 0; h < hidden_size; h++) h_sum += hidden[b * hidden_size + h] * w_hh[h * gate_size + g];
            gates[b * gate_size + g] += h_sum;
        }
    }
    
    #ifdef _OPENMP
    #pragma omp parallel for collapse(2)
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int h = 0; h < hidden_size; h++) {
            int g_idx = b * gate_size + h;
            float i = 1.0f / (1.0f + expf(-gates[g_idx]));
            float f = 1.0f / (1.0f + expf(-gates[g_idx + hidden_size]));
            float g = tanhf(gates[g_idx + 2 * hidden_size]);
            float o = 1.0f / (1.0f + expf(-gates[g_idx + 3 * hidden_size]));
            
            int c_idx = b * hidden_size + h;
            cell[c_idx] = f * cell[c_idx] + i * g;
            output[c_idx] = o * tanhf(cell[c_idx]);
        }
    }
    free(gates);
}

void gru_cell_forward(float* input, float* hidden, float* output, const float* w_ih, const float* w_hh, const float* b_ih, const float* b_hh, int batch_size, int input_size, int hidden_size) {
    int gate_size = 3 * hidden_size;
    float* gates = (float*)malloc(batch_size * gate_size * sizeof(float));
    if (!gates) return;
    
    matmul_bias(gates, input, w_ih, b_ih, batch_size, 2 * hidden_size, input_size);
    
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int g = 0; g < 2 * hidden_size; g++) {
            float h_sum = b_hh[g];
            for (int h = 0; h < hidden_size; h++) h_sum += hidden[b * hidden_size + h] * w_hh[h * gate_size + g];
            gates[b * gate_size + g] += h_sum;
        }
    }
    
    #ifdef _OPENMP
    #pragma omp parallel for collapse(2)
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int h = 0; h < 2 * hidden_size; h++) {
            int idx = b * gate_size + h;
            gates[idx] = 1.0f / (1.0f + expf(-gates[idx]));
        }
    }
    
    const float* w_ih_new = w_ih + (2 * hidden_size * input_size);
    const float* b_ih_new = b_ih + (2 * hidden_size);
    float* new_gates = gates + (2 * hidden_size);
    matmul_bias(new_gates, input, w_ih_new, b_ih_new, batch_size, hidden_size, input_size);
    
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int b = 0; b < batch_size; b++) {
        float* r_gate = gates + (b * gate_size);
        for (int h = 0; h < hidden_size; h++) {
            float r_sum = b_hh[2 * hidden_size + h];
            for (int h2 = 0; h2 < hidden_size; h2++) {
                r_sum += (hidden[b * hidden_size + h2] * r_gate[h2]) * w_hh[h2 * gate_size + 2 * hidden_size + h];
            }
            new_gates[b * gate_size + h] += r_sum;
        }
    }
    
    #ifdef _OPENMP
    #pragma omp parallel for collapse(2)
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int h = 0; h < hidden_size; h++) {
            int g_idx = b * gate_size;
            int o_idx = b * hidden_size + h;
            float update = gates[g_idx + hidden_size + h];
            output[o_idx] = (1.0f - update) * tanhf(new_gates[g_idx + h]) + update * hidden[o_idx];
        }
    }
    free(gates);
}

int rnn_forward(float* input, float* hidden, float* cell, float* output, float* final_hidden, float* final_cell, const float* weights, const float* biases, int batch_size, int seq_len, int input_size, int hidden_size, int num_layers, rnn_type_t rnn_type, int bidirectional) {
    if (bidirectional) return -1;
    int w_layer = (rnn_type == RNN_TYPE_LSTM ? 4 : (rnn_type == RNN_TYPE_GRU ? 3 : 1)) * hidden_size * (input_size + hidden_size);
    int b_layer = (rnn_type == RNN_TYPE_LSTM ? 8 : (rnn_type == RNN_TYPE_GRU ? 6 : 2)) * hidden_size;

    for (int layer = 0; layer < num_layers; layer++) {
        const float* lw = weights + layer * w_layer;
        const float* lb = biases + layer * b_layer;
        float* li = (layer == 0) ? input : output;
        int lis = (layer == 0) ? input_size : hidden_size;
        float* lh = hidden + layer * batch_size * hidden_size;
        float* lc = cell + layer * batch_size * hidden_size;

        for (int t = 0; t < seq_len; t++) {
            float* si = li + t * batch_size * lis;
            float* so = output + t * batch_size * hidden_size;
            if (rnn_type == RNN_TYPE_LSTM) {
                lstm_cell_forward(si, lh, lc, so, lw, lw + 4 * hidden_size * lis, lb, lb + 4 * hidden_size, batch_size, lis, hidden_size);
            } else if (rnn_type == RNN_TYPE_GRU) {
                gru_cell_forward(si, lh, so, lw, lw + 3 * hidden_size * lis, lb, lb + 3 * hidden_size, batch_size, lis, hidden_size);
            } else {
                matmul_bias(so, si, lw, lb, batch_size, hidden_size, lis);
                #ifdef _OPENMP
                #pragma omp parallel for
                #endif
                for (int b = 0; b < batch_size; b++) {
                    for (int h = 0; h < hidden_size; h++) {
                        float s = lb[hidden_size + h];
                        for (int h2 = 0; h2 < hidden_size; h2++) s += lh[b * hidden_size + h2] * lw[hidden_size * lis + h2 * hidden_size + h];
                        so[b * hidden_size + h] = tanhf(so[b * hidden_size + h] + s);
                    }
                }
            }
            memcpy(lh, so, batch_size * hidden_size * sizeof(float));
        }
        memcpy(final_hidden + layer * batch_size * hidden_size, lh, batch_size * hidden_size * sizeof(float));
        if (rnn_type == RNN_TYPE_LSTM) memcpy(final_cell + layer * batch_size * hidden_size, lc, batch_size * hidden_size * sizeof(float));
    }
    return 0;
}