#include <stdio.h>
#include <math.h>
#include <limits.h>
#include <complex.h>
#include <pthread.h>
#include <fftw3.h>
#include "jl_mem.h"
#include "soundfile.h"
#include "fx.h"

void fx_amp(FSample *s, double gain) {
    for(int i = 0; i < s->len; ++i) {
        s->L[i] *= gain;
        s->R[i] *= gain;
    }
}

void fx_balance(FSample *s, double balance) {
    if(balance == 0) {
        return;
    }

    if(balance < -1) {
        balance = -1;
    } else if(balance > 1) {
        balance = 1;
    }

    if(balance < 0) {
        for(int i = 0; i < s->len; ++i) {
            s->L[i] *= 1 + balance;
        }
    } else {
        for(int i = 0; i < s->len; ++i) {
            s->R[i] *= 1 - balance;
        }
    }
}

void fx_divide_rms(FSample *s, double t_ms) {
    int di = t_ms * SAMPLE_RATE / 1000.0;
    if(di > s->len) {
        di = s->len;
    }

    double rms = 0;
    double x;
    for(int i = 0; i < di; ++i) {
        x = s->L[i];
        rms += x*x;
        x = s->R[i];
        rms += x*x;
    }

    rms = sqrt(rms / (2.0 * di));
    fx_amp(s, 1.0 / rms);
}

void fx_fade_in(FSample *s, double fade_ms) {
    int n_fade = fade_ms * SAMPLE_RATE / 1000.0;
    for(int i = 0; i < n_fade && i < s->len; ++i) {
        double val = (double)i / (double)n_fade;
        s->L[i] *= val;
        s->R[i] *= val;
    }
}

pthread_mutex_t fftw_mutex = PTHREAD_MUTEX_INITIALIZER;

static fftw_complex *fft(int fft_len, int in_len, float *in_data) {

    double *in = jl_malloc_exit(fft_len * sizeof(double));
    for(int i = 0; i < in_len; ++i) {
        in[i] = in_data[i];
    }
    for(int i = in_len; i < fft_len; ++i) {
        in[i] = 0;
    }

    fftw_complex *out = jl_malloc_exit(fft_len * sizeof(fftw_complex));

    pthread_mutex_lock(&fftw_mutex);
    fftw_plan p = fftw_plan_dft_r2c_1d(fft_len, in, out, FFTW_ESTIMATE);
    pthread_mutex_unlock(&fftw_mutex);

    fftw_execute(p);


    pthread_mutex_lock(&fftw_mutex);
    fftw_destroy_plan(p);
    pthread_mutex_unlock(&fftw_mutex);

    jl_free(in);
    return out;
}

static float *rfft(int fft_len, fftw_complex *in, int out_len) {

    double *out = jl_malloc_exit(fft_len * sizeof(double));

    pthread_mutex_lock(&fftw_mutex);
    fftw_plan p = fftw_plan_dft_c2r_1d(fft_len, in, out, FFTW_ESTIMATE);
    pthread_mutex_unlock(&fftw_mutex);

    fftw_execute(p);

    pthread_mutex_lock(&fftw_mutex);
    fftw_destroy_plan(p);
    pthread_mutex_unlock(&fftw_mutex);

    float *f_out = jl_malloc_exit(out_len * sizeof(float));

    for(int i = 0; i < out_len; ++i) {
        f_out[i] = out[i] / (float)fft_len;
    }

    jl_free(out);
    return f_out;
}

void fx_ir(FSample *s, const char *path, double mix) {
    if(mix <= 0) {
        return;
    }

    FSample ir;
    ir.len = soundfile_load_stereo(path, &ir.L, &ir.R);

    int out_len = ir.len + s->len - 1;

    // Using a power of 2 fft length makes the computations MUCH faster.
    int fft_len = pow(2, ceil(log2(out_len)));

    fftw_complex *ir_L = fft(fft_len, ir.len, ir.L);
    fftw_complex *ir_R = fft(fft_len, ir.len, ir.R);

    fsample_free(&ir);

    fftw_complex *L = fft(fft_len, s->len, s->L);
    fftw_complex *R = fft(fft_len, s->len, s->R);

    // Convolve.
    for(int i = 0; i < fft_len; ++i) {
        L[i] *= ir_L[i];
        R[i] *= ir_R[i];
    }

    // Reverse fft.
    float *L_wet = rfft(fft_len, L, out_len);
    float *R_wet = rfft(fft_len, R, out_len);

    if(mix < 1) {
        for(int i = 0; i < out_len; ++i) {
            L_wet[i] *= mix;
            R_wet[i] *= mix;
        }
        for(int i = 0; i < s->len; ++i) {
            L_wet[i] += (1 - mix) * s->L[i];
            R_wet[i] += (1 - mix) * s->R[i];
        }
    }

    jl_free(ir_L);
    jl_free(ir_R);
    jl_free(L);
    jl_free(R);

    fsample_set_data(s, L_wet, R_wet, out_len);
}

void fx_mono(FSample *s) {
    for(int i = 0; i < s->len; ++i) {
        s->L[i] += s->R[i];
        s->L[i] /= 2.0;
        s->R[i] = s->L[i];
    }
}

void fx_pan(FSample *s, double pan) {
    if(pan < 0) {
        for(int i = 0; i < s->len; ++i) {
            s->L[i] += s->R[i] * -pan;
            s->R[i] *= (1 + pan);
        }
    } else {
        for(int i = 0; i < s->len; ++i) {
            s->R[i] += s->L[i] * pan;
            s->L[i] *= (1 - pan);
        }
    }
}

static float interp(float *v, double idx) {
    int i = (int)idx;
    double mu = idx - i;
    return mu * v[i+1] + (1 - mu) * v[i];
}

void fx_playback_speed(FSample *s, double speed) {
    int len = s->len / speed;
    float *L = jl_malloc_exit(len * sizeof(float));
    float *R = jl_malloc_exit(len * sizeof(float));

    double idx = 0;
    for(int i = 0; i < len; ++i) {
        L[i] = interp(s->L, idx);
        R[i] = interp(s->R, idx);
        idx += speed;
    }

    fsample_set_data(s, L, R, len);
}

void fx_pre_cut(FSample *s, double pct, double fade_ms) {
    float th = pct / 100.0;

    int i_cut = 0;

    for(i_cut = 0; i_cut < s->len; ++i_cut) {
        if(s->L[i_cut] > th ||
           s->R[i_cut] > th ||
           s->L[i_cut] < -th ||
           s->R[i_cut] < -th) {
            break;
        }
    }

    if(i_cut == s->len) {
        printf("Warning: pre-cut threshold not reached.\n");
        fx_fade_in(s, fade_ms);
        return;
    }

    i_cut -= fade_ms * SAMPLE_RATE / 1000.0;
    if(i_cut <= 0) {
        fx_fade_in(s, fade_ms);
        return;
    }

    if(i_cut > SAMPLE_RATE * 5 / 1000) {
        printf("Warning: cutting %i ms from sample.\n",
               i_cut * 1000 / SAMPLE_RATE);
    }

    int len = s->len - i_cut;

    // Allocate arrays.
    float *L = jl_malloc_exit(len * sizeof(float));
    float *R = jl_malloc_exit(len * sizeof(float));

    for(int i = 0; i < len; ++i) {
        L[i] = s->L[i_cut + i];
        R[i] = s->R[i_cut + i];
    }

    fsample_set_data(s, L, R, len);
    fx_fade_in(s, fade_ms);
}

static float *delay_array(float *in, int in_len, int out_len, int delay) {
    float *out = jl_malloc_exit(out_len * sizeof(float));
    for(int i = 0; i < delay; ++i) {
        out[i] = 0;
    }
    for(int i = 0; i < in_len; ++i) {
        out[delay + i] = in[i];
    }
    for(int i = delay + in_len; i < out_len; ++i) {
        out[i] = 0;
    }

    return out;
}

void fx_pre_delay(FSample *s, double left_ms, double right_ms) {
    int di_L = left_ms * SAMPLE_RATE / 1000;
    int di_R = right_ms * SAMPLE_RATE / 1000;

    int out_len = s->len;
    out_len += di_L > di_R ? di_L : di_R;

    float *L = delay_array(s->L, s->len, out_len, di_L);
    float *R = delay_array(s->R, s->len, out_len, di_R);

    fsample_set_data(s, L, R, out_len);
}

void fx_rc_highpass(FSample *s, double freq, int order) {
    fx_rc_lowshelf(s, freq, order, -1);
}

void fx_rc_lowpass(FSample *s, double freq, int order) {
    fx_rc_highshelf(s, freq, order, -1);
}

static double freq_3db(double freq, int order) {
     return freq / sqrt(pow(2.0, 1.0 / ((double)order)) - 1.0);
}

static void rc_highshelf1(float *data, int len, double freq, double gain) {
    double rc = 1.0 / (freq * 2.0 * M_PI);
    double dt = 1.0 / SAMPLE_RATE;
    double alpha = dt / (rc + dt);
    double prev = 0;

    for (int i = 0; i < len; ++i) {
        prev *= (1 - alpha);
        prev += alpha * data[i];
        data[i] += gain * (data[i] - prev);
    }
}

static void rc_lowshelf1(float *data, int len, double freq, double gain) {
    double rc = 1.0 / (freq * 2.0 * M_PI);
    double dt = 1.0 / SAMPLE_RATE;
    double alpha = dt / (rc + dt);
    double prev = 0;

    for (int i = 0; i < len; ++i) {
        prev *= (1 - alpha);
        prev += alpha * data[i];
        data[i] += gain * prev;
    }
}

void fx_rc_highshelf(FSample *s, double freq, int order, double gain) {
    freq = freq_3db(freq, order);
    for(int i = 0; i < order; ++i) {
        rc_highshelf1(s->L, s->len, freq, gain);
        rc_highshelf1(s->R, s->len, freq, gain);
    }
}

void fx_rc_lowshelf(FSample *s, double freq, int order, double gain) {
    freq = freq_3db(freq, order);
    for(int i = 0; i < order; ++i) {
        rc_lowshelf1(s->L, s->len, freq, gain);
        rc_lowshelf1(s->R, s->len, freq, gain);
    }
}

// See http://musicdsp.org/showArchiveComment.php?ArchiveID=256
void fx_stereo_width(FSample *s, double width) {
    float w2 = width * 0.5;
    float m, p;

    for(int i = 0; i < s->len; ++i) {
        m = (s->L[i] + s->R[i]) * 0.5;
        p = (s->R[i] - s->L[i]) * w2;
        s->L[i] = m - p;
        s->R[i] = m + p;
    }
}