---

#include "imx_ai_aecnr_core.h"
#include "mediastreamer2/msfilter.h"
#include "mediastreamer2/flowcontrol.h"

#define EC_DUMP
#define USE_CONVERT_NEON

typedef struct ImxAiAecState {
	struct imx_ai_aecnr_core *aec_core;
	MSBufferizer delayed_ref;
	MSFlowControlledBufferizer ref;
	MSBufferizer echo;
	char *state_str;
	int samplerate;
	int framesize;
	int delay_ms;
	int nominal_ref_samples;
	bool_t echostarted;
	bool_t bypass_mode;
	bool_t using_zeroes;
#ifdef EC_DUMP
	FILE *echofile;
	FILE *reffile;
	FILE *cleanfile;
	FILE *echofile_float;
	FILE *reffile_float;
	FILE *cleanfile_float;
#endif
} ImxAiAecState;

#ifdef USE_CONVERT_NEON
static void s16_to_float_neon(const int16_t* src, float* dst) {
    asm volatile (
        "ld1    {v0.8h}, [%[in]]        \n\t"
        "sxtl   v1.4s, v0.4h            \n\t"
        "sxtl2  v2.4s, v0.8h            \n\t"
        "scvtf  v1.4s, v1.4s            \n\t"
        "scvtf  v2.4s, v2.4s            \n\t"
        "st1    {v1.4s, v2.4s}, [%[out]] \n\t"
        : 
        : [in] "r" (src), [out] "r" (dst)
        : "v0", "v1", "v2", "memory"
    );
}

void float_to_s16_neon(const float* src, int16_t* dst) {
    asm volatile (
        "ld1    {v1.4s, v2.4s}, [%[in]] \n\t"
        "fcvtns v1.4s, v1.4s            \n\t"
        "fcvtns v2.4s, v2.4s            \n\t"
        "sqxtn  v0.4h, v1.4s            \n\t"
        "sqxtn2 v0.8h, v2.4s            \n\t"
        "st1    {v0.8h}, [%[out]]       \n\t"
        : 
        : [in] "r" (src), [out] "r" (dst)
        : "v0", "v1", "v2", "memory"
    );
}
#endif

static void imx_ai_aec_init(MSFilter *f) {
	ImxAiAecState *s = ms_new0(ImxAiAecState, 1);
	s->samplerate = AI_AECNR_SAMPLE_RATE; // 16000
	s->framesize = AI_AECNR_IO_SAMPLES;   // 256
	
	ms_bufferizer_init(&s->delayed_ref);
	ms_bufferizer_init(&s->echo);
	ms_flow_controlled_bufferizer_init(&s->ref, f, s->samplerate, 1);

	s->state_str = NULL;
	s->delay_ms = 0;
	s->echostarted = FALSE;
	s->bypass_mode = FALSE;
	s->using_zeroes = FALSE;
	
	f->data = s;
#ifdef EC_DUMP
	{
		char *fname = ms_strdup_printf("/tmp/imx-%p-echo.raw", f);
		s->echofile = fopen(fname, "wb");
		ms_free(fname);

		fname = ms_strdup_printf("/tmp/imx-%p-ref.raw", f);
		s->reffile = fopen(fname, "wb");
		ms_free(fname);

		fname = ms_strdup_printf("/tmp/imx-%p-clean.raw", f);
		s->cleanfile = fopen(fname, "wb");
		ms_free(fname);

		fname = ms_strdup_printf("/tmp/imx-%p-float_echo.raw", f);
		s->echofile_float = fopen(fname, "wb");
		ms_free(fname);

		fname = ms_strdup_printf("/tmp/imx-%p-float_ref.raw", f);
		s->reffile_float = fopen(fname, "wb");
		ms_free(fname);

		fname = ms_strdup_printf("/tmp/imx-%p-float_clean.raw", f);
		s->cleanfile_float = fopen(fname, "wb");
		ms_free(fname);
		ms_message("%s(%d) echofile=%p\n",__FUNCTION__,__LINE__,s->echofile);
		ms_message("%s(%d) reffile=%p\n",__FUNCTION__,__LINE__,s->reffile);
		ms_message("%s(%d) cleanfile=%p\n",__FUNCTION__,__LINE__,s->cleanfile);
		ms_message("%s(%d) echofile_float=%p\n",__FUNCTION__,__LINE__,s->echofile_float);
		ms_message("%s(%d) reffile_float=%p\n",__FUNCTION__,__LINE__,s->reffile_float);
		ms_message("%s(%d) cleanfile_float=%p\n",__FUNCTION__,__LINE__,s->cleanfile_float);
	}
#endif
ms_message("%s(%d) use samplerate=%u framesize=%u\n",__FUNCTION__,__LINE__,s->samplerate,s->framesize);
}

static void imx_ai_aec_preprocess(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	imx_ai_aecnr_config config;
	int delay_samples;
	mblk_t *m;

	config.use_small_model = 0; 
	s->aec_core = imx_ai_aecnr_core_open(&config);
	
	if (s->aec_core == NULL) {
		ms_error("imx_ai_aecnr_core_open() failed, bypass mode");
		s->bypass_mode = TRUE;
		return;
	}

	s->echostarted = FALSE;
	delay_samples = s->delay_ms * s->samplerate / 1000;
	s->nominal_ref_samples = delay_samples;

	ms_message("IMX AI AECNR initialized: %s, framesize=%i, delay=%i ms\n", 
				config.version, s->framesize, s->delay_ms);

	// flow control for synchronized thread's
	ms_flow_controlled_bufferizer_set_samplerate(&s->ref, s->samplerate);
	ms_flow_controlled_bufferizer_set_max_size_ms(&s->ref, s->delay_ms);

	/* bind begin delay with zero's */
	m = allocb(delay_samples * 2, 0);
	memset(m->b_wptr, 0, delay_samples * 2);
	m->b_wptr += delay_samples * 2;
	ms_bufferizer_put(&s->delayed_ref, m);
}

static void configure_flow_controlled_bufferizer(ImxAiAecState *s) {
	ms_flow_controlled_bufferizer_set_samplerate(&s->ref, s->samplerate);
	ms_flow_controlled_bufferizer_set_max_size_ms(&s->ref, s->delay_ms);
	ms_flow_controlled_bufferizer_set_granularity_ms(&s->ref, (s->framesize * 1000) / s->samplerate);
}

static void imx_ai_aec_process(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	int nsamples = s->framesize;
	int nbytes = s->framesize * sizeof(int16_t);
	mblk_t *refm;
	int16_t *ref_tmp = (int16_t *)alloca(nbytes);
	int16_t *echo_tmp = (int16_t *)alloca(nbytes);
	float *ref_float = (float *)alloca(nsamples * sizeof(float));
	float *echo_float = (float *)alloca(nsamples * sizeof(float));
	float *out_float = (float *)alloca(nsamples * sizeof(float));

	if (s->bypass_mode) {
		while ((refm = ms_queue_get(f->inputs[0])) != NULL) ms_queue_put(f->outputs[0], refm);
		while ((refm = ms_queue_get(f->inputs[1])) != NULL) ms_queue_put(f->outputs[1], refm);
		return;
	}


	// 1. Get ref signal (from remote)
	if (f->inputs[0] != NULL) {
		if (s->echostarted) {
			while ((refm = ms_queue_get(f->inputs[0])) != NULL) {
				ms_bufferizer_put(&s->delayed_ref, dupmsg(refm));
				ms_flow_controlled_bufferizer_put(&s->ref, refm);
			}
		} else {
			ms_queue_flush(f->inputs[0]);
		}
	}

	// 2. Get mic signal
	ms_bufferizer_put_from_queue(&s->echo, f->inputs[1]);

	// 3. Main cycle
	while (ms_bufferizer_get_avail(&s->echo) >= (size_t)nbytes) {
		mblk_t *out_clean = allocb(nbytes, 0);
		int avail;

		if (!s->echostarted) s->echostarted = TRUE;

		// Checking if there is data in the delay buffer
		avail = (int)ms_bufferizer_get_avail(&s->delayed_ref);
		if (avail < ((s->nominal_ref_samples * 2) + nbytes)) {
			// Silence injection if the reference is lagging behind
			avail = nbytes;
			refm = allocb(nbytes, 0);
			memset(refm->b_wptr, 0, nbytes);
			refm->b_wptr += nbytes;
			ms_bufferizer_put(&s->delayed_ref, refm);
			ms_queue_put(f->outputs[0], dupmsg(refm));
			if (!s->using_zeroes) {
				ms_warning("IMX AEC: Not enough ref samples, using zeroes");
				s->using_zeroes = TRUE;
			}
		} else {
			s->using_zeroes = FALSE;
			refm = allocb(nbytes, 0);
			ms_flow_controlled_bufferizer_read(&s->ref, refm->b_wptr, nbytes);
			refm->b_wptr += nbytes;
			ms_queue_put(f->outputs[0], refm);
		}

		// Reading aligned data
		ms_bufferizer_read(&s->echo, (uint8_t *)echo_tmp, nbytes);
		ms_bufferizer_read(&s->delayed_ref, (uint8_t *)ref_tmp, nbytes);
		avail -= nbytes;
#ifdef EC_DUMP
		if (s->reffile) fwrite(ref_tmp, nbytes, 1, s->reffile);
		if (s->echofile) fwrite(echo_tmp, nbytes, 1, s->echofile);
#endif
#ifdef USE_CONVERT_NEON
		for(int i = 0; i < nsamples; i+=8) {
			s16_to_float_neon(&ref_tmp[i], &ref_float[i]);
			s16_to_float_neon(&echo_tmp[i], &echo_float[i]);
		}
#else
		//Converts int16_t to float data
		for(int i = 0; i < nsamples; i++) {
			ref_float[i] = (float)ref_tmp[i] / 32768.0f;
			echo_float[i] = (float)echo_tmp[i] / 32768.0f;
		}
#endif
#ifdef EC_DUMP
		if (s->reffile_float) fwrite(ref_float, sizeof(float) * nsamples, 1, s->reffile_float);
		if (s->echofile_float) fwrite(echo_float, sizeof(float) * nsamples, 1, s->echofile_float);
#endif
		// AI Processing (AEC + NR)
		if (imx_ai_aecnr_core_process(s->aec_core, ref_float, echo_float, out_float) != 0)
		{
			ms_error("imx_ai_aecnr_core_process failed");
			//memcpy(out_clean->b_wptr, echo_tmp, nbytes); // Fallback
			for (int i = 0; i < nsamples; i++) out_float[i] = echo_float[i];
		}
#ifdef EC_DUMP
		if (s->cleanfile_float) fwrite(out_float, sizeof(float) * nsamples, 1, s->cleanfile_float);
#endif
		//Converts float data to int16_t
		int16_t *out_ptr = (int16_t *)out_clean->b_wptr;
#ifdef USE_CONVERT_NEON
		for(int i = 0; i < nsamples; i+=8) {
			float_to_s16_neon(&out_float[i], &out_ptr[i]);
		}
#else
		for(int i = 0; i < nsamples; i++) {
			float val = out_float[i] * 32768.0f;
			if(val > 32767.0f) val = 32767.0f;
			if(val < -32768.0f) val = -32768.0f;
			out_ptr[i] = (int16_t)val;
		}
#endif
#ifdef EC_DUMP
		if (s->cleanfile) fwrite(out_ptr, nbytes, 1, s->cleanfile);
#endif
		out_clean->b_wptr += nbytes;
		ms_queue_put(f->outputs[1], out_clean);
	}
}

static void imx_ai_aec_postprocess(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;

	if (s->aec_core) {
		imx_ai_aecnr_core_close(s->aec_core);
		s->aec_core = NULL;
	}
	ms_bufferizer_flush(&s->delayed_ref);
	ms_bufferizer_flush(&s->echo);
        ms_message("%s(%d)\n",__FUNCTION__,__LINE__);
#ifdef EC_DUMP
	if (s->echofile) fclose(s->echofile);
	if (s->reffile) fclose(s->reffile);
	if (s->cleanfile) fclose(s->cleanfile);
	if (s->echofile_float) fclose(s->echofile_float);
	if (s->reffile_float) fclose(s->reffile_float);
	if (s->cleanfile_float) fclose(s->cleanfile_float);
#endif
}

static void imx_ai_aec_uninit(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;

ms_message("%s(%d)\n",__FUNCTION__,__LINE__);
	if (s->state_str) ms_free(s->state_str);
	ms_bufferizer_uninit(&s->delayed_ref);
	ms_bufferizer_uninit(&s->echo);
	ms_free(s);
}

static int imx_ai_aec_set_sr(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
ms_message("%s(%d) Try set samplerate=%u\n",__FUNCTION__,__LINE__,*(int *)arg);
#if 0
	s->samplerate = *(int *)arg;
	configure_flow_controlled_bufferizer(s);
#else
	ms_message("%s(%d) Don't support set samplerate, use only %d samplerate\n",__FUNCTION__,__LINE__,AI_AECNR_SAMPLE_RATE);
#endif
	return 0;
}

static int imx_ai_aec_set_framesize(MSFilter *f, void *arg) {
#if 0
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	s->framesize_at_8000 = *(int *)arg;
#else
	ms_message("%s(%d) Don't support set framesize, use only %d framesize\n",__FUNCTION__,__LINE__,AI_AECNR_IO_SAMPLES);
#endif
	return 0;
}

static int imx_ai_aec_set_delay(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;

	s->delay_ms = *(int *)arg;
ms_message("%s(%d) delay_ms=%d\n",__FUNCTION__,__LINE__,s->delay_ms);
	configure_flow_controlled_bufferizer(s);
	return 0;
}

static int imx_ai_aec_set_tail_length(MSFilter *f, void *arg) {
	/* Do nothing because this is not needed by the IMX AI AEC echo canceller. */
	return 0;
}

static int imx_ai_aec_set_bypass_mode(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
ms_message("%s(%d)\n",__FUNCTION__,__LINE__);
	s->bypass_mode = *(bool_t *)arg;
	ms_message("set EC bypass mode to [%i]", s->bypass_mode);
	return 0;
}
static int imx_ai_aec_get_bypass_mode(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	*(bool_t *)arg = s->bypass_mode;
	return 0;
}

static int imx_ai_aec_set_state(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	s->state_str = ms_strdup((const char *)arg);
	return 0;
}

static int imx_ai_aec_get_state(MSFilter *f, void *arg) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	*(char **)arg = s->state_str;
	return 0;
}

static MSFilterMethod imx_ai_aec_methods[] = {{MS_FILTER_SET_SAMPLE_RATE, imx_ai_aec_set_sr},
                                            {MS_ECHO_CANCELLER_SET_TAIL_LENGTH, imx_ai_aec_set_tail_length},
                                            {MS_ECHO_CANCELLER_SET_DELAY, imx_ai_aec_set_delay},
                                            {MS_ECHO_CANCELLER_SET_FRAMESIZE, imx_ai_aec_set_framesize},
                                            {MS_ECHO_CANCELLER_SET_BYPASS_MODE, imx_ai_aec_set_bypass_mode},
                                            {MS_ECHO_CANCELLER_GET_BYPASS_MODE, imx_ai_aec_get_bypass_mode},
                                            {MS_ECHO_CANCELLER_GET_STATE_STRING, imx_ai_aec_get_state},
                                            {MS_ECHO_CANCELLER_SET_STATE_STRING, imx_ai_aec_set_state},
                                            {0, 0}};

MSFilterDesc ms_imx_ai_aec_desc = {.id = MS_SPEEX_EC_ID,
                                 .name = "MSIMXAiAec",
                                 .text = N_("Echo canceller using IMX AI AEC library"),
                                 .category = MS_FILTER_OTHER,
                                 .ninputs = 2,
                                 .noutputs = 2,
                                 .init = imx_ai_aec_init,
                                 .preprocess = imx_ai_aec_preprocess,
                                 .process = imx_ai_aec_process,
                                 .postprocess = imx_ai_aec_postprocess,
                                 .uninit = imx_ai_aec_uninit,
                                 .methods = imx_ai_aec_methods};

MS_FILTER_DESC_EXPORT(ms_imx_ai_aec_desc)