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

#define A 10

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) / norm_factor; 
        }
    }
}

// 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++) {
            int idx = k * qam->N + n;
            s[idx] = A * iq * cexp(2 * I * M_PI * qam->Fc * ((double)idx / qam->Fs));
        }
    }
}

// Demodulation QAM 
void demodulate(qam_system* qam, double complex* s, int nb_symbols, uint8_t* bits_hat, FILE *fp_constel) {
    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)(k * qam->N + n) / qam->Fs)) / A;
        }
        r /= qam->N;

        if (fp_constel) {
            fprintf(fp_constel, "% .8f % .8f\n", creal(r), cimag(r));
            fflush(fp_constel); 
        }

        // 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 = 0;
        int 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;
        }
    }
}

// PLL pour corriger le déphasage
void pll_qam_symbol(qam_system* qam, double complex* symbols_rx, double complex* r_corr, int nb_symbols, double Kp, double Ki, double alpha, FILE* fp_error) {
    double phase_est = 0.0;
    double integrator = 0.0;
    double filtered_error = 0.0;

    int sm = (int)sqrt(qam->M);
    int N = qam->N;

    for (int k = 0; k < nb_symbols; k++) {
        double complex r_symbol = 0;
        for (int n = 0; n < N; n++) {
            int idx = k * N + n;
            r_symbol += symbols_rx[idx] * cexp(-2.0 * I * M_PI * qam->Fc * ((double)idx / qam->Fs));
        }
        r_symbol /= N;

        r_symbol *= cexp(-I * phase_est);

        double min_d = INFINITY;
        double complex closest = 0;
        for (int i = 0; i < sm; i++) {
            for (int j = 0; j < sm; j++) {
                double d = cabs(r_symbol - qam->constellation[i][j]);
                if (d < min_d) {
                    min_d = d;
                    closest = qam->constellation[i][j];
                }
            }
        }

        double error = carg(r_symbol * conj(closest));

        filtered_error = (1.0 - alpha) * filtered_error + alpha * error;
        integrator += Ki * filtered_error;
        phase_est += Kp * filtered_error + integrator;

        // Écriture de l'erreur PLL dans le fichier
        if (fp_error) {
            fprintf(fp_error, "%d % .8f\n", k, 100 * filtered_error);
            fflush(fp_error);
        }

        for (int n = 0; n < N; n++) {
            int idx = k * N + n;
            r_corr[idx] = symbols_rx[idx] * cexp(-I * phase_est);
        }
    }
}


// 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);
}

double compare_bits(uint8_t* bits1, uint8_t* bits2, int nb_bits) {
    int errors = 0;
    for (int i = 0; i < nb_bits; i++) {
        if (bits1[i] != bits2[i]) errors++;
    }
    return (double)errors / nb_bits;
}

// Minimise le BER (si la pll s'est lockée de maniere déphasée de k*pi/2) 
void demodulate2(qam_system* qam, double complex* r_corr, int nb_symbols, uint8_t* input_bits, uint8_t* output_bits, FILE* fp_constel) {

    int nb_bits = nb_symbols * qam->k;
    double best_ber = INFINITY;
    double best_angle = 0.0;
    uint8_t* temp_bits = (uint8_t*)malloc(nb_bits * sizeof(uint8_t));

    for (int r = 0; r < 4; r++) {
        double angle = r * M_PI/2;

        double complex* rotated = (double complex*)malloc(sizeof(double complex) * nb_symbols * qam->N);
        for (int i = 0; i < nb_symbols * qam->N; i++) {
            rotated[i] = r_corr[i] * cexp(I * angle);
        }

        demodulate(qam, rotated, nb_symbols, temp_bits, fp_constel);

        double ber = compare_bits(input_bits, temp_bits, nb_bits);

        if (ber < best_ber) {
            best_ber = ber;
            best_angle = angle;
        }

        free(rotated);
    }

    for (int i = 0; i < nb_symbols * qam->N; i++) {
        r_corr[i] *= cexp(I * best_angle);
    }

    demodulate(qam, r_corr, nb_symbols, output_bits, fp_constel);

    free(temp_bits);
}


int main () {
    qam_system qam;
    qam.M = 16;
    qam.k = (int)log2((double)(qam.M));
    qam.Fs = 44100;
    //qam.Ts = 0.0003;
    //qam.N = (int)qam.Fs * qam.Ts;
    qam.Ts = 0.01;
    qam.N = (int)(qam.Fs * qam.Ts);
    qam.Fc = 2000;
    init_constellation(&qam);

    //int nb_bits = 1000;
    //int nb_symbols = nb_bits / qam.k;

    //uint8_t* input_bits = malloc(nb_bits * sizeof(uint8_t));
    //for (int i = 0; i < nb_bits; i++) {
    //    input_bits[i] = rand() % 2;
    //}
    char* texte = "Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux, Vif juge, trempez ce blond whisky aqueux";  
    int nb_chars = strlen(texte);
    int nb_bits = nb_chars * 8;
    int nb_symbols = (nb_bits + qam.k - 1) / qam.k;

    // Conversion du texte en bits
    uint8_t* input_bits = malloc(nb_bits * sizeof(uint8_t));
    for(int i = 0; i < nb_chars; i++){
        for(int b = 0; b < 8; b++){
            input_bits[i*8 + b] = (texte[i] >> (7-b)) & 1;
        }
    }
    
    // Conversion en symboles
    double complex* symbols = malloc(sizeof(double complex) * nb_symbols);
    bits_to_symbols(&qam, input_bits, nb_bits, symbols);

    // Modulation
    int total_samples = qam.N * nb_symbols;
    double complex* s = 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
    double snr_dB = 5; // SNR en dB
    double snr_lin = pow(10.0, snr_dB / 10.0);
    double sigma = sqrt(signal_power / snr_lin);
    add_noise(s, total_samples, 10);


    FILE *fp_ref = fopen("constellation_ref.dat", "w");
    int sm = (int)sqrt(qam.M);
    for (int i = 0; i < sm; i++) {
        for (int j = 0; j < sm; j++) {
            fprintf(fp_ref, "% .8f % .8f\n", creal(qam.constellation[i][j]), cimag(qam.constellation[i][j]));
        }
    }
    fclose(fp_ref);
    FILE *fp_constel = fopen("constellation.dat", "w");

    // Ajout de dephasage
    //double phase_offset = M_PI / 6.0;  // 30 degrés
    //for (int i = 0; i < total_samples; i++) {
    //    s[i] *= cexp(I * phase_offset);
    //}
    
    // AJout de decalage de fréquence
    double freq_offset = 1; // Hz de décalage
    for (int i = 0; i < total_samples; i++) {
        double t = (double)i / qam.Fs;
        s[i] *= cexp(I * 2 * M_PI * freq_offset * t);
    }
    
    // Ajout de decalage entre les symbole
    //int offset_samples = (int)(0.3 * qam.N); // décalage de 30% d’un symbole
    //memmove(s + offset_samples, s, (total_samples - offset_samples) * sizeof(double complex));

    double complex* r_corr = malloc(sizeof(double complex) * total_samples);
    double Kp = 0.2;
    double Ki = 0.02;
    double alpha = 0.1;
    FILE* fp_error = fopen("pll_error.dat", "w");
    pll_qam_symbol(&qam, s, r_corr, nb_symbols, Kp, Ki, alpha, fp_error);
    fclose(fp_error);

    // Démodulation
    uint8_t* output_bits = (uint8_t*)malloc(nb_bits * sizeof(uint8_t));
    demodulate2(&qam, r_corr, nb_symbols, input_bits, output_bits, fp_constel);
    //demodulate(&qam, r_corr, nb_symbols, output_bits, fp_constel);

    fclose(fp_constel);

    // Reconstruction du texte
    char* texte_recup = malloc(nb_chars + 1);
    for(int i = 0; i < nb_chars; i++){
        char c = 0;
        for(int b = 0; b < 8; b++){
            c |= output_bits[i*8 + b] << (7-b);
        }
        texte_recup[i] = c;
    }
    texte_recup[nb_chars] = '\0';
    printf("Texte original : %s\n\n", texte);
    printf("Texte demodulé : %s\n", texte_recup);

    // Calcul du BER
    double ber = compare_bits(input_bits, output_bits, nb_bits);
    printf("Taux d'erreur blind QAM: %.4f\n", ber * 100);

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

    return 0;
}