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

#include "correlator.h"

bool correlator_init(correlator_ctx_t *ctx, const float complex *search_seq, size_t nsymbols)
{
	// copy arguments and initialize variables
	ctx->search_sequence_len = nsymbols;

	ctx->input_history = NULL;
	ctx->search_sequence = NULL;

	ctx->history_ptr = 0;

	// Determine the smallest power of two greater than nsymbols
	ctx->buffer_size = 1;
	while(nsymbols > 0) {
		ctx->buffer_size <<= 1;
		nsymbols >>= 1;
	}

	// If buffer_size = 0b1000, then buffer_size_mask=0b0111
	ctx->buffer_size_mask = ctx->buffer_size - 1;

	// Allocate the buffers
	ctx->input_history = malloc(ctx->buffer_size * sizeof(ctx->input_history[0]));
	if(!ctx->input_history) {
		fprintf(stderr, "correlator: malloc() failed!\n");
		goto fail;
	}

	ctx->search_sequence = malloc(ctx->buffer_size * sizeof(ctx->search_sequence[0]));
	if(!ctx->search_sequence) {
		fprintf(stderr, "correlator: malloc() failed!\n");
		goto fail;
	}

	// Clear the history buffer
	memset(ctx->input_history, 0, ctx->buffer_size * sizeof(ctx->input_history[0]));

	// Prepare the search sequence.
	// The complex values are conjugated in the process.
	for(size_t i = 0; i < ctx->search_sequence_len; i++) {
		ctx->search_sequence[i] = conjf(search_seq[i]);
	}

	// Success!
	return true;

fail:
	correlator_free(ctx);
	return false;
}


void correlator_free(correlator_ctx_t *ctx)
{
	if(ctx->input_history) {
		free(ctx->input_history);
		ctx->input_history = NULL;
	}

	if(ctx->search_sequence) {
		free(ctx->search_sequence);
		ctx->input_history = NULL;
	}

	ctx->buffer_size = 0;
	ctx->search_sequence_len = 0;
}


float complex correlator_step(correlator_ctx_t *ctx, float complex sample)
{
	float complex result = 0;

	// increment and wrap the history pointer
	ctx->history_ptr++;
	ctx->history_ptr &= ctx->buffer_size_mask;

	// copy the input sample to the history
	ctx->input_history[ctx->history_ptr] = sample;

	size_t n = (ctx->history_ptr - ctx->search_sequence_len + 1) & ctx->buffer_size_mask;

	// calculate the correlation
	for(size_t m = 0; m < ctx->search_sequence_len; m++) {
		size_t nm = (n + m) & ctx->buffer_size_mask;

		result += ctx->search_sequence[m] * ctx->input_history[nm];
	}

	return result;
}


void correlator_get_input_history(correlator_ctx_t *ctx, float complex *history)
{
	size_t n = (ctx->history_ptr - ctx->search_sequence_len + 1) & ctx->buffer_size_mask;

	for(size_t m = 0; m < ctx->search_sequence_len; m++) {
		size_t nm = (n + m) & ctx->buffer_size_mask;
		history[m] = ctx->input_history[nm];
	}
}


const float complex* correlator_get_conj_search_sequence(correlator_ctx_t *ctx)
{
	return ctx->search_sequence;
}


float correlator_get_mean_phase_deviation(correlator_ctx_t *ctx)
{
	float complex mean_unrotated_symbol = 0;

	size_t n = (ctx->history_ptr - ctx->search_sequence_len + 1) & ctx->buffer_size_mask;

	for(size_t m = 0; m < ctx->search_sequence_len; m++) {
		size_t nm = (n + m) & ctx->buffer_size_mask;
		mean_unrotated_symbol += ctx->search_sequence[m] * ctx->input_history[nm];
	}

	// no need to divide by ctx->search_sequence_len because division does not change the angle
	return cargf(mean_unrotated_symbol);
}


/* Louise & Regiannini frequency estimator */
float correlator_get_mean_frequency_deviation(correlator_ctx_t *ctx)
{
	float complex z[ctx->search_sequence_len];

	size_t n = (ctx->history_ptr - ctx->search_sequence_len + 1) & ctx->buffer_size_mask;

	// remove the influence of the data from the received symbols
	for(size_t m = 0; m < ctx->search_sequence_len; m++) {
		size_t nm = (n + m) & ctx->buffer_size_mask;
		z[m] = ctx->search_sequence[m] * ctx->input_history[nm];
	}


	// Calculate averaged phase increments for <N> sub-sequences
	size_t N = ctx->search_sequence_len/2;
	float complex R_kappa[N];

	for(size_t kappa = 0; kappa < N; kappa++) {
		for(size_t k = 0; k < ctx->search_sequence_len - kappa; k++) {
			R_kappa[kappa] += z[k + kappa] * conj(z[k]);
		}

		R_kappa[kappa] /= ctx->search_sequence_len - kappa;
	}

	// Calculate phase estimate (in radians/sample)
	float complex sum_R_kappa = 0;

	for(size_t kappa = 0; kappa < N; kappa++) {
		sum_R_kappa += R_kappa[kappa];
	}

	float arg = cargf(sum_R_kappa);

	return arg / ((float)M_PI * (1 + N));
}