#	nextn - length selection for fft
#
#	Given an integer argument, "nextn" selects a slightly
#	larger n for which the computation of the FFT will be fast.
#
#	The strategy is as follows:  First find the first power of
#	2 exceeding the given n.  If the new n exceeds the old by
#	too much (see code) it is adjusted down by one of
#		9/16,  5/8,  or  3/4
#

nextn <-
function(n)
{
	nn <- 1
	while(nn < n)
		nn <- nn * 2

	if( nn>512 ) {
		ratio <- n/nn
		if( ratio<0.56 )
			nn <- (nn/16)*9
		else if( ratio<0.625 )
			nn <- (nn/8)*5
		else if( ratio<0.75 )
			nn <- (nn/4)*3
	}
	return(nn)
}


# 	Filters are vectors with an attached attribute "Lag".
#	Filtering with a(1), ... , a(L) takes place according
#	to the formula
#
#			  L
#		y(t)  =  SUM a(u) x(t-u)
#			 u=1
#
#	This introduces L NA values at the start of the series.
#	If the filter has a "Lag" attribute, the result is lagged
#	by this amount.

"Lag<-" <-
function(filter,lag)
{
	attr(filter,"Lag") <- lag
	return(filter)
}


Lag <-
function(filter)
	attr(filter,"Lag")


filter <-
function(x,filter) {

	nx <- length(x)
	nf <- length(filter)
	n <- nextn(nx)

	u <- rep(0,n)
	v <- rep(0,n)
	u[1:nx] <- x
	v[1:nf] <- f

	w <- (Re(fft(fft(u)*fft(v),inv=T))/n)[1:nx]
	w[1:nf] <- NA
	tsp(w) <- tsp(x)

	if( !is.null(Lag(filter)) )
		return( lag(w,Lag(filter)) )
	else
		return(w)
}


#	Auto and Cross Covariance and Correlation Computation
#	This code uses fast convolution rather than computing
#	convolution directly.

autocov <-
function(x,maxlag=length(x)-1,plot=F,main="Auto Covariance Function") {

	nxy <- length(x)
	n <- nextn(nxy+maxlag+1)

	xx <- rep(0,n)
	xx[1:nxy] <- (x-mean(x))
	
	z <- fft(xx)
	z <- (Re(fft(z*Conj(z),inv=T))/n)/nxy
	z <- z[1:(maxlag+1)]
	lag <- seq(from=0,len=length(z))/frequency(x)
	attr(z,"lags") <- lag

	if( plot ) {
		plot(
			x=lag,
			y=z,
			xlab="Lag",
			ylab="",
			main=main,
			type="h"
		)
	}

	return(invisible(z))
}


crosscov <-
function( x, y, maxlag=length(x)-1, plot=F, main="Cross Covariance Function") {

	if( length(x)!=length(y) )
		fatal("series length must be equal")

	nxy <- length(x)
	n <- nextn(nxy+maxlag+1)

	xx <- rep(0,n)
	yy <- rep(0,n)
	xx[1:nxy] <- (x-mean(x))
	yy[1:nxy] <- (y-mean(y))
	
	z <- (Re(fft(fft(xx)*Conj(fft(yy)),inv=T))/n)/nxy
	z <- c(z[(n-maxlag+1):n],z[1:(maxlag+1)])
	lag <- seq(from=-maxlag,len=length(z))/frequency(x)
	attr(z,"lags") <- lag

	if( plot ) {
		plot(
			x=lag,
			y=z,
			xlab="Lag",
			ylab="",
			main=main,
			type="h"
		)
	}
	
	return(invisible(z))
}


autocor <-
function( x, maxlag=length(x)-1, plot=F, main="Auto Correlation Function") {

	nxy <- length(x)
	n <- nextn(nxy+maxlag+1)

	xx <- rep(0,n)
	xx[1:nxy] <- (x-mean(x))
	
	z<-fft(xx)
	z <- (Re(fft(z*Conj(z),inv=T))/n)/nxy
	z <- z[1:(maxlag+1)]
	z <- z/z[1]
	lag <- seq(from=0,len=length(z))/frequency(x)
	attr(z,"lags") <- lag

	if( plot ) {
		plot(
			x=lag,
			y=z,
			xlab="Lag",
			ylab="",
			main=main,
			type="h"
		)
	}

	return(invisible(z))
}


crosscor <-
function( x, y, maxlag=length(x)-1, plot=F, main="Cross Correlation Function") {

	if( length(x)!=length(y) )
		fatal("series length must be equal")

	nxy <- length(x)
	n <- nextn(nxy+maxlag+1)

	xx <- rep(0,n)
	yy <- rep(0,n)
	xx[1:nxy] <- (x-mean(x))
	yy[1:nxy] <- (y-mean(y))
	
	z <- (Re(fft(fft(xx)*Conj(fft(yy)),inv=T))/n)/nxy
	z <- c(z[(n-maxlag+1):n],z[1:(maxlag+1)])
	z <- z/z[maxlag+1]
	lag <- seq(from=-maxlag,len=length(z))/frequency(x)
	attr(z,"lags") <- lag

	if( plot ) {
		plot(
			x=lag,
			y=z,
			xlab="Lag",
			ylab="",
			main=main,
			type="h"
		)
	}
	
	return(invisible(z))
}


#	Spectrum estimation and Cross-spectrum Analysis
#	The following functions provide the functionality of
#	David Brillingers "twospe" fortran subroutine.
#
#	Only the Daniel window is provided here.
#	(can you say "prewhiten" - I knew you could).
#	The code should be easy to change for other windows
#	because smoothing is done by fast convolution with
#	the smoothing kernel.

power.spectrum <- function(x,bw=0) {
	nx <- length(x)
	n <- nextn(nx)
	n2 <- n%/% 2 + 1
	u <- rep(0,n)
	u[1:nx] <- (x-mean(x))

	# periodogram

	Ixx <- Mod(fft(u))^2/(2*pi*length(x))
	Ixx[1] <- 0

	# smooth using fast convolution

	k <- floor((n*bw)/2)
	df <- 4*k+2
	if( k>0 ) {
		v <- rep(0,n)
		v[c(1:(k+1),(n-k+1):n)] <- 1/(2*k+1)
		fxx <- (Re(fft(fft(Ixx)*fft(v),inv=T))/n)[1:n2]
	} else {
		fxx <- Ixx[1:n2]
	}

	freq <- ((1:n2)-1)/n
	return(list(freq=freq,fxx=fxx,bw=bw,df=df,n=nx))
}


cross.spectrum <- function(x,y,bw=0)
{
	nx <- length(x)
	n <- nextn(nx)
	n2 <- n %/% 2 + 1
	u <- rep(0,n)
	v <- rep(0,n)
	u[1:nx] <- (y-mean(y))
	v[1:nx] <- (x-mean(x))

	# periodograms and cross-periodogram

	Iyy <- Mod(fft(u))^2/(2*pi*length(x))
	Ixx <- Mod(fft(v))^2/(2*pi*length(x))
	Iyx <- fft(u)*Conj(fft(v))/(2*pi*length(x))
	Iyy[1] <- 0
	Ixx[1] <- 0

	# smoothing

	k <- floor((n*bw)/2)
	df <- 4*k+2
	if( k>0 ) {
		v[1:n] <- 0
		v[c(1:(k+1),(n-k+1):n)] <- 1/(2*k+1)
		fyy <- (Re(fft(fft(Iyy)*fft(v),inv=T))/n)[1:n2]
		fxx <- (Re(fft(fft(Ixx)*fft(v),inv=T))/n)[1:n2]
		fyx <- (fft(fft(Iyx)*fft(v),inv=T)/n)[1:n2]
	} else {
		fyy <- Iyy[1:n2]
		fxx <- Ixx[1:n2]
		fyx <- Iyx[1:n2]
	}

	freq <- ((1:n2)-1)/n
	return(list(freq=freq,fxx=fxx,,fyy=fyy,fyx=fyx,bw=bw,df=df,n=nx))
}


twospe <- function(x,y,bw=0,cb=0,xname="X",yname="Y")
{
	if( missing(x) ) fatal("Need at least one input series")
	if( missing(y) )
		z <- power.spectrum(x,bw)
	else
		z <- cross.spectrum(x,y,bw)

	plot(
		z$freq, ifelse(z$fxx>0,z$fxx,NA),
		log="y", type="l",
		main=paste("Power Spectrum of",xname),
		xlab=paste("Bandwidth =",bw,"    Degrees of freedom =",z$df),
		ylab="Power Spectrum"
	)
	if( cb>0 ) {
		se <- 10^(cb*0.4342945/sqrt(z$df/2))
		lines(z$freq,z$fxx*se,lty=2)
		lines(z$freq,z$fxx/se,lty=2)
	}

	if( missing(y) ) return(invisible())

        gain <- Mod(z$fyx/z$fxx)
        phase<- Arg(z$fyx/z$fxx)
        rxy2 <- Mod(z$fyx)^2/(z$fyy*z$fxx)
        fee <- z$fyy-Mod(z$fyx)^2/z$fxx

	plot(
		z$freq, ifelse(z$fyy>0,z$fyy,NA),
		log="y", type="l",
		main=paste("Power Spectrum of",yname),
		xlab=paste("Bandwidth =",bw,"    Degrees of freedom =",z$df),
		ylab="Power Spectrum"
	)
	if( cb>0 ) {
		lines(z$freq,z$fyy*se,lty=2)
		lines(z$freq,z$fyy/se,lty=2)
	}

	plot(
		z$freq, ifelse(gain>0,gain,NA), log="y", type="l",
		main=paste("Gain of",yname,"on",xname),
		xlab="",
		ylab="Gain"
	)
	if( cb>0 ) {
		se <- 10^((cb*0.4342945/sqrt(z$df/4))*sqrt(1/rxy2-1))
		lines(z$freq,gain*se,lty=2)
		lines(z$freq,gain/se,lty=2)
	}
		

	plot(
		z$freq, phase, pch="+",
		ylim=c(-pi,pi),
		main=paste("Phase of",yname,"on",xname),
		xlab="",
		ylab="Phase"
	)
	if( cb>0 ) {
		se <- (cb*0.4342945/sqrt(z$df/4))*sqrt(1/rxy2-1)
		points(z$freq,phase+se,pch="-")
		points(z$freq,phase-se,pch="-")
	}

	plot(
		z$freq, rxy2, type="l",
		ylim=c(0,1),
		main=paste("Coherence of",xname,"and",yname),
		xlab=ifelse(z$df>2,
			"The horizontal line gives the 95% null point",
			""),
		ylab="Coherence"
	)
	if( z$df>2 ) abline(h=1-0.05^(1/(z$df/2-1)),lty=2)

	plot(
		z$freq, ifelse(fee>0,fee,NA),
		log="y", type="l",
		main="Residual Spectrum",
		xlab="",
		ylab="Residual Spectrum"
	)
	if( cb>0 ) {
		se <- 10^(cb*0.4342945/sqrt(z$df/2-1))
		lines(z$freq,fee*se,lty=2)
		lines(z$freq,fee/se,lty=2)
	}

	A <- z$fyx/z$fxx
	A <- c(A,Conj(rev(A[2:(length(A)-1)])))
	n <- length(A)
	coef <- Re(fft(A,inv=T))/n
	coef <- coef[c((n%/%2+1):n,1:(n%/%2))]
	time <- 1:n-n%/%2-1
	plot(
		time,
		coef,
		main="Impulse Response",
		xlab="Lag",
		ylab="Coefficient",
		type="h"
	)

	return(invisible())
}
