// --------------------------------------------------------------
#include <string.h>
#include <malloc.h>
#include <math.h>
#include "equalizer.h"

// --------------------------------------------------------------
// Helpers
// --------------------------------------------------------------
void CoeffComplexUnitAt(cpx_t *pW, uint32_t N, uint32_t dly)
{
	if (dly >= N)
		dly = 0;

	memset(pW, 0, N*sizeof(cpx_t));
	pW[dly].real = 1.0;
	pW[dly].imag = 0.0;
}

void CoeffComplexConv(cpx_t *pSrc, cpx_t *pSrcDst, uint32_t N, uint32_t dly)
{
	uint32_t i, j, N2;
	cpx_t *pX, *pW, t, energy = {0};
	radio_float_t energy_correct;

	N2 = N+dly;

	pX = (cpx_t*)malloc(N2*sizeof(cpx_t));
	memset(pX, 0, N2*sizeof(cpx_t));
	pW = (cpx_t*)malloc(N2*sizeof(cpx_t));
	memset(pW, 0, N2*sizeof(cpx_t));

	CpxCopy(pSrc, pX, N);

	for(j=0; j < N; j++)
		energy = CpxAddS(energy, pSrcDst[j]);

	energy_correct = (radio_float_t)(1.0/CpxMagS(energy));

	for(i=0; i < N2; i++)
	{
		for(j=0; j < N; j++)
		{
			t = CpxMulS(pX[N-j+i-1], pSrcDst[j]);
			pW[i] = CpxScaleRealS(CpxAddS(pW[i], t), energy_correct);
		}
	}
	CpxCopy(&pW[N/2], pSrcDst, N);
}

// --------------------------------------------------------------
// Equalizer
// --------------------------------------------------------------
void EQComplexInit(eq_cpx_t *pObj, uint32_t N)
{
	pObj->N = N;
	pObj->pX = (cpx_t*)malloc(pObj->N*sizeof(cpx_t));
	pObj->pW = (cpx_t*)malloc(pObj->N*sizeof(cpx_t));
	pObj->pW_conj = (cpx_t*)malloc(pObj->N*sizeof(cpx_t));
	memset(pObj->pX, 0, pObj->N*sizeof(cpx_t));
	memset(pObj->pW, 0, pObj->N*sizeof(cpx_t));
	memset(pObj->pW_conj, 0, pObj->N*sizeof(cpx_t));
}

void EQComplexFree(eq_cpx_t *pObj)
{
	free(pObj->pX);
	free(pObj->pW);
	free(pObj->pW_conj);
}

cpx_t EQComplexProcess(eq_cpx_t *pObj, cpx_t x)
{
	uint32_t i;
	cpx_t y;

	for(i=pObj->N-1; (int)i > 0; i--)
	{
		pObj->pX[i] = pObj->pX[i-1];
	}

	// fill M top positions with M new data
	pObj->pX[i] = x;

	y = CpxMulAdd(Cpx(0, 0), pObj->pX, pObj->pW_conj, pObj->N);

	return y;
}

// --------------------------------------------------------------
// CMA
// --------------------------------------------------------------
void CMAInit(cma_t *pObj, uint32_t ntaps)
{
	pObj->N = ntaps;
	pObj->Px = 0;
	pObj->Px_last = 0;

	EQComplexInit(&pObj->eq, pObj->N);
}

void CMAFree(cma_t *pObj)
{
	EQComplexFree(&pObj->eq);
}

cpx_t CMAProcess(cma_t *pObj, cpx_t x)
{
	cpx_t y;
	radio_float_t x2, x1;

	y = EQComplexProcess(&pObj->eq, x);
	x1 = CpxMagS(x);
	x2 = CpxMagS(pObj->eq.pX[pObj->eq.N-1]);

	pObj->Px = dmax(0, pObj->Px + (x1*x1 - pObj->Px_last));
	pObj->Px_last = x2*x2;

	return y;
}

// Proakis: page 698
void CMATrainGodard(cma_t *pObj, cpx_t sym_s_eq, radio_float_t mag_h, radio_float_t R2, radio_float_t mu)
{
	cpx_t d;
	cpx_t e;
	radio_float_t mag_s, t;

	mag_s = CpxMagS(sym_s_eq);
	
	t = (radio_float_t)((mag_h + R2*mag_s - mag_s*mag_s*mag_s)/(1E-4+mag_h));
	d = CpxScaleRealS(sym_s_eq, t);

	e = CpxSubS(d, sym_s_eq);

	LMSUpdateWeigths(&pObj->eq, e, mu/((radio_float_t)1+pObj->Px));
}

// "Sliced Multi-modulus Blind Equalization Algorithm",
// ETRI Journal, Volume 27, Number 3, June 2005,
// Shafayat Abrar and Roy A. Axford Jr.
void MMATrain(cma_t *pObj, cpx_t sym_s_eq, cpx_t R, radio_float_t mu)
{
	cpx_t e;

	e.real = sym_s_eq.real*(R.real - sym_s_eq.real*sym_s_eq.real);
	e.imag = sym_s_eq.imag*(R.imag - sym_s_eq.imag*sym_s_eq.imag);

	LMSUpdateWeigths(&pObj->eq, e, mu/((radio_float_t)1+pObj->Px));
}

// "Sliced Multi-modulus Blind Equalization Algorithm",
// ETRI Journal, Volume 27, Number 3, June 2005,
// Shafayat Abrar and Roy A. Axford Jr.
void SMMATrain(cma_t *pObj, cpx_t sym_s_eq, radio_float_t mag_h, cpx_t R, radio_float_t mu)
{
	cpx_t e;

	e.real = sym_s_eq.real*(mag_h*R.real - sym_s_eq.real*sym_s_eq.real);
	e.imag = sym_s_eq.imag*(mag_h*R.imag - sym_s_eq.imag*sym_s_eq.imag);

	LMSUpdateWeigths(&pObj->eq, e, mu/((radio_float_t)1+pObj->Px));
}

// --------------------------------------------------------------
// DFE
// --------------------------------------------------------------
void DFEComplexInit(dfe_cpx_t *pObj, uint32_t K)
{
	pObj->N = 2*K+1;
	pObj->N_ff = K+1;
	pObj->N_fb = K;

	EQComplexInit(&pObj->eq, pObj->N);

}

void DFEComplexFree(dfe_cpx_t *pObj)
{
	EQComplexFree(&pObj->eq);
}

cpx_t DFEComplexProcess(dfe_cpx_t *pObj, cpx_t xs, cpx_t xh)
{
	int32_t i;
	cpx_t y, *pW_conj, *pX;

	pX = pObj->eq.pX;
	pW_conj = pObj->eq.pW_conj;

	// Shift Feedback part
	for(i=pObj->N-1; i > (int32_t)pObj->N_ff; i--)
	{
		pX[i] = pX[i-1];
	}

	// put last x into buffer
	pX[i] = xh;

	// Shift Feed forward part
	for(i=pObj->N_ff-1; i > 0; i--)
	{
		pX[i] = pX[i-1];
	}

	// put current x into buffer
	pX[i] = xs;

	y = CpxMulAdd(Cpx(0, 0), pX, pW_conj, pObj->N);

	return y;

}

void DFEAdaptLMS(dfe_cpx_t *pObj, cpx_t e, radio_float_t mu)
{
	LMSUpdateWeigths(&pObj->eq, e, mu);
}

// --------------------------------------------------------------
// LMS
// --------------------------------------------------------------
void LMSUpdateWeigths(eq_cpx_t *pObj, cpx_t e, radio_float_t mu)
{
	uint32_t i;

	e = CpxScaleRealS(CpxConjS(e), mu);

	CpxScale(pObj->pX, e, pObj->pW_conj, pObj->N);
	CpxAdd(pObj->pW_conj, pObj->pW, pObj->N);

	for(i=0; i < pObj->N; i++)
	{
		pObj->pW_conj[i] = CpxConjS(pObj->pW[i]);
	}
}

// --------------------------------------------------------------
// RLS
// --------------------------------------------------------------
void RLSInit(rls_t *pObj, uint32_t N, uint32_t M)
{
	uint32_t i, j;
	radio_float_t *pR;

	pObj->w = (radio_float_t*)malloc(N*sizeof(radio_float_t));
	memset(pObj->w, 0, N*sizeof(radio_float_t));
	pObj->w[0] = 1.0f;

	pObj->pW = (cpx_t*)malloc(N*sizeof(cpx_t));
	memset(pObj->pW, 0, N*sizeof(cpx_t));

	pObj->z = (radio_float_t*)malloc(N*sizeof(radio_float_t));
	memset(pObj->z, 0, N*sizeof(radio_float_t));
	
	pObj->x = (radio_float_t*)malloc(N*sizeof(radio_float_t));
	memset(pObj->x, 0, N*sizeof(radio_float_t));
	
	pObj->pX = (cpx_t*)malloc(N*M*sizeof(cpx_t));
	memset(pObj->pX, 0, N*M*sizeof(cpx_t));

	pObj->R = (radio_float_t*)malloc(N*N*sizeof(radio_float_t));
	memset(pObj->R, 0, N*N*sizeof(radio_float_t));

	pObj->T = (radio_float_t*)malloc(N*N*sizeof(radio_float_t));
	memset(pObj->T, 0, N*N*sizeof(radio_float_t));

	pObj->T2 = (radio_float_t*)malloc(N*N*sizeof(radio_float_t));
	memset(pObj->T2, 0, N*N*sizeof(radio_float_t));

	pR = pObj->R;
	for(i=0; i < N; i++)
		for(j=0; j < N; j++)
			pR[i+j*N] = (radio_float_t)1E6*(i==j);

	pObj->N = N;
	pObj->M = M;
}

void RLSFree(rls_t *pObj)
{
	if(pObj->w)
		free(pObj->w);

	if(pObj->z)
		free(pObj->z);

	if(pObj->x)
		free(pObj->x);

	if(pObj->R)
		free(pObj->R);

	if(pObj->T)
		free(pObj->T);

	if(pObj->T2)
		free(pObj->T2);
}
/*
void RLSProcess(rls_t *pObj, radio_float_t rho, radio_float_t mu, radio_float_t x, radio_float_t d, radio_float_t *pE, radio_float_t *pDD)
{
	int i, j, k, N;
	radio_float_t *pR, *pX, *pZ, *pT, *pT2, v, t, rho_i;

	N = pObj->N;

	FIR(pObj->w, pObj->x, N, &x, pDD, 1);

	*pE = d - *pDD;

	// Whitening
	pR = pObj->R;
	pX = pObj->x;
	for(i=0; i < N; i++)
	{
		t = 0;
		for(j=0; j < N; j++)
			t += pR[j*N+i] * pX[j];

		pObj->z[i] = t;
	}

	// Berechnunng der Normierungskonstante
	v = rho;
	for(i=0; i < N; i++)
		v += pObj->x[i] * pObj->z[i];

	v = (radio_float_t)1.0/v;

	// Normierung
	for(i=0; i < N; i++)
		pObj->z[i] *= v;


	// Filter update
	for(i=0; i < N; i++)
		pObj->w[i] += (mu * (*pE)*pObj->z[i]);


	// Matrix update
	pZ = pObj->z;
	pX = pObj->x;
	pT = pObj->T;
	for(i=0; i < N; i++)
	{
		for(j=0; j < N; j++)
			pT[j*N+i] = pZ[i] * pX[j];
	}

	pT = pObj->T;
	pT2 = pObj->T2;
	pR = pObj->R;

	for (i=0;i<N;i++)
	{
		for (j=0;j<N;j++)
		{
			pT2[j*N+i] = 0;
			for (k=0;k<N;k++)
			{
				pT2[j*N+i] += pT[k*N+i] * pR[j*N+k];
			}
		}
	}

	for(i=0; i < N*N; i++)
		pR[i] -= pT2[i];

	rho_i = (radio_float_t)1.0/rho;
	for(i=0; i < N*N; i++)
		pR[i] *= rho_i;
}
*/
cpx_t RLSProcess(rls_t *pObj, radio_float_t i, radio_float_t q)
{
	uint32_t j;
	cpx_t y, t;

	y.real = 0;
	y.imag = 0;
	for(j=pObj->N-1; j != 0; j--)
	{
		pObj->pX[j] = pObj->pX[j-1];
	}
	pObj->pX[j].real = i;
	pObj->pX[j].imag = q;

	for(j=0; j < pObj->N; j++)
	{
		pObj->pW[j].real = pObj->w[j];
		pObj->pW[j].imag = pObj->w[j];
	}

///	for(j=0; j < pObj->N; j++)
//	{
//		t = CpxMuls(pObj->pX[M*j], CpxConjs(pObj->pW[j]));
//		y = CpxAdds(y, t);
//	}

	for(j=0; j < pObj->N; j++)
	{
		t.real = pObj->pX[j].real * pObj->pW[j].real;
		t.imag = pObj->pX[j].imag * pObj->pW[j].imag;
		
		y = CpxAddS(y, t);
	}

	for(j=0; j < pObj->N; j++)
	{
		pObj->x[j] = pObj->pX[j].real;
	}

	return y;

}

radio_float_t RLSUpdate(rls_t *pObj, radio_float_t rho, radio_float_t mu, radio_float_t y, radio_float_t d)
{
	int i, j, k, N;
	radio_float_t *pR, *pX, *pZ, *pT, *pT2, v, t, rho_i, e;

	N = pObj->N;

	e = d - y;

	pR = pObj->R;
	pX = pObj->x;

//	for(j=0; j < pObj->N; j++)
//	{
//		pObj->x[j] = pObj->pX[j*pObj->M].real;
//	}

	// Whitening
	for(i=0; i < N; i++)
	{
		t = 0;
		for(j=0; j < N; j++)
			t += pR[j*N+i] * pX[j];

		pObj->z[i] = t;
	}

	// Berechnunng der Normierungskonstante
	v = rho;
	for(i=0; i < N; i++)
		v += pObj->x[i] * pObj->z[i];

	v = (radio_float_t)1.0/v;

	// Normierung
	for(i=0; i < N; i++)
		pObj->z[i] *= v;


	// Filter update
	for(i=0; i < N; i++)
		pObj->w[i] += (mu * e * pObj->z[i]);


	// Matrix update
	pZ = pObj->z;
	pX = pObj->x;
	pT = pObj->T;
	for(i=0; i < N; i++)
	{
		for(j=0; j < N; j++)
			pT[j*N+i] = pZ[i] * pX[j];
	}

	pT = pObj->T;
	pT2 = pObj->T2;
	pR = pObj->R;

	for (i=0;i<N;i++)
	{
		for (j=0;j<N;j++)
		{
			pT2[j*N+i] = 0;
			for (k=0;k<N;k++)
			{
				pT2[j*N+i] += pT[k*N+i] * pR[j*N+k];
			}
		}
	}

	for(i=0; i < N*N; i++)
		pR[i] -= pT2[i];

	rho_i = (radio_float_t)1.0/rho;
	for(i=0; i < N*N; i++)
		pR[i] *= rho_i;

	return e;
}

// --------------------------------------------------------------