#ifndef LOMB_SCARGLE_SPECTRUM
#define LOMB_SCARGLE_SPECTRUM


#ifndef PI
#define PI 3.141592653589793238462643
#endif

#ifndef EPSILON
#define EPSILON 1e-15
#endif





// Calculates the Lomb-Scargle spectrum (LSS) and Lomb-Scargle periodogram LSP (aka LS power spectrum) of a real (!) time series (an estimate for the spectrum of the generating process)
// 		Sampling times can be arbitrary, but must be given in increasing order.
//		Consider this as a (typically slight) improvement over leastSquaresPowerSpectrum() defined above.
//		In the case of evenly sampled data, the Lomb-Scargle spectrum reduces to the discrete spectrum returned by discretePowerSpectrum().
// See:
//		[Scargle (1982) - Studies in astronomical time series analysis. II-Statistical aspects of spectral analysis of unevenly spaced data, eq. (10)]
//		[Scargle (1989) - Studies in astronomical time series analysis. III. Fourier transforms, autocorrelation functions, and cross-correlation functions of unevenly spaced data, eq. (I.1)]
//		[Rehfeld (2011) - Comparison of correlation analysis techniques for irregularly sampled time series, §2.2]
//		Note that the normalization used here is for continuous-time stochastic process and is different than in the original articles by Scargle (you need to multiply power spectrum by N/T to get the original Lomb-Scargle normalization).
//		To get the power spectrum normalization used for stochastic processes, multiply the obtained spectrum by the (total signal count N)/(total time span T) (for discrete sum version, same as original Lomb-Scargle version).
//		Also, df is in frequency units. To get angular frequencies, replace df by df*2PI.
//		Also note that here the spectrum is the complex conjugate of the spectrum defined by Lomb & Scangler (for consistency with the conventional Fourier transform).
// Returns false on error
template<class TYPE>
bool LombScargleSpectrum(	const vector<TYPE> 		&signal,
							const vector<double> 	&times,
							long 					start, 			// start index in data series (ignore anything prior)
							long 					end, 			// end index in data series (ignore anything after)
							double 					maxFrequency, 	// (input) limit the considered spectrum band. Set to non-positive for "average" Nyquist–Shannon bound.
							vector<TYPE>* 			spectrumReal,	// (output) either pre-allocated, or NULL to skip this part
							vector<TYPE>* 			spectrumImag,	// (output) either pre-allocated, or NULL to skip this part
							vector<TYPE>* 			powerSpectrum,	// (output) either pre-allocated, or NULL to skip this part
							double 					&df){ 			//(output) fundamental frequency
	if((times.size()<=end) || (signal.size() <=end)) return false;
	double T = times[end] - times[start]; //data time span
	if(T<=0) return false;
	long n,m, N = end - start + 1;
	if(N<2) return false;
	double S, C, phi, tau, f, mindt, maxdt, meandt=T/(N-1);
	TYPE FC, FS, FRe, FIm, Xmean, zero=signal[start]*0;
	
	for(n=start, mindt=maxdt=times[start+1]-times[start]; n<end; ++n){
		mindt = min(mindt, times[n+1]-times[n]);
		maxdt = max(maxdt, times[n+1]-times[n]);
	}
	bool uniform = (maxdt-mindt < 0.00001*meandt);
	df = 1.0/(N*meandt);
	long maxMode = (maxFrequency>0 ? min((long)(maxFrequency/df), N/2 - 1) : N/2 - 1); //upper limit given by Nyquist frequency
	
	if(powerSpectrum!=NULL){
		(*powerSpectrum).clear();
		(*powerSpectrum).reserve(maxMode+1);
	}
	if(spectrumReal!=NULL){
		(*spectrumReal).clear();
		(*spectrumReal).reserve(maxMode+1);
	}
	if(spectrumImag!=NULL){
		(*spectrumImag).clear();
		(*spectrumImag).reserve(maxMode+1);
	}
	
	for(m=0; m<=maxMode; ++m){
		if(uniform){
			//use classical DFT formula
			for(n=0, FRe=FIm=zero; n<N; ++n){
				FRe = FRe + signal[start+n]*cos(-(2.0*PI*m*n)/N);
				FIm = FIm + signal[start+n]*sin(-(2.0*PI*m*n)/N);
			}
			FRe = FRe/N;
			FIm = FIm/N;
			
		}else{
		
			//resort to Lomb-Scargle estimator
			if(m==0){
				for(n=start, Xmean=zero; n<=end; ++n){
					Xmean = Xmean + signal[n];
				}
				Xmean = Xmean/N;
				FRe = Xmean;
				FIm = zero;

			}else{
				f = m*df;
				
				//calculate time shift
				for(n=start, S=C=0; n<=end; ++n){
					C += cos(4*PI*f*times[n]);
					S += sin(4*PI*f*times[n]);
				}
				tau = atan(S/C)/(4*PI*f);
				
				//calculate spectrum
				for(n=start, S=C=0, FC=FS=zero; n<=end; ++n){
					C += SQ(cos(2*PI*f*(times[n]-tau)));
					S += SQ(sin(2*PI*f*(times[n]-tau)));
					FC = FC + (signal[n]-Xmean)*cos(-2*PI*f*(times[n]-tau));
					FS = FS + (signal[n]-Xmean)*sin(-2*PI*f*(times[n]-tau));
				}
				FC = FC/sqrt(2*C*N);
				FS = FS/sqrt(2*S*N);
				
				//shift phase of complex spectrum as required by Lomb-Scangler formula
				phi = 2*PI*f*(times[start]-tau);
				FRe = FC*cos(phi) - FS*sin(phi);
				FIm = FC*sin(phi) + FS*cos(phi);
			}
		}
		
		if(powerSpectrum!=NULL) (*powerSpectrum).push_back(T*(FRe*FRe + FIm*FIm));
		if(spectrumReal!=NULL) (*spectrumReal).push_back(sqrt(T)*FRe);
		if(spectrumImag!=NULL) (*spectrumImag).push_back(sqrt(T)*FIm);
	}
	return true;
}





template<class TYPE>
bool averagedLombScargleSpectrum(	const vector<TYPE> 		&signal,
									const vector<double> 	&times,
									long 					start, 				// start index in data series (ignore anything prior)
									long 					end, 				// end index in data series (ignore anything after)
									long 					segmentsAim,		// (input) number of segments to split the time series into (if possible)
									long					minNumberOfSegments,// (input) if segmentsAchieved ends up lower than this, abort and return false
									long					maxPointsToOmit,	// (input) maximum number of data points allowed to be omitted in order to make time series points divisible by segmentsAim. If zero, it might happen that only one segment is achieved.
									double 					maxFrequency, 		// (input) limit the considered spectrum band. Set to non-positive for "average" Nyquist–Shannon bound.
									long					&segmentsAchieved,	// (output) number of segments the time series was actually split into
									vector<TYPE> 			&powerSpectrum,		// (output)
									double 					&df,				// (output)
									double					&dtPerSegment,		// (output)
									long					&NperSegment){
	if(end<start || start<0) return false;
	const long N = end-start+1;
	if(N<2) return false;
	long Nreduced = N;
	segmentsAchieved = max(1l,segmentsAim);
	while(Nreduced % segmentsAchieved != 0){
		if(Nreduced>N-maxPointsToOmit){ --Nreduced; }
		else if(segmentsAchieved>minNumberOfSegments){ --segmentsAchieved; Nreduced=N; }
		else{ return false; }
	}
	if(segmentsAchieved<minNumberOfSegments) return false;
	NperSegment = Nreduced/segmentsAchieved;
	const double Treduced = times[start+Nreduced-1] - times[start]; //evaluated time span
	if(Treduced<=0) return false;
	dtPerSegment = segmentsAchieved*Treduced/(Nreduced-1.0); // dt for time points of each segment
	const double TperSegment = dtPerSegment*(NperSegment-1);
	df = 1.0/(NperSegment*dtPerSegment);
	const long maxMode = (maxFrequency>0 ? min((long)(maxFrequency/df), NperSegment/2 - 1) : NperSegment/2 - 1); //upper limit given by Nyquist frequency
		
	double S, C, phi, tau, f;
	TYPE FC, FS, FRe, FIm, Xmean, zero=signal[start]*0;
	long segmentStart, segmentEnd, segmentStep, n;
	
	powerSpectrum.clear();
	powerSpectrum.assign(maxMode+1,zero);

	for(long s=0; s<segmentsAchieved; ++s){
		for(long m=0; m<=maxMode; ++m){
			segmentStart = start + s;
			segmentStep = segmentsAchieved;
			segmentEnd = segmentStart + segmentStep*(NperSegment-1);
			if(m==0){
				for(n=segmentStart, Xmean=zero; n<=segmentEnd; n+=segmentStep){
					Xmean = Xmean + signal[n];
				}
				Xmean = Xmean/NperSegment;
				FRe = Xmean;
				FIm = zero;
	
			}else{
				f = m*df;
				
				//calculate time shift
				for(n=segmentStart, S=C=0; n<=segmentEnd; n+=segmentStep){
					C += cos(4*PI*f*times[n]);
					S += sin(4*PI*f*times[n]);
				}
				tau = atan(S/C)/(4*PI*f);
				
				//calculate spectrum
				for(n=segmentStart, S=C=0, FC=FS=zero; n<=segmentEnd; n+=segmentStep){
					C += SQ(cos(2*PI*f*(times[n]-tau)));
					S += SQ(sin(2*PI*f*(times[n]-tau)));
					FC = FC + (signal[n]-Xmean)*cos(-2*PI*f*(times[n]-tau));
					FS = FS + (signal[n]-Xmean)*sin(-2*PI*f*(times[n]-tau));
				}
				FC = FC/sqrt(2*C*NperSegment);
				FS = FS/sqrt(2*S*NperSegment);
				
				//shift phase of complex spectrum as required by Lomb-Scangler formula
				phi = 2*PI*f*(times[start]-tau);
				FRe = FC*cos(phi) - FS*sin(phi);
				FIm = FC*sin(phi) + FS*cos(phi);
			}
			
			powerSpectrum[m] += (TperSegment*(FRe*FRe + FIm*FIm));
		}
	}
	for(long m=0; m<powerSpectrum.size(); ++m){
		powerSpectrum[m] /= segmentsAchieved;
	}
	return true;
}








// returns the expected Lomb-Scargle periodogram of a time series obtained by iid normal values of variance sigma^2 (the discrete-time version of white noise)
// the scaling is consistent with LombScargleSpectrum(..)
double LombScargleSpectrum_of_normal_errors(double sigma, double T, long N){
	return SQ(sigma)*T/N;
}


// only accurate for large N
double LombScargleSpectrum_of_OU_at_zero_mode(double SD, double lambda, double T, long N){
	double dt = T/(N-1);
	if(lambda*dt<EPSILON) return 2*SQ(SD)/lambda;
	return T*SQ(SD)/N * (1+exp(-lambda*dt))/(1-exp(-lambda*dt));
}



#endif