/*
 * Copyright (c) 2010-2022 Belledonne Communications SARL.
 *
 * This file is part of mediastreamer2
 * (see https://gitlab.linphone.org/BC/public/mediastreamer2).
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include <mediastreamer2/dsptools.h>
#include <mediastreamer2/msequalizer.h>

#include <math.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#ifdef MS_FIXED_POINT
#define GAIN_ZERODB 20000
#else
#define GAIN_ZERODB 1.0f
#endif

#define EQUALIZER_DEFAULT_RATE 8000

typedef struct _EqualizerState {
	int rate;
	int nfft; /*number of fft points in time*/
	ms_word16_t *fft_cpx;
	int fir_len;
	ms_word16_t *fir;
	ms_mem_t *mem; /*memories for filtering computations*/
	bool_t needs_update;
	bool_t active;
} EqualizerState;

static void equalizer_state_flatten(EqualizerState *s) {
	int i;
	ms_word16_t val = (ms_word16_t)(GAIN_ZERODB / s->nfft);
	s->fft_cpx[0] = val;
	for (i = 1; i < s->nfft; i += 2)
		s->fft_cpx[i] = val;
}

static void equalizer_rate_update(EqualizerState *s, int rate) {
	int nFFT;

	if (rate < 16000) {
		nFFT = 128;
	} else if (rate < 32000) {
		nFFT = 256;
	} else {
		nFFT = 512;
	}

	s->rate = rate;
	s->nfft = nFFT;
	if (s->fft_cpx != NULL) ms_free(s->fft_cpx);
	s->fft_cpx = (ms_word16_t *)ms_new0(ms_word16_t, s->nfft);
	equalizer_state_flatten(s);
	s->fir_len = s->nfft;
	if (s->fir != NULL) ms_free(s->fir);
	s->fir = (ms_word16_t *)ms_new0(ms_word16_t, s->fir_len);
	if (s->mem != NULL) ms_free(s->mem);
	s->mem = (ms_mem_t *)ms_new0(ms_mem_t, s->fir_len);
	s->needs_update = TRUE;
}

static EqualizerState *equalizer_state_new(int rate) {
	EqualizerState *s = (EqualizerState *)ms_new0(EqualizerState, 1);
	equalizer_rate_update(s, rate);
	s->active = TRUE;
	return s;
}

static void equalizer_state_destroy(EqualizerState *s) {
	ms_free(s->fft_cpx);
	ms_free(s->fir);
	ms_free(s->mem);
	ms_free(s);
}

static int equalizer_state_hz_to_index(EqualizerState *s, int hz) {
	int ret;
	if (hz < 0) {
		ms_error("Bad frequency value %i", hz);
		return -1;
	}
	if (hz > (s->rate / 2)) {
		hz = (s->rate / 2);
	}
	/*round to nearest integer*/
	ret = ((hz * s->nfft) + (s->rate / 2)) / s->rate;
	if (ret == s->nfft / 2) ret = (s->nfft / 2) - 1;
	return ret;
}

static int equalizer_state_index2hz(EqualizerState *s, int index) {
	return (index * s->rate + s->nfft / 2) / s->nfft;
}

static float gain_float(ms_word16_t val) {
	return (float)val / GAIN_ZERODB;
}

static float equalizer_state_get(EqualizerState *s, int freqhz) {
	int idx = equalizer_state_hz_to_index(s, freqhz);
	if (idx >= 0) return gain_float(s->fft_cpx[idx * 2]) * s->nfft;
	return 0;
}

/* The natural peaking equalizer amplitude transfer function is multiplied to the discrete f-points.
 * Note that for PEQ no sqrt is needed for the overall calculation, applying it to gain yields the
 * same response.
 */
static float equalizer_compute_gainpoint(int f, int freq_0, float sqrt_gain, int freq_bw) {
	float k1, k2;
	k1 = ((float)(f * f) - (float)(freq_0 * freq_0));
	k1 *= k1;
	k2 = (float)(f * freq_bw);
	k2 *= k2;
	return (k1 + k2 * sqrt_gain) / (k1 + k2 / sqrt_gain);
}

static void equalizer_point_set(EqualizerState *s, int i, int f, float gain) {
	int index = 1 + ((i - 1) * 2);
	if (index >= 0 && index < s->nfft) {
		ms_message("Setting gain %f for freq_index %i (%i Hz)\n", gain, i, f);
		s->fft_cpx[index] = (s->fft_cpx[index] * (int)(gain * 32768)) / 32768;
	} else {
		ms_error("equalizer: invalid index %i for fft table of size %i", index, s->nfft);
	}
}

static void equalizer_state_set(EqualizerState *s, int freq_0, float gain, int freq_bw) {
	// int low,high;
	int i, f;
	int delta_f = equalizer_state_index2hz(s, 1);
	float sqrt_gain = (float)sqrt(gain);
	int mid = equalizer_state_hz_to_index(s, freq_0);
	freq_bw -= delta_f / 2; /* subtract a constant - compensates for limited fft steepness at low f */
	if (freq_bw < delta_f / 2) freq_bw = delta_f / 2;
	i = mid;
	f = equalizer_state_index2hz(s, i);
	equalizer_point_set(s, i, f, gain); /* gain according to argument */
	do { /* note: to better accomodate limited fft steepness, -delta is applied in f-calc ... */
		i++;
		f = equalizer_state_index2hz(s, i);
		gain = equalizer_compute_gainpoint(f - delta_f, freq_0, sqrt_gain, freq_bw);
		equalizer_point_set(s, i, f, gain);
	} while (i < s->nfft / 2 && (gain > 1.1 || gain < 0.9));
	i = mid;
	do { /* ... and here +delta, as to  */
		i--;
		f = equalizer_state_index2hz(s, i);
		gain = equalizer_compute_gainpoint(f + delta_f, freq_0, sqrt_gain, freq_bw);
		equalizer_point_set(s, i, f, gain);
	} while (i >= 0 && (gain > 1.1 || gain < 0.9));
	s->needs_update = TRUE;
}

static void dump_table(ms_word16_t *t, int len) {
	int i;
	for (i = 0; i < len; i++)
#ifdef MS_FIXED_POINT
		ms_message("[%i]\t%i", i, t[i]);
#else
		ms_message("[%i]\t%f", i, t[i]);
#endif
}

static void time_shift(ms_word16_t *s, int len) {
	int i;
	int half = len / 2;
	ms_word16_t tmp;
	for (i = 0; i < half; ++i) {
		tmp = s[i];
		s[i] = s[i + half];
		s[i + half] = tmp;
	}
}

/*
 *hamming:
 * 0.54 - 0.46*cos(2*M_PI*t/T)
 *
 * blackman
 * 0.42 - 0.5*cos(2*M_PI*t/T) + 0.08*cos(4*M_PI*t/T)
 */

static void norm_and_apodize(ms_word16_t *s, int len) {
	int i;
	float x;
	float w;
	for (i = 0; i < len; ++i) {
		x = (float)((float)i * 2 * M_PI / (float)len);
		w = (float)(0.54 - (0.46 * cos(x)));
		// w=0.42 - (0.5*cos(x)) + (0.08*cos(2*x));
		s[i] = w * (float)s[i];
	}
}

static void equalizer_state_compute_impulse_response(EqualizerState *s) {
	void *fft_handle = ms_fft_init(s->nfft);

	ms_message("Equalizer rate: %d, selecting %d steps for FFT", s->rate, s->nfft);

	ms_message("Spectral domain:");
	dump_table(s->fft_cpx, s->nfft);
	ms_ifft(fft_handle, s->fft_cpx, s->fir);
	ms_fft_destroy(fft_handle);
	/*
	ms_message("Inverse fft result:");
	dump_table(s->fir,s->fir_len);
	*/
	time_shift(s->fir, s->fir_len);
	/*
	ms_message("Time shifted:");
	dump_table(s->fir,s->fir_len);
	*/
	norm_and_apodize(s->fir, s->fir_len);
	ms_message("Apodized impulse response:");
	dump_table(s->fir, s->fir_len);
	s->needs_update = FALSE;
}

#ifdef MS_FIXED_POINT
#define INT16_TO_WORD16(i, w, l) w = (i)
#define WORD16_TO_INT16(i, w, l) i = (w)
#else

static void int16_to_word16(const int16_t *is, ms_word16_t *w, int l) {
	int i;
	for (i = 0; i < l; ++i) {
		w[i] = (ms_word16_t)is[i];
	}
}

static void word16_to_int16(const ms_word16_t *w, int16_t *is, int l) {
	int i;
	for (i = 0; i < l; ++i)
		is[i] = (int16_t)w[i];
}

#define INT16_TO_WORD16(i, w, l)                                                                                       \
	w = (ms_word16_t *)alloca(sizeof(ms_word16_t) * (l));                                                              \
	int16_to_word16(i, w, l)
#define WORD16_TO_INT16(w, i, l) word16_to_int16(w, i, l)
#endif

static void equalizer_state_run(EqualizerState *s, int16_t *samples, int nsamples) {
	ms_word16_t *w;
	if (s->needs_update) equalizer_state_compute_impulse_response(s);
	INT16_TO_WORD16(samples, w, nsamples);
	ms_fir_mem16(w, s->fir, w, nsamples, s->fir_len, s->mem);
	WORD16_TO_INT16(w, samples, nsamples);
}

static void equalizer_init(MSFilter *f) {
	f->data = equalizer_state_new(EQUALIZER_DEFAULT_RATE);
}

static void equalizer_uninit(MSFilter *f) {
	equalizer_state_destroy((EqualizerState *)f->data);
}

static void equalizer_process(MSFilter *f) {
	mblk_t *m;
	EqualizerState *s = (EqualizerState *)f->data;
	while ((m = ms_queue_get(f->inputs[0])) != NULL) {
		if (s->active) {
			equalizer_state_run(s, (int16_t *)m->b_rptr, (int)((m->b_wptr - m->b_rptr) / 2));
		}
		ms_queue_put(f->outputs[0], m);
	}
}

static int equalizer_set_gain(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	MSEqualizerGain *d = (MSEqualizerGain *)data;
	equalizer_state_set(s, (int)d->frequency, d->gain, (int)d->width);
	return 0;
}

static int equalizer_get_gain(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	MSEqualizerGain *d = (MSEqualizerGain *)data;
	d->gain = equalizer_state_get(s, (int)d->frequency);
	d->width = 0;
	return 0;
}

static int equalizer_set_rate(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	equalizer_rate_update(s, *(int *)data);
	return 0;
}

static int equalizer_set_active(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	s->active = *(bool_t *)data;
	return 0;
}

static int equalizer_dump(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	float *t = (float *)data;
	int i;
	*t = s->fft_cpx[0];
	t++;
	for (i = 1; i < s->nfft; i += 2) {
		*t = ((float)s->fft_cpx[i] * (float)s->nfft) / (float)GAIN_ZERODB;
		t++;
	}
	return 0;
}

static int equalizer_get_nfreqs(MSFilter *f, void *data) {
	EqualizerState *s = (EqualizerState *)f->data;
	*(int *)data = s->nfft / 2;
	return 0;
}

static MSFilterMethod equalizer_methods[] = {{MS_EQUALIZER_SET_GAIN, equalizer_set_gain},
                                             {MS_EQUALIZER_GET_GAIN, equalizer_get_gain},
                                             {MS_EQUALIZER_SET_ACTIVE, equalizer_set_active},
                                             {MS_FILTER_SET_SAMPLE_RATE, equalizer_set_rate},
                                             {MS_EQUALIZER_DUMP_STATE, equalizer_dump},
                                             {MS_EQUALIZER_GET_NUM_FREQUENCIES, equalizer_get_nfreqs},
                                             {0, NULL}};

#ifdef _MSC_VER

MSFilterDesc ms_equalizer_desc = {MS_EQUALIZER_ID,
                                  "MSEqualizer",
                                  N_("Parametric sound equalizer."),
                                  MS_FILTER_OTHER,
                                  NULL,
                                  1,
                                  1,
                                  equalizer_init,
                                  NULL,
                                  equalizer_process,
                                  NULL,
                                  equalizer_uninit,
                                  equalizer_methods};

#else

MSFilterDesc ms_equalizer_desc = {.id = MS_EQUALIZER_ID,
                                  .name = "MSEqualizer",
                                  .text = N_("Parametric sound equalizer."),
                                  .category = MS_FILTER_OTHER,
                                  .ninputs = 1,
                                  .noutputs = 1,
                                  .init = equalizer_init,
                                  .process = equalizer_process,
                                  .uninit = equalizer_uninit,
                                  .methods = equalizer_methods};

#endif

bctbx_list_t *ms_parse_equalizer_string(const char *str) {
	bctbx_list_t *eq_list = NULL;
	do {
		int bytes;
		MSEqualizerGain g;
		if (sscanf(str, "%f:%f:%f %n", &g.frequency, &g.gain, &g.width, &bytes) == 3) {
			MSEqualizerGain *gain = ms_new(MSEqualizerGain, 1);
			*gain = g;
			eq_list = bctbx_list_append(eq_list, gain);
			str += bytes;
		} else break;
	} while (1);
	return eq_list;
}

MS_FILTER_DESC_EXPORT(ms_equalizer_desc)
