/***************************************************************************
		Copyright (C) 2002-2015 Kentaro Kitagawa
		                   kitag@issp.u-tokyo.ac.jp
		
		This program is free software; you can redistribute it and/or
		modify it under the terms of the GNU Library General Public
		License as published by the Free Software Foundation; either
		version 2 of the License, or (at your option) any later version.
		
		You should have received a copy of the GNU Library General 
		Public License and a list of authors along with this program; 
		see the files COPYING and AUTHORS.
***************************************************************************/
#include "ar.h"

template <class Context>
YuleWalkerCousin<Context>::YuleWalkerCousin(tfuncIC ic) : SpectrumSolver(), m_funcARIC(ic) {	
}

template <class Context>
void
YuleWalkerCousin<Context>::genSpectrum(const std::vector<std::complex<double> >& memin, std::vector<std::complex<double> >& memout,
	int t0, double tol, FFT::twindowfunc windowfunc, double windowlength) {
	int t = memin.size();
	int n = memout.size();
	
	int wpoints = lrint(numberOfNoises(memin)); //# of fittable data in freq. domain.
	wpoints = std::min(std::max(wpoints, t/100 + 1), t);
//	fprintf(stderr, "# of data points = %d\n", wpoints);
	
	int taps_div = t - 1;
	taps_div = std::max(taps_div / 10, 1);
	shared_ptr<Context> context(new Context);
	context->a.resize(t);
	std::fill(context->a.begin(), context->a.end(), std::complex<double>(0.0));
	context->a[0] = 1.0;
	context->t = t;
	context->sigma2 = 0.0;
	for(unsigned int i = 0; i < t; i++) {
		context->sigma2 += std::norm(memin[i]) / (double)t;
	}
	context->p = 0;
	first(memin, context);
	unsigned int taps = 0;
	double ic = 1e99;
	for(unsigned int p = 0; p < t - 2; p++) {
		if(p % taps_div == 0) {
            m_contexts.push_back(std::make_shared<Context>(*context));
			if(taps + taps_div/2 < p)
				break;
		}
		step(context);
		context->p++;
		if(context->sigma2 < 0)
			break;
		double new_ic = arIC(context->sigma2, context->p, wpoints);
		if(new_ic < ic) {
			taps = p + 1;
			ic = new_ic;
		}
	}
	taps = std::min((int)lrint(windowlength * taps), (int)t - 2);
	int cidx = std::min((int)taps / taps_div, (int)m_contexts.size() - 1);
	context = m_contexts[cidx];
	m_contexts.clear();
	for(unsigned int p = cidx * taps_div; p < taps; p++) {
		assert(context->p == p);
		step(context);
		context->p++;		
	}
	dbgPrint(formatString("MEM/AR: t=%d, taps=%d, IC_min=%g, IC=%g\n", t, taps, ic, m_funcARIC(context->sigma2, context->p, t)));

	std::vector<std::complex<double> > zfbuf(n), fftbuf(n);
	std::fill(zfbuf.begin(), zfbuf.end(), std::complex<double>(0.0));
	for(int i = 0; i < taps + 1; i++) {
		zfbuf[i] = context->a[i];
	}
	m_fftN->exec(zfbuf, fftbuf);

	//Power spectrum density.
	std::vector<double> psd(n);
	for(int i = 0; i < n; i++) {
		double z = t * context->sigma2 / (std::norm(fftbuf[i]));
		psd[i] = std::max(z, 0.0);
	}
	//Least-Square Phase Estimation.
	double coeff = lspe(memin, t0, psd, memout, tol, true, windowfunc);

	// Derivative of denominator.
	std::fill(zfbuf.begin(), zfbuf.end(), std::complex<double>(0.0));
	for(int i = 0; i < taps + 1; i++) {
		zfbuf[i] = context->a[i] * (double)((i >= n/2) ? (i - n) : i) * std::complex<double>(0, -1);
	}
	std::vector<std::complex<double> > fftbuf2(n);
	m_fftN->exec(zfbuf, fftbuf2);
	//Peak detection. Sub-resolution detection for smooth curves.
	double dy_old = std::real(fftbuf2[0] * std::conj(fftbuf[0]));
	for(int ip = 0; ip < n; ip++) {
		int in = (ip + 1) % n;
		double dy = std::real(fftbuf2[in] * std::conj(fftbuf[in]));
		if((dy_old < 0) && (dy > 0)) {
			double dx = - dy_old / (dy - dy_old);
			if((dx >= 0) && (dx <= 1.0)) {
				std::complex<double> z = 0.0, xn = 1.0,
					x = std::polar(1.0, -2 * M_PI * (dx + ip) / (double)n);
				for(int i = 0; i < taps + 1; i++) {
					z += context->a[i] * xn;
					xn *= x;
				}
				double r = coeff * sqrt(std::max(t * context->sigma2 / std::norm(z), 0.0));
				m_peaks.push_back(std::pair<double, double>(r, dx + ip));
			}
		}
		dy_old = dy;
	}
}

template class YuleWalkerCousin<ARContext>;
template class YuleWalkerCousin<MEMBurgContext>;

void
MEMBurg::first(
	const std::vector<std::complex<double> >& memin, const shared_ptr<MEMBurgContext> &context) {
	unsigned int t = context->t;
	context->epsilon.resize(t);
	context->eta.resize(t);
	std::copy(memin.begin(), memin.end(), context->epsilon.begin());
	std::copy(memin.begin(), memin.end(), context->eta.begin());
}
void
MEMBurg::step(const shared_ptr<MEMBurgContext> &context) {
	unsigned int t = context->t;
	unsigned int p = context->p;
	std::complex<double> x = 0.0, y = 0.0;
	for(unsigned int i = p + 1; i < t; i++) {
		x += context->epsilon[i] * std::conj(context->eta[i-1]);
		y += std::norm(context->epsilon[i]) + std::norm(context->eta[i-1]);
	}
	std::complex<double> alpha = x / y;
	alpha *= -2;
	for(unsigned int i = t - 1; i >= p + 1; i--) {
		context->eta[i] = context->eta[i-1] + std::conj(alpha) * context->epsilon[i];
		context->epsilon[i] += alpha * context->eta[i-1];
	}
	std::vector<std::complex<double> > a_next(p + 2);
	for(unsigned int i = 0; i < p + 2; i++) {
		a_next[i] = context->a[i] + alpha * std::conj(context->a[p + 1 - i]);
	}
	std::copy(a_next.begin(), a_next.end(), context->a.begin());
	context->sigma2 *= 1 - std::norm(alpha);
}
void
YuleWalkerAR::first(
	const std::vector<std::complex<double> >& memin, const shared_ptr<ARContext> &context) {
	unsigned int t = context->t;
	m_rx.resize(t);
	autoCorrelation(memin, m_rx);
}
void
YuleWalkerAR::step(const shared_ptr<ARContext> &context) {
	unsigned int p = context->p;
	std::complex<double> delta = 0.0;
	for(unsigned int i = 0; i < p + 1; i++) {
		delta += context->a[i] * m_rx[p + 1 - i];
	}
	std::vector<std::complex<double> > a_next(p + 2);
	for(unsigned int i = 0; i < p + 2; i++) {
		a_next[i] = context->a[i] - delta/context->sigma2 * std::conj(context->a[p + 1 - i]);
	}
	std::copy(a_next.begin(), a_next.end(), context->a.begin());
	context->sigma2 += - std::norm(delta) / context->sigma2;
}

