#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "unity.h"
#include "neuralNetwork.h"

#define FILE_HEADER_STRING "__info2_neural_network_file_format__"


// --------------------------
// Hilfsfunktion zum Erstellen einer Test-Datei für das Netzwerk
// --------------------------
static void prepareNeuralNetworkFile(const char *path, const NeuralNetwork nn)
{
    FILE *file = fopen(path, "wb");
    if(file == NULL) return;

    // 1. Schreibe den Datei-Header
    //    Dieser Header wird beim Laden überprüft, um das Dateiformat sicherzustellen
    fwrite(FILE_HEADER_STRING, sizeof(char), strlen(FILE_HEADER_STRING), file);

    // 2. Schreibe alle Layer des Netzwerks
    for(unsigned int i = 0; i < nn.numberOfLayers; i++)
    {
        Layer layer = nn.layers[i];
        int inputDim = (int)layer.weights.cols;
        int outputDim = (int)layer.weights.rows;

        // Schreibe die Dimensionen des Layers
        fwrite(&inputDim, sizeof(int), 1, file);
        fwrite(&outputDim, sizeof(int), 1, file);

        // Schreibe die Gewichtsmatrix (outputDim x inputDim)
        fwrite(layer.weights.buffer, sizeof(MatrixType), outputDim * inputDim, file);

        // Schreibe den Bias-Vektor (outputDim x 1)
        fwrite(layer.biases.buffer, sizeof(MatrixType), outputDim, file);
    }

    // 3. Schreibe zwei Nullen, um das Ende der Layer anzuzeigen
    int zero = 0;
    fwrite(&zero, sizeof(int), 1, file); // inputDim = 0
    fwrite(&zero, sizeof(int), 1, file); // outputDim = 0

    fclose(file);
}



// --------------------------
// Test: Prüft, ob loadModel richtige Anzahl Layer lädt
// --------------------------
void test_loadModelReturnsCorrectNumberOfLayers(void)
{
    const char *path = "test_nn_file.info2";
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    prepareNeuralNetworkFile(path, nn);

    NeuralNetwork loaded = loadModel(path);
    TEST_ASSERT_EQUAL_INT(1, loaded.numberOfLayers);
    clearModel(&loaded);
    remove(path);
}

// --------------------------
// Test: Prüft Dimensionen der Gewichte
// --------------------------
void test_loadModelReturnsCorrectWeightDimensions(void)
{
    const char *path = "test_nn_file.info2";
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    prepareNeuralNetworkFile(path, nn);

    NeuralNetwork loaded = loadModel(path);
    TEST_ASSERT_EQUAL_INT(3, loaded.layers[0].weights.rows);
    TEST_ASSERT_EQUAL_INT(2, loaded.layers[0].weights.cols);
    clearModel(&loaded);
    remove(path);
}

// --------------------------
// Test: Prüft Dimensionen der Biases
// --------------------------
void test_loadModelReturnsCorrectBiasDimensions(void)
{
    const char *path = "test_nn_file.info2";
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    prepareNeuralNetworkFile(path, nn);

    NeuralNetwork loaded = loadModel(path);
    TEST_ASSERT_EQUAL_INT(3, loaded.layers[0].biases.rows);
    TEST_ASSERT_EQUAL_INT(1, loaded.layers[0].biases.cols);
    clearModel(&loaded);
    remove(path);
}

// --------------------------
// Test: Prüft, dass Gewichte korrekt geladen werden
// --------------------------
void test_loadModelReturnsCorrectWeights(void)
{
    const char *path = "test_nn_file.info2";
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    prepareNeuralNetworkFile(path, nn);

    NeuralNetwork loaded = loadModel(path);
    int n = loaded.layers[0].weights.rows * loaded.layers[0].weights.cols;
    TEST_ASSERT_EQUAL_INT_ARRAY(wBuf, loaded.layers[0].weights.buffer, n);
    clearModel(&loaded);
    remove(path);
}

// --------------------------
// Test: Prüft, dass Bias korrekt geladen werden
// --------------------------
void test_loadModelReturnsCorrectBiases(void)
{
    const char *path = "test_nn_file.info2";
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    prepareNeuralNetworkFile(path, nn);

    NeuralNetwork loaded = loadModel(path);
    int n = loaded.layers[0].biases.rows * loaded.layers[0].biases.cols;
    TEST_ASSERT_EQUAL_INT_ARRAY(bBuf, loaded.layers[0].biases.buffer, n);
    clearModel(&loaded);
    remove(path);
}

// --------------------------
// Test: predict Funktion
// --------------------------
void test_predictReturnsCorrectLabels(void)
{
    GrayScalePixelType img1[] = {10,20,30,40};
    GrayScalePixelType img2[] = {5,15,25,35};
    GrayScaleImage images[] = {
        {.buffer=img1, .width=2, .height=2},
        {.buffer=img2, .width=2, .height=2}
    };

    // Dummy Network für test: ReLU-ähnlich
    MatrixType w1[] = {1,0,0,1,1,0,0,1};
    MatrixType b1[] = {0,0};
    Layer layers[] = {{.weights={w1,2,4}, .biases={b1,2,1}, .activation=NULL}};
    NeuralNetwork nn = {layers,1};

    unsigned char *labels = predict(nn, images, 2);
    TEST_ASSERT_NOT_NULL(labels);
    free(labels);
}

// --------------------------
// Test: clearModel setzt Pointer auf NULL
// --------------------------
void test_clearModelSetsMembersToNull(void)
{
    MatrixType wBuf[] = {1,2,3,4,5,6};
    MatrixType bBuf[] = {1,2,3};
    Layer layers[] = {{.weights={wBuf,3,2}, .biases={bBuf,3,1}}};
    NeuralNetwork nn = {layers,1};

    clearModel(&nn);
    TEST_ASSERT_NULL(nn.layers);
    TEST_ASSERT_EQUAL_INT(0, nn.numberOfLayers);
}

// --------------------------
// Test: Fehlerhafte Datei (Header falsch)
// --------------------------
void test_loadModelFailsOnWrongFileTag(void)
{
    const char *path = "wrong_nn_file.info2";
    FILE *file = fopen(path, "wb");
    if(file != NULL)
    {
        const char *wrongTag = "wrong_header_string";
        fwrite(wrongTag, sizeof(char), strlen(wrongTag), file);
        fclose(file);
    }

    NeuralNetwork nn = loadModel(path);
    TEST_ASSERT_NULL(nn.layers);
    TEST_ASSERT_EQUAL_INT(0, nn.numberOfLayers);
    remove(path);
}

// --------------------------
// Unity Setup / Teardown
// --------------------------
void setUp(void) {}
void tearDown(void) {}

// --------------------------
// Hauptfunktion zum Ausführen der Tests
// --------------------------
int main(void)
{
    UNITY_BEGIN();

    printf("\n============================\nNeural network tests\n============================\n");
    RUN_TEST(test_loadModelReturnsCorrectNumberOfLayers);
    RUN_TEST(test_loadModelReturnsCorrectWeightDimensions);
    RUN_TEST(test_loadModelReturnsCorrectBiasDimensions);
    RUN_TEST(test_loadModelReturnsCorrectWeights);
    RUN_TEST(test_loadModelReturnsCorrectBiases);
    RUN_TEST(test_predictReturnsCorrectLabels);
    RUN_TEST(test_clearModelSetsMembersToNull);
    RUN_TEST(test_loadModelFailsOnWrongFileTag);

    return UNITY_END();
}