#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdint.h>
#include "neuralNetwork.h"

#define BUFFER_SIZE 200
#define FILE_HEADER_STRING "__info2_neural_network_file_format__"

static void softmax(Matrix *matrix)
{
    if(matrix->cols > 0)
    {
        double *colSums = (double *)calloc((size_t)matrix->cols, sizeof(double));
        if(colSums == NULL) return;

        for(int colIdx = 0; colIdx < matrix->cols; colIdx++)
        {
            for(int rowIdx = 0; rowIdx < matrix->rows; rowIdx++)
            {
                MatrixType expValue = (MatrixType)exp(getMatrixAt(*matrix, rowIdx, colIdx));
                setMatrixAt(expValue, *matrix, rowIdx, colIdx);
                colSums[colIdx] += expValue;
            }
        }

        for(int colIdx = 0; colIdx < matrix->cols; colIdx++)
        {
            double s = colSums[colIdx];
            if(s == 0.0) s = 1.0;
            for(int rowIdx = 0; rowIdx < matrix->rows; rowIdx++)
            {
                MatrixType normalizedValue = (MatrixType)(getMatrixAt(*matrix, rowIdx, colIdx) / s);
                setMatrixAt(normalizedValue, *matrix, rowIdx, colIdx);
            }
        }

        free(colSums);
    }
}

static void relu(Matrix *matrix)
{
    for(int i = 0; i < matrix->rows * matrix->cols; i++)
    {
        matrix->buffer[i] = matrix->buffer[i] >= 0 ? matrix->buffer[i] : 0;
    }
}

/* Prüft den Dateikopf. Liefert 1 bei Erfolg, 0 bei Fehler. */
static int checkFileHeader(FILE *file)
{
    if(file == NULL) return 0;

    size_t headerLen = strlen(FILE_HEADER_STRING);
    if(headerLen == 0 || headerLen >= BUFFER_SIZE) return 0;

    char buffer[BUFFER_SIZE];
    if(fseek(file, 0, SEEK_SET) != 0) return 0;
    if(fread(buffer, sizeof(char), headerLen, file) != headerLen) return 0;
    if(memcmp(buffer, FILE_HEADER_STRING, headerLen) != 0) return 0;

    return 1;
}

/* Liest eine Dimension (wie vom Test-Writer mit sizeof(int) geschrieben) und prüft Plausibilität.
   Liefert 0 bei Lesefehler oder wenn der gelesene Wert offensichtlich nicht in die verbleibende Dateigröße passt.
*/
static unsigned int readDimension(FILE *file)
{
    if (file == NULL) return 0;

    int value = 0;
    if (fread(&value, sizeof(int), 1, file) != 1) {
        return 0;
    }

    if (value < 0) return 0;
    return (unsigned int)value;
}


/* Liest eine Matrix rows x cols; wenn Einlese-Fehler: leere Matrix zurückgeben */
static Matrix readMatrix(FILE *file, unsigned int rows, unsigned int cols)
{
    Matrix matrix = createMatrix(rows, cols);
    if(matrix.buffer == NULL) return matrix;

    size_t toRead = (size_t)rows * cols;
    if(toRead > 0)
    {
        size_t readCount = fread(matrix.buffer, sizeof(MatrixType), toRead, file);
        if(readCount != toRead)
        {
            clearMatrix(&matrix);
        }
    }

    return matrix;
}

static Layer readLayer(FILE *file, unsigned int inputDimension, unsigned int outputDimension)
{
    Layer layer;
    layer.activation = NULL;
    layer.weights = readMatrix(file, outputDimension, inputDimension);
    layer.biases  = readMatrix(file, outputDimension, 1);
    return layer;
}

static int isEmptyLayer(const Layer layer)
{
    return layer.biases.cols == 0 || layer.biases.rows == 0 || layer.biases.buffer == NULL ||
           layer.weights.rows == 0 || layer.weights.cols == 0 || layer.weights.buffer == NULL;
}

static void clearLayer(Layer *layer)
{
    if(layer == NULL) return;
    clearMatrix(&layer->weights);
    clearMatrix(&layer->biases);
    layer->activation = NULL;
}

static void assignActivations(NeuralNetwork model)
{
    if(model.numberOfLayers == 0) return;
    for(int i = 0; i < (int)model.numberOfLayers - 1; i++)
        model.layers[i].activation = relu;
    model.layers[model.numberOfLayers - 1].activation = softmax;
}

NeuralNetwork loadModel(const char *path)
{
    NeuralNetwork model = {NULL, 0};
    FILE *file = fopen(path, "rb");
    if(file == NULL) return model;

    if(fseek(file, 0, SEEK_SET) != 0)
    {
        fclose(file);
        return model;
    }

    if(!checkFileHeader(file))
    {
        fclose(file);
        return model;
    }

    /* Lese erstes Dimensions-Paar (input, output) */
    unsigned int inputDimension = readDimension(file);
    unsigned int outputDimension = readDimension(file);

    fprintf(stderr, "[loadModel] first dims: input=%u output=%u\n", inputDimension, outputDimension);

    while (inputDimension > 0 && outputDimension > 0)
    {
        Layer layer = readLayer(file, inputDimension, outputDimension);

        if (isEmptyLayer(layer))
        {
            clearLayer(&layer);
            clearModel(&model);
            break;
        }

        Layer *tmp = (Layer *)realloc(model.layers, (model.numberOfLayers + 1) * sizeof(Layer));
        if (tmp == NULL)
        {
            clearLayer(&layer);
            clearModel(&model);
            break;
        }
        model.layers = tmp;
        model.layers[model.numberOfLayers] = layer;
        model.numberOfLayers++;

        fprintf(stderr, "[loadModel] loaded layer %d: weights %u x %u, biases %u x %u\n",
                model.numberOfLayers,
                layer.weights.rows, layer.weights.cols,
                layer.biases.rows, layer.biases.cols);

        /* Lese das nächste Dimensions-Paar (writer schreibt für jede Schicht ein Paar) */
        unsigned int nextInput = readDimension(file);
        unsigned int nextOutput = readDimension(file);

        fprintf(stderr, "[loadModel] next raw dims read: nextInput=%u nextOutput=%u\n", nextInput, nextOutput);

        /* Wenn nächstes Paar (0,0) -> Ende */
        if (nextInput == 0 || nextOutput == 0)
        {
            inputDimension = 0;
            outputDimension = 0;
            break;
        }

        /* Setze für die nächste Iteration */
        inputDimension = nextInput;
        outputDimension = nextOutput;

        fprintf(stderr, "[loadModel] next dims: input=%u output=%u\n", inputDimension, outputDimension);
    }

    fclose(file);
    assignActivations(model);
    return model;
}

static Matrix imageBatchToMatrixOfImageVectors(const GrayScaleImage images[], unsigned int count)
{
    Matrix matrix;
    matrix.buffer = NULL;
    matrix.rows = 0;
    matrix.cols = 0;

    if(count > 0 && images != NULL)
    {
        matrix = createMatrix(images[0].height * images[0].width, count);
        if(matrix.buffer != NULL)
        {
            for(unsigned int i = 0; i < count; i++)
            {
                for(unsigned int j = 0; j < images[i].width * images[i].height; j++)
                    setMatrixAt((MatrixType)images[i].buffer[j], matrix, j, i);
            }
        }
    }

    return matrix;
}

static Matrix forward(const NeuralNetwork model, Matrix inputBatch)
{
    Matrix result = inputBatch;
    if(result.buffer == NULL) return result;

    for(int i = 0; i < model.numberOfLayers; i++)
    {
        Matrix weightResult = multiply(model.layers[i].weights, result);
        clearMatrix(&result);

        Matrix biasResult = add(weightResult, model.layers[i].biases);
        clearMatrix(&weightResult);

        if(model.layers[i].activation != NULL)
            model.layers[i].activation(&biasResult);

        result = biasResult;
    }

    return result;
}

unsigned char *argmax(const Matrix matrix)
{
    if(matrix.rows == 0 || matrix.cols == 0) return NULL;

    unsigned char *maxIdx = (unsigned char *)malloc((size_t)matrix.cols * sizeof(unsigned char));
    if(maxIdx == NULL) return NULL;

    for(int colIdx = 0; colIdx < matrix.cols; colIdx++)
    {
        int best = 0;
        for(int rowIdx = 1; rowIdx < matrix.rows; rowIdx++)
        {
            if(getMatrixAt(matrix, rowIdx, colIdx) > getMatrixAt(matrix, best, colIdx))
                best = rowIdx;
        }
        maxIdx[colIdx] = (unsigned char)best;
    }
    return maxIdx;
}

unsigned char *predict(const NeuralNetwork model, const GrayScaleImage images[], unsigned int numberOfImages)
{
    Matrix inputBatch = imageBatchToMatrixOfImageVectors(images, numberOfImages);
    Matrix outputBatch = forward(model, inputBatch);
    unsigned char *result = argmax(outputBatch);
    clearMatrix(&outputBatch);
    return result;
}

void clearModel(NeuralNetwork *model)
{
    if(model == NULL) return;
    for(int i = 0; i < model->numberOfLayers; i++)
        clearLayer(&model->layers[i]);
    free(model->layers);
    model->layers = NULL;
    model->numberOfLayers = 0;
}