#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <complex.h>
#include "../WAV/wav.h"

#define A 10000

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;

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] = A * (ip + I * qp) / norm_factor; 
        }
    }
}

// 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, int* 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];
    }
}

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

void demodulate(qam_system* qam, double complex* s, int nb_symbols, int* bits_hat, double sigma) {
    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
        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;
                }
            }
        }

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

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

int main () {
    qam_system qam;
    qam.M = 256;
    qam.k = (int)log2((double)(qam.M));
    qam.Fs = 44100;
    qam.Ts = 0.00003;
    qam.N = (int)qam.Fs * qam.Ts;
    qam.Fc = 2000;
    init_constellation(&qam);

    // Nombre de bit multiple de k sinon remplir de zero jusqu'a ce que ce le soit
    int bits[16] = {1,0,1,1,  0,1,1,0,  1,1,0,0,  0,0,0,1};
    int nb_bits = 16;
    int nb_symbols = 16 / qam.k;


    double complex symbols[nb_symbols];
    bits_to_symbols(&qam, bits, nb_bits, symbols);

    int total_samples = qam.N * nb_symbols;
    double complex* s = (double complex*)malloc(sizeof(double complex) * total_samples);
    modulate(&qam, symbols, nb_symbols, s);

    int bits_hat[nb_bits];
    demodulate(&qam, s, nb_symbols, bits_hat, 0.0);

    for (int i = 0; i < nb_bits; i++) 
        printf("%d", bits[i]);
    printf("\n");

    for (int i = 0; i < nb_bits; i++) 
        printf("%d", bits_hat[i]);
    printf("\n");

    double* si = (double*)malloc(sizeof(double) * total_samples);

    for (int i = 0; i < total_samples; i++) {
       si[i] = cimag(s[i]); 
    }

    write_wav("s.wav", si, total_samples);

    free_constellation(&qam);
    free(s);
    free(si);

    return 0;
}