#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <complex.h>
#include "../files/files.h"
#include <string.h>

#define A 1 

struct qam_system_s {
    int M; // Nombre de symboles M-QAM
    int k; // Nombre de bits/symboles
    double Fs; // Fréquence d'échantillionage
    double Ts; // Temps d'échantillionage
    int N; // Nombre d'échantillions
    double Fc; // Fréquence de la porteuse
    double complex** constellation; // Tableau de symboles I + j Q
};
typedef struct qam_system_s qam_system;

// Initialisation de la constellation (double tableau de taille sqrt(M)), 
// ToDo : changer à un tableau à 1 dimension pour éviter de calculer sqrt(M)
void init_constellation (qam_system* qam) {
    int sm = (int)sqrt(qam->M);
    qam->constellation = (double complex**)malloc(sizeof(double complex*) * sm);
    
    for (int i = 0; i < sm; i++) {
       qam->constellation[i] = (double complex*)malloc(sizeof(double complex) * sm);
    }

    double norm_factor = sqrt((double)(qam->M - 1) / 3.0); // Pour puissance unitaire

    for (int i = 0; i < sm; i++) {
        double complex ip = -(sm - 1) + 2 * i;
        for (int j = 0; j < sm; j++) {
            double complex qp = -(sm - 1) + 2 * j;
            qam->constellation[i][j] = (ip + I * qp); 
        }
    }
}

// Calcul du bruit gaussien pour un sigma donné
// Formule de Box-Muller
double gaussian_noise (double sigma) {
    double u1 = (rand() + 1) / ((double)RAND_MAX + 2);
    double u2 = (rand() + 1) / ((double)RAND_MAX + 2);
    return sigma * sqrt(-2 * log(u1)) * cos(2 * M_PI * u2);
}

// Ajout du bruit
void add_noise (double complex* s, int len, double sigma) {
    for (int i = 0; i < len; i++) {
        double nr = gaussian_noise(sigma);
        double ni = gaussian_noise(sigma);
        s[i] += nr + I * ni;
    }
}

// Changer le tableau de bits en boolen ou alors la represenation binaire et shifter pour extraire les bits (pas bien si M plus grand)
void bits_to_symbols (qam_system* qam, uint8_t* bits, int nb_bits, double complex* symbols) {
    int nb_symbols = nb_bits / qam->k;
    int sm = sqrt(qam->M);
    for (int k = 0; k < nb_symbols; k++) {
        int id = 0;
        for (int b = 0 ; b < qam->k; b++) {
            id = id * 2 + bits[k * qam->k + b];
        }
        int i = id / sm; 
        int j = id % sm;
        symbols[k] = qam->constellation[i][j];
    }
}

// Modulation QAM 
void modulate (qam_system* qam, double complex* symbols, int nb_symbols, double complex* s) {
    for (int k = 0; k < nb_symbols; k++) {
        double complex iq = symbols[k];
        for (int n = 0; n < qam->N; n++) {
            s[k * qam->N + n] = iq * cexp(2 * I * M_PI * qam->Fc * ((double)n / qam->Fs)); 
        }
    }
}

// Demodulation QAM 
void demodulate(qam_system* qam, double complex* s, int nb_symbols, uint8_t* bits_hat) {
    for (int k = 0; k < nb_symbols; k++) {
        double complex r = 0;
        for (int n = 0; n < qam->N; n++) {
            r += s[k * qam->N + n] * cexp(-2 * I * M_PI * qam->Fc * ((double)n / qam->Fs));
        }
        r /= qam->N;

        // Distance euclidien de Ir et Qr pour avoir le point le plus proche de la constellation (lent)
        int sm = (int)sqrt(qam->M);
        double min_d = INFINITY;
        int i_cl, j_cl = 0;
        for (int i = 0; i < sm; i++) {
            for (int j = 0; j < sm; j++) {
                double d = cabs(r - qam->constellation[i][j]);
                if (d < min_d) {
                    min_d = d;
                    i_cl = i;
                    j_cl = j;
                }
            }
        }

        // index du symbole (id) : même mappage que dans bits_to_symbols()
        int id = i_cl * sm + j_cl;

        for (int b = 0; b < qam->k; b++) {
            bits_hat[k * qam->k + b] = (id >> (qam->k - 1 - b)) & 1;
        }
    }
}

// Demodulation QAM avec porteuse estimé
void demodulate2(qam_system* qam, double complex* s, int nb_symbols, uint8_t* bits_hat) {
    for (int k = 0; k < nb_symbols; k++) {
        double complex r = 0;
        for (int n = 0; n < qam->N; n++) {
            r += s[k * qam->N + n];
        }
        r /= qam->N;

        // Distance euclidien de Ir et Qr pour avoir le point le plus proche de la constellation (lent)
        int sm = (int)sqrt(qam->M);
        double min_d = INFINITY;
        int i_cl, j_cl = 0;
        for (int i = 0; i < sm; i++) {
            for (int j = 0; j < sm; j++) {
                double d = cabs(r - qam->constellation[i][j]);
                if (d < min_d) {
                    min_d = d;
                    i_cl = i;
                    j_cl = j;
                }
            }
        }

        // index du symbole (id) : même mappage que dans bits_to_symbols()
        int id = i_cl * sm + j_cl;

        for (int b = 0; b < qam->k; b++) {
            bits_hat[k * qam->k + b] = (id >> (qam->k - 1 - b)) & 1;
        }
    }
}
/*
// Pour j = 1 dans la def de wiki sur l'autocorrelation
double fc_autocorrelation_1_rad(double complex* s, int N) {
    double complex r = 0;
    for (int n = 0; n < N - 1; n++) {
        r += s[n] * conj(s[n + 1]);
    }
    r /= (N - 1);
    double om = -carg(r); // rad / sample
    return om;
}

// Autocorrelation pour trouver la Fc (fréquence de la porteuse)
double fc_autocorrelation_multilag_rad(double complex* s, int N, int max_lag) {
    double sum_phase = 0;
    for (int k = 1; k <= max_lag; k++) {
        double complex rk = 0;
        for (int n = 0; n < N - k; n++) {
            rk += s[n] * conj(s[n + k]);
        }
        rk /= (double)(N - k);
        sum_phase += -carg(rk) / k;
    }
    return sum_phase / max_lag;  // rad/sample
}

// Correction du signal pour enlever l'offset de la porteuse
void signal_correction(double complex* s, int total_samples, double om_hat) {
    for (int n = 0; n < total_samples; n++) {
        s[n] *= cexp(-I * om_hat * n);
    }
}
*/

void blind_carrier(double complex* s, int N, int M) {
    int k = (int)sqrt(M);  // puissance à élever
    double complex sum = 0;
    for(int n = 0; n < N; n++)
        sum += cpow(s[n], k);

    double phi_hat = carg(sum) / k;

    for(int n = 0; n < N; n++)
        s[n] *= cexp(-I * phi_hat);
}

// Libération de la mémoire
void free_constellation(qam_system* qam) {
    int sm = (int)sqrt(qam->M);
    for (int i = 0; i < sm; i++)
        free(qam->constellation[i]);
    free(qam->constellation);
}

// Compare deux tableaux de bits (0/1) et retourne le pourcentage de fiabilité.
double taux_erreur_bits(const uint8_t *in_bits, size_t nb_bits_in,
                        const uint8_t *out_bits, size_t nb_bits_out) {
    if (!in_bits || !out_bits)
        return 1.0; 

    size_t n_min = nb_bits_in < nb_bits_out ? nb_bits_in : nb_bits_out;
    size_t n_max = nb_bits_in > nb_bits_out ? nb_bits_in : nb_bits_out;
    size_t err = n_max - n_min;

    for (size_t i = 0; i < n_min; i++)
        err += ((in_bits[i] ^ out_bits[i]) & 1);

    return n_max ? (double)err / n_max : 0.0;
}


int main (int argc, char *argv[]) {
    if (argc < 2) {
        fprintf(stderr, "Utilisation: %s <fichier_entree>\n", argv[0]);
        return 1;
    }

    qam_system qam;
    qam.M = 16;
    qam.k = (int)log2((double)(qam.M));
    qam.Fs = 44100;
    qam.Ts = 0.3;
    qam.N = (int)qam.Fs * qam.Ts;
    qam.Fc = 2000;
    init_constellation(&qam);

    printf("Lecture du fichier...\n");
    // Lecture du fichier et conversion en bits
    const char *input_filename = argv[1];
    bit_array input_bits = file_to_bits(input_filename);
    size_t nb_symbols = input_bits.nb_bits / qam.k;

    printf("Mise en forme des symboles...\n");
    // Mise en forme des symboles
    double complex *symbols = malloc(sizeof(double complex) * nb_symbols);
    bits_to_symbols(&qam, input_bits.bits, input_bits.nb_bits, symbols);

    printf("Modulation...\n");
    // Modulation QAM 
    int total_samples = qam.N * nb_symbols;
    double complex* s = (double complex*)malloc(sizeof(double complex) * total_samples);
    modulate(&qam, symbols, nb_symbols, s);

    // Ajout du bruit 
    double signal_power = (2.0/3.0)*(qam.M-1); // puissance moyenne avant échelle
    double snr_dB = 10; // Signal to noise ratio
    double snr_lin = pow(10.0, snr_dB / 10.0);
    double sigma = sqrt(signal_power / snr_lin);
    printf("Ajout du bruit... \n       puissance du signal : %f\n       SNR db : %f\n       sigma : %f\n", signal_power, snr_dB, sigma);
    add_noise(s, total_samples, 0);


    // Demodulation QAM
    //printf("Demodulation...\n");
    //bit_array output_bits;
    //output_bits.nb_bits = input_bits.nb_bits;
    //output_bits.bits = (uint8_t*)malloc(output_bits.nb_bits * sizeof(uint8_t));
    //demodulate(&qam, s, nb_symbols, output_bits.bits);

    //printf("Ecriture...\n");
    // Ecriture du fichier de Demodulation
    //char *output_filename = make_output_filename(input_filename);
    //bits_to_file(output_filename, &output_bits);

    //double erreurs = taux_erreur_bits(input_bits.bits, input_bits.nb_bits, output_bits.bits, output_bits.nb_bits);

    //printf("Comparaison :\n");
    //printf("    Erreurs       : %f\n", erreurs * 100);

    /*
    double om_hat = fc_autocorrelation_1_rad(s, total_samples);
    printf("Estimation de fc 1: %f\n", (om_hat * qam.Fs) / (2 * M_PI));

    om_hat = fc_autocorrelation_multilag_rad(s, total_samples, 10);
    printf("Estimation de fc multilag: %f\n", (om_hat * qam.Fs) / (2 * M_PI));

    // Correction du signal avec Fc (en rad/sample) trouvé
    signal_correction(s, total_samples, om_hat);

    printf("Demodulation...\n");
    demodulate2(&qam, s, nb_symbols, output_bits.bits);
    */

    // Avant la demodulation
    printf("Test de blind carrier correction...\n");

    // Paramètres CMA simples
    double mu = 0.001;       // pas d'adaptation
    int num_iter = 100;      // nombre d'itérations

    // Appel de la fonction blind carrier correction
    blind_carrier(s, total_samples, qam.M);

    // Ensuite, utilise ton démodulateur actuel
    bit_array output_bits;
    output_bits.nb_bits = input_bits.nb_bits;
    output_bits.bits = (uint8_t*)malloc(output_bits.nb_bits * sizeof(uint8_t));

    demodulate2(&qam, s, nb_symbols, output_bits.bits);

    // Comparaison avec bits originaux
    double erreurs = taux_erreur_bits(input_bits.bits, input_bits.nb_bits, output_bits.bits, output_bits.nb_bits);
    printf("Taux d'erreur après blind carrier correction: %f %%\n", erreurs * 100);

    // Ecriture du fichier de Demodulation
    printf("Ecriture...\n");
    char *output_filename = make_output_filename(input_filename);
    bits_to_file(output_filename, &output_bits);

    // Libération mémoire
    free_bit_array(&input_bits);
    free_bit_array(&output_bits);
    free(symbols);
    free(s);
    free_constellation(&qam);
    free(output_filename);

    return 0;
}