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

#define PI 3.14159265358979323846

#pragma pack(push, 1)
typedef struct {
    char riff[4];
    unsigned int fileSize;
    char wave[4];
    char fmt[4];
    unsigned int fmtSize;
    unsigned short audioFormat;
    unsigned short numChannels;
    unsigned int sampleRate;
    unsigned int byteRate;
    unsigned short blockAlign;
    unsigned short bitsPerSample;
    char data[4];
    unsigned int dataSize;
} WavHeader;
#pragma pack(pop)

typedef struct {
    double real;
    double imag;
} Complex;

/* ===================== HELPERS ===================== */

int next_pow2(int n)
{
    int p = 1;
    while (p < n)
        p <<= 1;
    return p;
}

Complex complex_create(double real, double imag)
{
    Complex c;
    c.real = real;
    c.imag = imag;
    return c;
}

Complex complex_add(Complex a, Complex b)
{
    return complex_create(a.real + b.real, a.imag + b.imag);
}

Complex complex_sub(Complex a, Complex b)
{
    return complex_create(a.real - b.real, a.imag - b.imag);
}

Complex complex_mul(Complex a, Complex b)
{
    return complex_create(
        a.real * b.real - a.imag * b.imag,
        a.real * b.imag + a.imag * b.real
    );
}

/* ===================== FFT ===================== */

void fft(Complex *x, int n)
{
    int bits = 0;
    int temp = n;

    while (temp > 1) {
        temp >>= 1;
        bits++;
    }

    int *rev = malloc(n * sizeof(int));

    for (int i = 0; i < n; i++) {
        int j = i;
        int r = 0;

        for (int k = 0; k < bits; k++) {
            r = (r << 1) | (j & 1);
            j >>= 1;
        }

        rev[i] = r;
    }

    for (int i = 0; i < n; i++) {
        if (rev[i] > i) {
            Complex tmp = x[i];
            x[i] = x[rev[i]];
            x[rev[i]] = tmp;
        }
    }

    free(rev);

    for (int s = 1; s <= bits; s++) {
        int m = 1 << s;

        for (int k = 0; k < n; k += m) {
            for (int j = 0; j < m / 2; j++) {

                double angle = -2.0 * PI * j / m;
                Complex w = complex_create(cos(angle), sin(angle));

                Complex t = complex_mul(w, x[k + j + m / 2]);
                Complex u = x[k + j];

                x[k + j] = complex_add(u, t);
                x[k + j + m / 2] = complex_sub(u, t);
            }
        }
    }
}

/* ===================== IFFT ===================== */

void ifft(Complex *x, int n)
{
    for (int i = 0; i < n; i++)
        x[i].imag = -x[i].imag;

    fft(x, n);

    for (int i = 0; i < n; i++) {
        x[i].real /= n;
        x[i].imag = -x[i].imag / n;
    }
}

/* ===================== WAV READ ===================== */

int read_wav(const char *filename,
             int *sampleRate,
             int *numSamples,
             double **signal)
{
    FILE *file = fopen(filename, "rb");

    if (!file) {
        fprintf(stderr, "Error opening file\n");
        return 1;
    }

    WavHeader header;
    fread(&header, sizeof(WavHeader), 1, file);

    if (memcmp(header.riff, "RIFF", 4) != 0 ||
        memcmp(header.wave, "WAVE", 4) != 0) {
        fprintf(stderr, "Invalid WAV file\n");
        return 1;
    }

    if (header.audioFormat != 1 ||
        header.numChannels != 1 ||
        header.bitsPerSample != 16) {
        fprintf(stderr, "Only 16-bit mono PCM supported\n");
        return 1;
    }

    *sampleRate = header.sampleRate;
    *numSamples = header.dataSize / 2;

    *signal = malloc(*numSamples * sizeof(double));

    for (int i = 0; i < *numSamples; i++) {
        unsigned char b0 = fgetc(file);
        unsigned char b1 = fgetc(file);

        short s = (short)((b1 << 8) | b0);
        (*signal)[i] = (double)s;
    }

    fclose(file);
    return 0;
}

/* ===================== WAV WRITE ===================== */

int write_wav(const char *filename,
              int sampleRate,
              int numSamples,
              double *signal)
{
    FILE *file = fopen(filename, "wb");

    if (!file) {
        fprintf(stderr, "Failed to open output file\n");
        return 1;
    }

    int dataSize = numSamples * 2;

    WavHeader header;

    memcpy(header.riff, "RIFF", 4);
    header.fileSize = 36 + dataSize;
    memcpy(header.wave, "WAVE", 4);
    memcpy(header.fmt, "fmt ", 4);
    header.fmtSize = 16;
    header.audioFormat = 1;
    header.numChannels = 1;
    header.sampleRate = sampleRate;
    header.byteRate = sampleRate * 2;
    header.blockAlign = 2;
    header.bitsPerSample = 16;
    memcpy(header.data, "data", 4);
    header.dataSize = dataSize;

    fwrite(&header, sizeof(WavHeader), 1, file);

    for (int i = 0; i < numSamples; i++) {

        double s = signal[i];

        if (s > 32767.0) s = 32767.0;
        if (s < -32768.0) s = -32768.0;

        short sample = (short)s;
        unsigned short us = (unsigned short)sample;

        fputc(us & 0xff, file);
        fputc((us >> 8) & 0xff, file);
    }

    fclose(file);
    return 0;
}

/* ===================== FIXED BANDPASS FILTER ===================== */

void bandpass_filter(Complex *spectrum, int n, int sampleRate)
{
    int hums[] = {60, 120, 180, 240, 300};
    int width = 5;

    for (int i = 0; i < n; i++) {

        double freq = (double)i * sampleRate / n;

        for (int h = 0; h < 5; h++) {

            if ((freq >= hums[h] - width &&
                 freq <= hums[h] + width) ||

                (freq >= sampleRate - hums[h] - width &&
                 freq <= sampleRate - hums[h] + width))
            {
                spectrum[i].real = 0.0;
                spectrum[i].imag = 0.0;
            }
        }
    }
}

/* ===================== NORMALIZATION ===================== */

void normalize_signal(double *signal, int n)
{
    double peak = 0.0;

    for (int i = 0; i < n; i++) {
        double v = fabs(signal[i]);
        if (v > peak) peak = v;
    }

    if (peak == 0) return;

    double scale = 32767.0 / peak;

    for (int i = 0; i < n; i++)
        signal[i] *= scale;
}

/* ===================== MAIN ===================== */

int main(int argc, char **argv)
{
    if (argc != 3) {
        fprintf(stderr, "Usage: %s input.wav output.wav\n", argv[0]);
        return 1;
    }

    int sr, n;
    double *signal;

    if (read_wav(argv[1], &sr, &n, &signal)) {
        return 1;
    }

    int N = next_pow2(n);

    double *buf = calloc(N, sizeof(double));

    for (int i = 0; i < n; i++)
        buf[i] = signal[i];

    Complex *x = malloc(N * sizeof(Complex));

    for (int i = 0; i < N; i++) {
        x[i].real = buf[i];
        x[i].imag = 0;
    }

    fft(x, N);
    bandpass_filter(x, N, sr);
    ifft(x, N);

    for (int i = 0; i < N; i++)
        buf[i] = x[i].real;

    // normalize_signal(buf, n);

    write_wav(argv[2], sr, n, buf);

    printf("Done: filtered + reconstructed audio written.\n");

    free(signal);
    free(buf);
    free(x);

    return 0;
}