#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <string.h>

char *file_pesi = "rete_mnist.bin";

#include "mnist/mnist_manager.h"
// #include "cifar-10/cifar10_manager.h";

typedef unsigned char byte;

float LRE = 0.01;
float soglia_sigmoide = 0.5;
#define TOLLERANZA 99.5

// Metto 0 per sigmoide, 1 per ReLU, 2 per gradino
#define FUNZIONE_ATTIVAZIONE 1

typedef struct
{
    float *pesi;
    float bias;
    int size;
} Percettrone;

typedef struct
{
    Percettrone *percettroni;
    int size;
} Layer;

typedef struct
{
    Layer *layers;
    int size;
} ReteNeurale;

// Inizializzazioni
Percettrone inzializza_percettrone(int);
Layer inizializza_layer(int, int);
ReteNeurale inizializza_rete_neurale(int, int, int, int);

// Previsioni
float attivazione(Percettrone p, float *);
float derivata_attivazione(float);
void softmax(float *, int);
int previsione_softmax(float *, int);

// Forward
float **elabora_attivazioni(ReteNeurale *, Istanza);

// Discesa del gradiente stocastico
float **elabora_gradienti(ReteNeurale *, byte, float **);
void discesa_gradiente(ReteNeurale *, float **, float **);
float calcola_gradiente_disceso(ReteNeurale *, int, int, float **);

// Correzioni
void aggiorna_pesi(ReteNeurale *, float **, float **, Istanza);
void correggi_pesi_percettrone_float(Percettrone *, int, float **, float);
void correggi_pesi_percettrone_byte(Percettrone *, Istanza, float, int);

// Addestramento
byte addestra(ReteNeurale *, Dataset);

// Import/Export
void salvaReteNeurale(const char *, ReteNeurale *);
ReteNeurale *caricaReteNeurale(const char *);

// Stampa
void stampa_pesi_rete(ReteNeurale *);

/*############# INIZIALIZZAZIONI #########################*/

// INIZIALIZZAZIONE HE
Percettrone inizializza_percettrone(int n_pesi)
{
    Percettrone p;
    p.pesi = (float *)malloc(sizeof(float) * n_pesi);

    // Deviazione standard
    float he_std = sqrt(2.0 / n_pesi);

    for (int i = 0; i < n_pesi; i++)
    {
        // Teorema centrale del limite
        float somma = 0.0;
        for (int i = 0; i < 12; i++)
        {
            somma += (float)rand() / RAND_MAX;
        }
        p.pesi[i] = (somma - 6.0) * he_std;

        p.bias = 0.0; // Bias a zero per inizializzazione He
    }

    p.size = n_pesi;

    return p;
}

Layer inizializza_layer(int n_percettroni, int n_pesi)
{
    Layer layer;
    layer.percettroni = (Percettrone *)malloc(sizeof(Percettrone) * n_percettroni);

    for (int i = 0; i < n_percettroni; i++)
    {
        layer.percettroni[i] = inizializza_percettrone(n_pesi);
    }

    layer.size = n_percettroni;

    return layer;
}

ReteNeurale inizializza_rete_neurale(int numero_input, int numero_layers, int numero_percettroni_iniziali, int numero_percettroni_finali)
{
    srand(time(NULL));
    ReteNeurale r;
    r.layers = (Layer *)malloc(sizeof(Layer) * numero_layers);
    r.size = numero_layers;
    int somma_parametri = 0;

    for (int livello = 0; livello < numero_layers; livello++)
    {
        int numero_percettroni_livello = 1;

        if (livello == numero_layers - 1)
            numero_percettroni_livello = numero_percettroni_finali;
        else
        {
            float frazione = (float)(numero_layers - livello) / (float)numero_layers;
            numero_percettroni_livello = (int)((float)(numero_percettroni_iniziali * frazione));
        }

        printf("Layer %d -> percettroni: %d\n", livello, numero_percettroni_livello);

        if (livello == 0)
        {
            r.layers[livello] = inizializza_layer(numero_percettroni_livello, numero_input);
            somma_parametri += (numero_percettroni_livello * numero_input);
        }
        else
        {
            r.layers[livello] = inizializza_layer(numero_percettroni_livello, r.layers[livello - 1].size);
            somma_parametri += (numero_percettroni_livello * r.layers[livello - 1].size);
        }
    }

    printf("Rete neurale da %d parametri\n", somma_parametri);

    return r;
}

/*
    ################# PREVISIONI ################################
*/

float attivazione(Percettrone p, float *valori)
{
    float sommatoria = 0.0;
    for (int i = 0; i < p.size; i++)
    {
        sommatoria += (valori[i] * p.pesi[i]);
    }

    sommatoria += p.bias;
    float risultato;

    if (FUNZIONE_ATTIVAZIONE == 0)
        risultato = 1.0 / (1.0 + exp(-sommatoria));
    else if (FUNZIONE_ATTIVAZIONE == 1)
        risultato = sommatoria > 0 ? sommatoria : 0.0;
    else
        risultato = sommatoria > 0 ? 1.0 : 0.0;

    return risultato;
}

float derivata_attivazione(float valore)
{
    float derivata;

    if (FUNZIONE_ATTIVAZIONE == 0)
        derivata = valore * (1.0 - valore);
    else if (FUNZIONE_ATTIVAZIONE == 1)
        derivata = valore > 0 ? 1.0 : 0.0;
    else
        derivata = 0.0;

    return derivata;
}

void softmax(float *input, int size)
{
    float max = input[0];

    for (int i = 1; i < size; i++)
        if (input[i] > max)
            max = input[i];

    float somma = 0.0;
    for (int i = 0; i < size; i++)
    {
        input[i] = expf(input[i] - max);
        somma += input[i];
    }

    for (int i = 0; i < size; i++)
        input[i] /= somma;
}

int previsione_softmax(float *livello_percettroni, int size)
{

    softmax(livello_percettroni, size);
    int max = 0;

    for (int i = 1; i < size; i++) {
        if (livello_percettroni[i] > livello_percettroni[max])
            max = i;
    }

    return max;
}

/*
    ################# FORWARD ################################
*/

float **elabora_attivazioni(ReteNeurale *rete, Istanza istanza)
{
    float **attivazioni = (float **)malloc(sizeof(float *) * rete->size);
    float *inputs = (float *)malloc(sizeof(float) * N_INPUTS);
    
    for (int i = 0; i < N_INPUTS; i++) {
        inputs[i] = (float)istanza.dati[i] / 255.0;
    }

    for (int indice_layer = 0; indice_layer < rete->size; indice_layer++)
    {
        attivazioni[indice_layer] = (float *)malloc(sizeof(float) * rete->layers[indice_layer].size);
        
        for (int indice_percettrone = 0; indice_percettrone < rete->layers[indice_layer].size; indice_percettrone++)
        {
            if(indice_layer == 0)
                attivazioni[indice_layer][indice_percettrone] = attivazione(rete->layers[indice_layer].percettroni[indice_percettrone], inputs);
            else
                attivazioni[indice_layer][indice_percettrone] = attivazione(rete->layers[indice_layer].percettroni[indice_percettrone], attivazioni[indice_layer - 1]);
        }
    }

    free(inputs);
    return attivazioni;
}

/*
    ################# DISCESA DEL GRADIENTE STOCASTICO ################################
*/

float **elabora_gradienti(ReteNeurale *rete_neurale, byte output, float **attivazioni)
{
    float **gradienti = (float **)malloc(sizeof(float *) * rete_neurale->size);

    for (int indice_layer = 0; indice_layer < rete_neurale->size; indice_layer++)
    {
        gradienti[indice_layer] = (float *)malloc(sizeof(float) * rete_neurale->layers[indice_layer].size);
    }

    // ULTIMO LAYER
    for (int indice_percettrone = 0; indice_percettrone < rete_neurale->layers[rete_neurale->size - 1].size; indice_percettrone++)
        if (indice_percettrone == output)
            gradienti[rete_neurale->size - 1][indice_percettrone] = 1 - attivazioni[rete_neurale->size - 1][indice_percettrone];
        else
            gradienti[rete_neurale->size - 1][indice_percettrone] = 0 - attivazioni[rete_neurale->size - 1][indice_percettrone];

    //RESTANTI LAYER
    discesa_gradiente(rete_neurale, attivazioni, gradienti);

    return gradienti;
}

void discesa_gradiente(ReteNeurale *rete, float **attivazioni, float **gradienti)
{
    for (int indice_layer = rete->size - 2; indice_layer >= 0; indice_layer--)
    {
        for (int indice_percettrone = 0; indice_percettrone < rete->layers[indice_layer].size; indice_percettrone++)
        {
            float derivata = derivata_attivazione(attivazioni[indice_layer][indice_percettrone]);
            float gradiente_disceso = calcola_gradiente_disceso(rete, indice_layer + 1, indice_percettrone, gradienti);

            gradienti[indice_layer][indice_percettrone] = gradiente_disceso * derivata;
        }
    }
}

float calcola_gradiente_disceso(ReteNeurale *rete, int livello, int indice_peso, float **gradienti)
{
    float sommatoria = 0.0;
    for (int indice_percettrone = 0; indice_percettrone < rete->layers[livello].size; indice_percettrone++)
    {
        sommatoria += (gradienti[livello][indice_percettrone] * rete->layers[livello].percettroni[indice_percettrone].pesi[indice_peso]);
    }

    return sommatoria;
}

/*
    ################# CORREZIONI ################################
*/

void aggiorna_pesi(ReteNeurale *rete_neurale, float **attivazioni, float **gradienti, Istanza istanza)
{

    for (int indice_percettrone = 0; indice_percettrone < rete_neurale->layers[rete_neurale->size - 1].size; indice_percettrone++)
    {

        for (int indice_peso = 0; indice_peso < rete_neurale->layers[rete_neurale->size - 1].percettroni[indice_percettrone].size; indice_peso++)
        {
            float gradiente_peso = gradienti[rete_neurale->size - 1][indice_percettrone] * attivazioni[rete_neurale->size - 2][indice_peso];
            rete_neurale->layers[rete_neurale->size - 1].percettroni[indice_percettrone].pesi[indice_peso] += gradiente_peso * LRE;
        }
        rete_neurale->layers[rete_neurale->size - 1].percettroni[indice_percettrone].bias += gradienti[rete_neurale->size - 1][indice_percettrone] * LRE;
    }

    for (int indice_layer = rete_neurale->size - 2; indice_layer >= 0; indice_layer--)
    {
        for (int indice_percettrone = 0; indice_percettrone < rete_neurale->layers[indice_layer].size; indice_percettrone++)
        {
            if (indice_layer != 0)
            {
                correggi_pesi_percettrone_float(&rete_neurale->layers[indice_layer].percettroni[indice_percettrone], indice_layer, attivazioni, gradienti[indice_layer][indice_percettrone]);
            }
            else
            {
                correggi_pesi_percettrone_byte(&rete_neurale->layers[0].percettroni[indice_percettrone], istanza, gradienti[0][indice_percettrone], indice_percettrone);
            }
        }
    }
}

void correggi_pesi_percettrone_float(Percettrone *p, int layer, float **input, float gradiente_percettrone)
{

    for (int indice_peso = 0; indice_peso < p->size; indice_peso++)
    {
        float gradiente_peso = gradiente_percettrone * input[layer - 1][indice_peso];
        p->pesi[indice_peso] += (gradiente_peso * LRE);
    }

    p->bias += (gradiente_percettrone * LRE);
}

void correggi_pesi_percettrone_byte(Percettrone *p, Istanza input, float gradiente_percettrone, int indice_percettrone)
{
    for (int indice_peso = 0; indice_peso < p->size; indice_peso++)
    {
        float gradiente_peso = gradiente_percettrone * ((float)input.dati[indice_peso] / 255.0);
        p->pesi[indice_peso] += (gradiente_peso * LRE);
    }

    p->bias += (gradiente_percettrone * LRE);
}

/*
    ################### ADDESTRAMENTO #########################
*/

byte addestra(ReteNeurale *rete_neurale, Dataset set)
{
    int corrette = 0;
    for (int indice_set = 0; indice_set < set.size; indice_set++)
    {
        byte output_corretto = set.istanze[indice_set].classificazione;

        float **attivazioni = elabora_attivazioni(rete_neurale, set.istanze[indice_set]);

        float **gradienti = elabora_gradienti(rete_neurale, output_corretto, attivazioni);
        int previsto = previsione_softmax(attivazioni[rete_neurale->size - 1], rete_neurale->layers[rete_neurale->size - 1].size);

        if (previsto == output_corretto) {
            corrette++;
        }
        
        aggiorna_pesi(rete_neurale, attivazioni, gradienti, set.istanze[indice_set]);

        for (int i = 0; i < rete_neurale->size; i++)
            free(gradienti[i]);
        free(gradienti);

        for (int i = 0; i < rete_neurale->size; i++)
            free(attivazioni[i]);
        free(attivazioni);
        //free(attivazioni_softmax);
    }

    float percentuale = ((corrette * 100.0) / set.size);

    printf("Risposte corrette: %.2f%%\n", percentuale);

    if (percentuale >= TOLLERANZA)
        return 1;
    else
        return 0;
}

/*
    ################# IMPORT EXPORT ################################
*/
void salvaReteNeurale(const char *filename, ReteNeurale *rete)
{
    FILE *file = fopen(filename, "wb");
    if (!file)
    {
        perror("Errore nell'apertura del file");
        exit(EXIT_FAILURE);
    }

    fwrite(&rete->size, sizeof(int), 1, file);

    for (int i = 0; i < rete->size; i++)
    {
        Layer *layer = &rete->layers[i];
        fwrite(&layer->size, sizeof(int), 1, file);

        for (int j = 0; j < layer->size; j++)
        {
            Percettrone *perc = &layer->percettroni[j];
            fwrite(&perc->size, sizeof(int), 1, file);
            fwrite(perc->pesi, sizeof(float), perc->size, file);
            fwrite(&perc->bias, sizeof(float), 1, file);
        }
    }

    fclose(file);
}

ReteNeurale *caricaReteNeurale(const char *filename)
{
    FILE *file = fopen(filename, "rb");

    if (!file)
    {
        perror("Errore nell'apertura del file");
        return NULL;
    }

    ReteNeurale *rete = malloc(sizeof(ReteNeurale));
    if (!rete)
    {
        perror("Errore nell'allocazione della memoria");
        return NULL;
    }

    fread(&rete->size, sizeof(int), 1, file);
    rete->layers = malloc(rete->size * sizeof(Layer));
    if (!rete->layers)
    {
        perror("Errore nell'allocazione della memoria");
        return NULL;
    }

    for (int i = 0; i < rete->size; i++)
    {
        Layer *layer = &rete->layers[i];
        fread(&layer->size, sizeof(int), 1, file);
        layer->percettroni = malloc(layer->size * sizeof(Percettrone));
        if (!layer->percettroni)
        {
            perror("Errore nell'allocazione della memoria");
            return NULL;
        }

        for (int j = 0; j < layer->size; j++)
        {
            Percettrone *perc = &layer->percettroni[j];
            fread(&perc->size, sizeof(int), 1, file);
            perc->pesi = malloc(perc->size * sizeof(float));
            if (!perc->pesi)
            {
                perror("Errore nell'allocazione della memoria");
                return NULL;
            }
            fread(perc->pesi, sizeof(float), perc->size, file);
            fread(&perc->bias, sizeof(float), 1, file);
        }
    }

    fclose(file);
    return rete;
}

/*
    ################ STAMPE ############################
*/
void stampa_pesi_rete(ReteNeurale *rete)
{
    for (int indice_layer = 0; indice_layer < rete->size; indice_layer++)
    {
        printf("\nLivello %d", indice_layer);
        for (int indice_percettrone = 0; indice_percettrone < rete->layers[indice_layer].size; indice_percettrone++)
        {
            printf("\n\tPercettrone %d", indice_percettrone);
            for (int indice_peso = 0; indice_peso < rete->layers[indice_layer].percettroni[indice_percettrone].size; indice_peso++)
            {
                printf("\n\t\tPeso %d: %f", indice_peso, rete->layers[indice_layer].percettroni[indice_percettrone].pesi[indice_peso]);
            }
        }
    }
    printf("\n");
}