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

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

void softmax_forward(float* x, int batch_size, int seq_len) {
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int b = 0; b < batch_size; b++) {
        float* batch_scores = x + b * seq_len;
        float max_score = batch_scores[0];
        for (int i = 1; i < seq_len; i++) {
            if (batch_scores[i] > max_score) max_score = batch_scores[i];
        }
        float sum_exp = 0.0f;
        for (int i = 0; i < seq_len; i++) {
            batch_scores[i] = expf(batch_scores[i] - max_score);
            sum_exp += batch_scores[i];
        }
        if (sum_exp > 0.0f) {
            for (int i = 0; i < seq_len; i++) batch_scores[i] /= sum_exp;
        }
    }
}

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]);
}

void matrix_multiply(float* C, const float* A, const float* B, 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 = 0.0f;
            for (int l = 0; l < k; l++) sum += A[i * k + l] * B[l * n + j];
            C[i * n + j] = sum;
        }
    }
}

int bahdanau_attention_forward(float* decoder_hidden, float* encoder_outputs,
                              float* attention_weights, float* context_vector,
                              const float* w1, const float* w2, const float* v,
                              int batch_size, int seq_len, int hidden_size,
                              attention_type_t attention_type) {
    
    float* attention_scores = (float*)malloc(batch_size * seq_len * sizeof(float));
    if (!attention_scores) return -1;
    
    switch (attention_type) {
        case ATTENTION_CONCAT: {
            if (!w1 || !w2 || !v) { free(attention_scores); return -2; }
            #ifdef _OPENMP
            #pragma omp parallel for collapse(2)
            #endif
            for (int b = 0; b < batch_size; b++) {
                for (int s = 0; s < seq_len; s++) {
                    float score = 0.0f;
                    for (int h = 0; h < hidden_size; h++) {
                        float enc_contrib = 0.0f, dec_contrib = 0.0f;
                        for (int h2 = 0; h2 < hidden_size; h2++) {
                            enc_contrib += encoder_outputs[b * seq_len * hidden_size + s * hidden_size + h2] * w1[h2 * hidden_size + h];
                            dec_contrib += decoder_hidden[b * hidden_size + h2] * w2[h2 * hidden_size + h];
                        }
                        score += v[h] * tanhf(enc_contrib + dec_contrib);
                    }
                    attention_scores[b * seq_len + s] = score;
                }
            }
            break;
        }
        case ATTENTION_GENERAL: {
            if (!w1) { free(attention_scores); return -3; }
            int status = 0;
            #ifdef _OPENMP
            #pragma omp parallel for reduction(|:status)
            #endif
            for (int b = 0; b < batch_size; b++) {
                if (status) continue;
                float* transformed_enc = (float*)malloc(seq_len * hidden_size * sizeof(float));
                if (!transformed_enc) {
                    status = 1;
                } else {
                    matrix_multiply(transformed_enc, encoder_outputs + (b * seq_len * hidden_size), w1, seq_len, hidden_size, hidden_size);
                    for (int s = 0; s < seq_len; s++) {
                        float score = 0.0f;
                        for (int h = 0; h < hidden_size; h++) {
                            score += decoder_hidden[b * hidden_size + h] * transformed_enc[s * hidden_size + h];
                        }
                        attention_scores[b * seq_len + s] = score;
                    }
                    free(transformed_enc);
                }
            }
            if (status) { free(attention_scores); return -1; }
            break;
        }
        case ATTENTION_DOT:
        default: {
            #ifdef _OPENMP
            #pragma omp parallel for collapse(2)
            #endif
            for (int b = 0; b < batch_size; b++) {
                for (int s = 0; s < seq_len; s++) {
                    float score = 0.0f;
                    for (int h = 0; h < hidden_size; h++) {
                        score += decoder_hidden[b * hidden_size + h] * encoder_outputs[b * seq_len * hidden_size + s * hidden_size + h];
                    }
                    attention_scores[b * seq_len + s] = score;
                }
            }
            break;
        }
    }
    softmax_forward(attention_scores, batch_size, seq_len);
    memcpy(attention_weights, attention_scores, batch_size * seq_len * sizeof(float));
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (int b = 0; b < batch_size; b++) {
        for (int h = 0; h < hidden_size; h++) context_vector[b * hidden_size + h] = 0.0f;
        for (int s = 0; s < seq_len; s++) {
            float weight = attention_weights[b * seq_len + s];
            for (int h = 0; h < hidden_size; h++) {
                context_vector[b * hidden_size + h] += weight * encoder_outputs[b * seq_len * hidden_size + s * hidden_size + h];
            }
        }
    }
    free(attention_scores);
    return 0;
}