#include "ekd_blas.h"

void print_vector(float* vector, uint32_t size) {
    for (uint32_t i=0; i<size; i++) printf("%.2f ", vector[i]);
}

void print_matrix(float** matrix, uint32_t rows, uint32_t cols) {
    for (uint32_t i=0; i<rows; i++) {
        printf("\n");
        for (uint32_t j=0; j<cols; j++) printf("%.2f ", matrix[i][j]);
    }
}

void copy_vector(float* input_vector, float* resultant_vector, uint32_t size) {
    for (uint32_t i=0; i<size; i++) resultant_vector[i] = input_vector[i];
}

void copy_matrix(float** input_matrix, float** resultant_matrix, uint32_t rows, uint32_t cols) {
    for (uint32_t i=0; i<rows; i++){
        for (uint32_t j=0; j<cols; j++) resultant_matrix[i][j] = input_matrix[i][j];
    }
}

float trace_matrix(float** input_matrix, uint32_t size) {
    float trace = 0.0f;
    for (uint32_t i=0; i<size; i++) {
        for (uint32_t j=0; j< size; j++) {
            if (i == j) trace += input_matrix[i][j];
        }
   }

   return trace;
}

void shift_matrix(float** input_matrix, uint32_t size, float trace, shift_matrix_operation_e operation) {
    for (uint32_t i=0; i<size; i++) {
        for (uint32_t j=0; j<size; j++) {
            if (i == j) {
                if (operation == SHIFT_ADD) {
                    input_matrix[i][j] += trace;
                } else if (operation == SHIFT_SUBTRACT) {
                    input_matrix[i][j] -= trace;
                }
            }
        }
    }
}

void eye(float** input_matrix, uint32_t size) {
    for (uint32_t i=0; i<size; i++) {
        for (uint32_t j=0; j<size; j++) input_matrix[i][j] = (i == j) ? 1 : 0;
    }
}

void outer_product(float* input_vector, uint32_t size, float** resultant_matrix) {
    // for (uint32_t i=0; i<size; i++) {
    //     for (uint32_t j=0; j<size; j++) resultant_matrix[i][j] = input_vector[i] * input_vector[j]; 
    // }
    for (uint32_t i = 0; i < size; i++) {
        for (uint32_t j = i; j < size; j++) {
            float value = input_vector[i] * input_vector[j];
            resultant_matrix[i][j] = value;
            resultant_matrix[j][i] = value;
        }
    }
}

float inner_product(float* input_vector_1, float* input_vector_2, uint32_t size) {
    float resultant_scalar = 0.0;
    for (uint32_t i=0; i<size; i++) resultant_scalar += input_vector_1[i] * input_vector_2[i];
    return resultant_scalar;
}

void transpose(float** input_matrix, uint32_t rows, uint32_t cols, float** transposed_matrix) {
    for (uint32_t i=0; i<rows; i++){
        for (uint32_t j=0; j<cols; j++) transposed_matrix[j][i] = input_matrix[i][j];
    }
}

void matrix_operation(float** input_matrix_A, uint32_t rows_A, uint32_t cols_A, float** input_matrix_B, uint32_t rows_B, uint32_t cols_B, float** resultant_matrix, matrix_matrix_operation_e operation) {
        if (operation == MM_ADD) {
            for (uint32_t i=0; i<rows_A; i++) {
                for (uint32_t j=0; j<cols_A; j++) resultant_matrix[i][j] = input_matrix_A[i][j] + input_matrix_B[i][j];
            }
        } else if (operation == MM_SUBTRACT) {
            for (uint32_t i=0; i<rows_A; i++) {
                for (uint32_t j=0; j<cols_A; j++) resultant_matrix[i][j] = input_matrix_A[i][j] - input_matrix_B[i][j];
            }
        } else if (operation == MM_MULTIPLY) {
            for (uint32_t i=0; i<rows_A; i++) {
                for (uint32_t j=0; j<cols_B; j++) {
                    resultant_matrix[i][j] = 0;
                    for (uint32_t k=0; k<cols_A; k++) resultant_matrix[i][j] += input_matrix_A[i][k] * input_matrix_B[k][j];
                }
            }
        } else if (operation == MM_MULTIPLY_OMP) {
            uint32_t a, b, c;
            #pragma omp parallel for private(a, b, c) shared(input_matrix_A, input_matrix_B, resultant_matrix)
            for (a=0; a<rows_A; a++) {
                for (b=0; b<cols_B; b++) {
                resultant_matrix[a][b] = 0;
                    for (c=0; c<cols_A; c++) resultant_matrix[a][b] += input_matrix_A[a][c] * input_matrix_B[c][b];
                }
             }
        } else if (operation == MM_MULTIPLY_GEMM) {
            float* temp_matrix_A = malloc(rows_A * cols_A * sizeof(float));
            float* temp_matrix_B = malloc(rows_B * cols_B * sizeof(float));
            float* temp_resultant_matrix = malloc(rows_A * cols_B * sizeof(float));

            for (uint32_t i = 0; i < rows_A; i++) {
                for (uint32_t j = 0; j < cols_A; j++) {
                    temp_matrix_A[i * cols_A + j] = input_matrix_A[i][j];
                }
            }

            for (uint32_t i = 0; i < rows_B; i++) {
                for (uint32_t j = 0; j < cols_B; j++) {
                    temp_matrix_B[i * cols_B + j] = input_matrix_B[i][j];
                }
            }

            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, rows_A, cols_B, cols_A, 1.0, temp_matrix_A, cols_A, temp_matrix_B, cols_B, 0.0, temp_resultant_matrix, cols_B);

            for (uint32_t i = 0; i < rows_A; i++) {
                for (uint32_t j = 0; j < cols_B; j++) {
                    resultant_matrix[i][j] = temp_resultant_matrix[i * cols_B + j];
                }
            }

            free(temp_matrix_A);
            free(temp_matrix_B);
            free(temp_resultant_matrix);
        }
}

uint8_t sign(float num) {
    return (num > 0) - (num < 0);
}

void vector_normalize(float* vector, uint32_t size) {
    float norm = 0.0;
    for (uint32_t i=0; i<size; i++) norm += vector[i] * vector[i];
    norm = sqrt(norm);
    for (uint32_t i=0; i<size; i++) vector[i] = (norm == 0) ? 0 : vector[i]/norm;
}

float vector_norm(float* vector, uint32_t size) {
    float norm = 0.0;
    for (uint32_t i=0; i<size; i++) norm += vector[i] * vector[i];

    return sqrt(norm);
}

float matrix_norm(float** input_matrix_A1, float** input_matrix_A, uint32_t rows, uint32_t cols) {
    float norm = 0.0;
    for (uint32_t i=0; i<rows; i++) {
        for (uint32_t j=0; j<cols; j++) norm += fabs(input_matrix_A1[i][j] - input_matrix_A[i][j]) * fabs(input_matrix_A1[i][j] - input_matrix_A[i][j]);
    }

    return sqrt(norm);
}

float cosine_similarity(float* curr_vector, float* new_vector, uint32_t size) {
    float dot_product = 0.0f;
    for (uint32_t i=0; i<size; i++) dot_product += curr_vector[i] * new_vector[i];

    return (vector_norm(curr_vector, size) == 0 || vector_norm(new_vector, size) == 0) ? 1e-6 : dot_product / (vector_norm(curr_vector, size) * vector_norm(new_vector, size));
}

void matrix_vector_multiply(float** input_matrix, uint32_t rows, uint32_t cols, float* input_vector, float* resultant_vector, matrix_vector_operation_e operation) {
    if (operation == MV_MULTIPLY) {
        for (uint32_t i=0; i<rows; i++) {
            resultant_vector[i] = 0;
            for (uint32_t j=0; j<cols; j++) {
                resultant_vector[i] += input_matrix[i][j] * input_vector[j];
            }
        }
    } else if (operation == MV_MULTIPLY_GEMV) {
        float* temp_matrix = malloc(rows * cols * sizeof(float));

        for (uint32_t i = 0; i < rows; i++) {
            for (uint32_t j = 0; j < cols; j++) {
                temp_matrix[i * cols + j] = input_matrix[i][j];
            }
        }

        cblas_sgemv(CblasRowMajor, CblasNoTrans, rows, cols, 1.0f, temp_matrix, cols, input_vector, 1, 0.0f, resultant_vector, 1);

        free(temp_matrix);
    }
}