## based on code by Martyn Plummer, plus kernel code by Adrian Trapletti
spectrum <- function (..., method = c("pgram", "ar"))
{
    switch(match.arg(method),
	   pgram = spec.pgram(...),
	   ar	 = spec.ar(...)
	   )
}

## spec.taper based on code by Kurt Hornik
spec.taper <- function (x, p = 0.1)
{
    if (any(p < 0) || any(p > 0.5))
        stop("p must be between 0 and 0.5")
    a <- attributes(x)
    x <- as.matrix(x)
    nc <- ncol(x)
    if (length(p) == 1)
        p <- rep(p, nc)
    else if (length(p) != nc)
        stop("length of p must be 1 or equal the number of columns of x")
    nr <- nrow(x)
    for (i in 1:nc) {
        m <- floor(nr * p[i])
        if(m == 0) next
        w <- 0.5 * (1 - cos(pi * seq(1, 2 * m - 1, by = 2)/(2 * m)))
        x[, i] <- c(w, rep(1, nr - 2 * m), rev(w)) * x[, i]
    }
    attributes(x) <- a
    x
}

spec.ar <- function(x, n.freq, order = NULL, plot = TRUE,
                    na.action = na.fail, method = "yule-walker", ...)
{
    ## can be called with a ts or a result of an AR fit.
    if(!is.list(x)) {
        series <- deparse(substitute(x))
        x <- na.action(as.ts(x))
        xfreq <- frequency(x)
        n <- NROW(x)
        nser <- NCOL(x)
        x <- ar(x, is.null(order), order, na.action=na.action, method=method)
    } else {
        cn <- match(c("ar", "var.pred", "order"), names(x))
        if(any(is.na(cn)))
            stop("x must be a time series or an ar() fit")
        series <- x$series
        xfreq <- x$frequency
        if(is.array(x$ar)) nser <- dim(x$ar)[2] else nser <- 1
    }
    n <- x$n.used
    order <- x$order
    if(missing(n.freq)) n.freq <- 500
    freq <- seq(0, 0.5, length = n.freq)
    if (nser == 1) {
        coh <- phase <- NULL
        if(order >= 1) {
            cs <- outer(freq, 1:order, function(x, y) cos(2*pi*x*y)) %*% x$ar
            sn <- outer(freq, 1:order, function(x, y) sin(2*pi*x*y)) %*% x$ar
            spec <- x$var.pred/(xfreq*((1 - cs)^2 + sn^2))
        } else
            spec <- rep(x$var.pred/(xfreq), length(freq))
    } else .NotYetImplemented()
    spg.out <- list(freq = freq, spec = spec, coh = coh, phase = phase,
                    n.used = nrow(x), series = series,
                    method = paste("AR (", order, ") spectrum ", sep="")
                    )
    class(spg.out) <- "spec"
    if(plot) {
	plot(spg.out, ci = 0, ...)
        return(invisible(spg.out))
    } else return(spg.out)
}

spec.pgram <-
    function (x, spans = NULL, kernel = NULL, taper = 0.1,
              pad = 0, fast = TRUE,
              demean = FALSE, detrend = TRUE,
              plot = TRUE, na.action = na.fail, ...)
{
    ## Estimate spectral density from (smoothed) periodogram.
    series <- deparse(substitute(x))
    x <- na.action(as.ts(x))
    xfreq <- frequency(x)
    x <- as.matrix(x)
    N <- nrow(x)
    nser <- ncol(x)
    if(!is.null(spans)) # allow user to mistake order of args
        if(is.tskernel(spans)) kernel <- spans
        else kernel <- kernel("modified.daniell", spans %/% 2)
    if(!is.null(kernel) && !is.tskernel(kernel))
        stop("must specify spans or a valid kernel")
    if (detrend) {
        t <- 1:N - (N + 1)/2
        sumt2 <- N * (N^2 - 1)/12
        for (i in 1:ncol(x))
            x[, i] <- x[, i] - mean(x[, i]) - sum(x[, i] * t) * t/sumt2
    }
    else if (demean) {
        x <- sweep(x, 2, apply(x, 2, mean))
    }
    x <- spec.taper(x, taper)
    ## to correct for tapering: Bloomfield (1976, p. 194)
    ## Total taper is taper*2
    u2 <- (1 - (5/8)*taper*2)
    u4 <- (1 - (93/128)*taper*2)
    if (pad > 0) {
        x <- rbind(x, matrix(0, nrow = N * pad, ncol = ncol(x)))
        N <- nrow(x)
    }
    NewN <- if(fast) nextn(N) else N
    x <- rbind(x, matrix(0, nrow = (NewN - N), ncol = ncol(x)))
    N <- nrow(x)
    Nspec <- floor(N/2)
    freq <- seq(from = xfreq/N, by = xfreq/N, length = Nspec)
    xfft <- mvfft(x)
    pgram <- array(NA, dim = c(N, ncol(x), ncol(x)))
    for (i in 1:ncol(x)) {
        for (j in 1:ncol(x)) {
            pgram[, i, j] <- xfft[, i] * Conj(xfft[, j])/(N*xfreq)
        }
    }
    ## value at zero is invalid as mean has been removed, so interpolate
    pgram[1, i, j] <- 0.5*(pgram[2, i, j] + pgram[N, i, j])
    if(!is.null(kernel)) {
        for (i in 1:ncol(x)) for (j in 1:ncol(x))
                pgram[, i, j] <- kernapply(pgram[, i, j], kernel,
                                           circular = TRUE)
        df <- df.kernel(kernel)/(u4/u2^2)
        bandwidth <- bandwidth.kernel(kernel) * xfreq/N
    } else {
        df <- 2/(u4/u2^2)
        bandwidth <- sqrt(1/12) * xfreq/N
    }
    pgram <- pgram[1+(1:Nspec),,, drop=FALSE]
    spec <- matrix(NA, nrow = Nspec, ncol = nser)
    for (i in 1:nser) spec[, i] <- Re(pgram[1:Nspec, i, i])
    if (nser == 1) {
        coh <- phase <- NULL
    } else {
        coh <- phase <- matrix(NA, nrow = Nspec, ncol = nser * (nser - 1)/2)
        for (i in 1:(nser - 1)) {
            for (j in (i + 1):nser) {
                coh[, i + (j - 1) * (j - 2)/2] <-
                    Mod(pgram[, i, j])^2/(spec[, i] * spec[, j])
                phase[, i + (j - 1) * (j - 2)/2] <- Arg(pgram[, i, j])
            }
        }
    }
    ## correct for tapering
    for (i in 1:nser) spec[, i] <- spec[, i]/u2
    spec <- drop(spec)
    spg.out <-
        list(freq = freq, spec = spec, coh = coh, phase = phase,
             kernel = kernel, df = df,
             bandwidth = bandwidth, n.used = nrow(x),
             series = series, snames = colnames(x),
             method = ifelse(!is.null(kernel), "Smoothed Periodogram",
                             "Raw Periodogram"),
             taper = taper, pad = pad, detrend = detrend, demean = demean)
    class(spg.out) <- "spec"
    if(plot) {
	plot(spg.out, ...)
        return(invisible(spg.out))
    } else return(spg.out)
}

plot.spec <-
    function (x, add = FALSE, ci = 0.95, log = c("yes", "dB", "no"),
              xlab = "frequency", ylab = NULL,
              type = "l", ci.col="blue", main = NULL, sub = NULL,
              plot.type = c("marginal", "coherency", "phase"), ...)
{
    spec.ci <- function (spec.obj, coverage = 0.95)
    {
        ## A utility function for plot.spec which calculates the confidence
        ## interval (centred around zero). We use a conditional argument to
        ## ensure that the ci always contains zero.

        if (coverage < 0 || coverage >= 1)
            stop("coverage probability out of range [0,1)")
        tail <- (1 - coverage)
        df <- spec.obj$df
        upper.quantile <- 1 - tail * (1 - pchisq(df, df))
        lower.quantile <- tail * pchisq(df, df)
        1/(qchisq(c(upper.quantile, lower.quantile), df)/df)
    }

    plot.type <- match.arg(plot.type)
    m <- match.call()
    if(plot.type == "coherency") {
        m[[1]] <- as.name("plot.spec.coherency")
        m$plot.type <- m$log <- m$add <- NULL
        return(eval(m, parent.frame()))
    }
    if(plot.type == "phase") {
        m[[1]] <- as.name("plot.spec.phase")
        m$plot.type <- m$log <- m$add <- NULL
        return(eval(m, parent.frame()))
    }
    if(is.null(ylab))
        ylab <- if(log == "dB") "spectrum (dB)" else "spectrum"
    if(is.logical(log))
        log <- if(log) "yes" else "no"
    if(missing(log) && getOption("ts.S.compat")) log <- "dB"
    log <- match.arg(log)
    ylog <- ""
    if(log=="dB") x$spec <- 10 * log10(x$spec)
    if(log=="yes") ylog <- "y"
    if(add) {
        matplot(x$freq, x$spec, type = type, add=TRUE, ...)
    } else {
        matplot(x$freq, x$spec, xlab = xlab, ylab = ylab, type = type,
                log = ylog, ...)
        is.ar <- !is.na(pmatch("AR", x$method))
        if (ci <= 0 || log == "no" || is.ar) {
            ## No confidence limits
            ci.text <- ""
        } else {
            ## The position of the error bar has no meaning: only the width
            ## and height. It is positioned in the top right hand corner.
            ##
            conf.lim <- spec.ci(x, coverage = ci)
            if(log=="dB") {
                conf.lim <- 10*log10(conf.lim)
                conf.y <- max(x$spec) - conf.lim[2]
                conf.x <- max(x$freq) - x$bandwidth
                lines(rep(conf.x, 2), conf.y + conf.lim, col=ci.col)
                lines(conf.x + c(-0.5, 0.5) * x$bandwidth, rep(conf.y, 2),
                      col=ci.col)
                ci.text <- paste(", ", round(100*ci, 2),  "% C.I. is (",
                                 paste(format(conf.lim, digits = 3),
                                       collapse = ","), ")dB", sep="")
            } else {
                ci.text <- ""
                conf.y <- max(x$spec) / conf.lim[2]
                conf.x <- max(x$freq) - x$bandwidth
                lines(rep(conf.x, 2), conf.y * conf.lim, col=ci.col)
                lines(conf.x + c(-0.5, 0.5) * x$bandwidth, rep(conf.y, 2),
                      col=ci.col)
            }
        }
        if (is.null(main))
            main <- paste(paste("Series:", x$series), x$method, sep = "\n")
        if (is.null(sub) && !is.ar)
             sub <- paste("bandwidth = ", format(x$bandwidth, digits = 3),
                         ci.text, sep="")
        title(main = main, sub = sub)
    }
    invisible(x)
}

## based on code in Venables & Ripley
plot.spec.coherency <-
    function(x, ci = 0.95,
             xlab = "frequency", ylab = "squared coherency", ylim=c(0,1),
             type = "l", main = NULL, ci.lty = 3, ci.col="blue", ...)
{
    nser <- NCOL(x$spec)
    ## Formulae from Bloomfield (1976, p.225)
    gg <- 2/x$df
    se <- sqrt(gg/2)
    z <- -qnorm((1-ci)/2)
    if (is.null(main))
        main <- paste(paste("Series:", x$series),
                      "Squared Coherency", sep = " --  ")
    if(nser == 2) {
        plot(x$freq, x$coh, type=type, xlab=xlab, ylab=ylab, ylim=ylim, ...)
        coh <- pmin(0.99999, sqrt(x$coh))
        lines(x$freq, (tanh(atanh(coh) + z*se))^2, lty=ci.lty, col=ci.col)
        lines(x$freq, (pmax(0, tanh(atanh(coh) - z*se)))^2,
              lty=ci.lty, col=ci.col)
        title(main)
    } else {
        opar <- par(mfrow = c(nser-1, nser-1), mar = c(1.5, 1.5, 0.5, 0.5),
                    oma = c(4, 4, 6, 4))
        on.exit(par(opar))
        plot.new()
        for (j in 2:nser) for (i in 1:(j-1)) {
            par(mfg=c(j-1,i, nser-1, nser-1))
            ind <- i + (j - 1) * (j - 2)/2
            plot(x$freq, x$coh[, ind], type=type, ylim=ylim, axes=FALSE,
                 xlab="", ylab="", ...)
            coh <- pmin(0.99999, sqrt(x$coh[, ind]))
            lines(x$freq, (tanh(atanh(coh) + z*se))^2, lty=ci.lty, col=ci.col)
            lines(x$freq, (pmax(0, tanh(atanh(coh) - z*se)))^2,
                  lty=ci.lty, col=ci.col)
            box()
            if (i == 1) {
                axis(2, xpd = NA)
                title(ylab=x$snames[j], xpd = NA)
            }
            if (j == nser) {
                axis(1, xpd = NA)
                title(xlab=x$snames[i], xpd = NA)
            }
            mtext(main, 3, 3, TRUE, 0.5,
                  cex = par("cex.main"), font = par("font.main"))
        }
    }
}

plot.spec.phase <-
    function(x, ci = 0.95,
             xlab = "frequency", ylab = "phase", ylim=c(-pi, pi),
             type = "l", main = NULL, ci.lty = 3, ci.col="blue", ...)
{
    nser <- NCOL(x$spec)
    ## Formulae from Bloomfield (1976, p.225)
    gg <- 2/x$df
    if (is.null(main))
        main <- paste(paste("Series:", x$series),
                      "Phase spectrum", sep = "  -- ")
    if(nser == 2) {
        plot(x$freq, x$phase, type=type, xlab=xlab, ylab=ylab, ylim=ylim, ...)
        coh <- sqrt(x$coh)
        cl <- asin( pmin( 0.9999, qt(ci, 2/gg-2)*
                         sqrt(gg*(coh^{-2} - 1)/(2*(1-gg)) ) ) )
        lines(x$freq, x$phase + cl, lty=ci.lty, col=ci.col)
        lines(x$freq, x$phase - cl, lty=ci.lty, col=ci.col)
        title(main)
    } else {
        opar <- par(mfrow = c(nser-1, nser-1), mar = c(1.5, 1.5, 0.5, 0.5),
                    oma = c(4, 4, 6, 4))
        on.exit(par(opar))
        plot.new()
        for (j in 2:nser) for (i in 1:(j-1)) {
            par(mfg=c(j-1,i, nser-1, nser-1))
            ind <- i + (j - 1) * (j - 2)/2
            plot(x$freq, x$phase[, ind], type=type, ylim=ylim, axes=FALSE,
                 xlab="", ylab="", ...)
            coh <- sqrt(x$coh[, ind])
            cl <- asin( pmin( 0.9999, qt(ci, 2/gg-2)*
                             sqrt(gg*(coh^{-2} - 1)/(2*(1-gg)) ) ) )
            lines(x$freq, x$phase[, ind] + cl, lty=ci.lty, col=ci.col)
            lines(x$freq, x$phase[, ind] - cl, lty=ci.lty, col=ci.col)
            box()
            if (i == 1) {
                axis(2, xpd = NA)
                title(ylab=x$snames[j], xpd = NA)
            }
            if (j == nser) {
                axis(1, xpd = NA)
                title(xlab=x$snames[i], xpd = NA)
            }
            mtext(main, 3, 3, TRUE, 0.5,
                  cex = par("cex.main"), font = par("font.main"))
        }
    }
}