AMR - adaptive mesh refinement

amr <- function(x, ...) UseMethod("amr", x)
is.amr <- function(x) inherits(x, "amr")

#temp
amr_vals <- function(node) if(length(node[["refinement"]])>0) {
    # browser()
    unlist(lapply(node[["refinement"]], amr_vals))
    } else node[["value"]]

# x must be a matrix with two columns
amr.matrix <- function(x, weights=NULL, xlim=NULL, ylim=NULL, ...) {
    if(is.null(xlim)) xlim <- c(min(x[,1]), max(x[,1]))
    if(is.null(ylim)) ylim <- c(min(x[,2]), max(x[,2]))
    idx <- which(!is.na(x[,1]) & !is.na(x[,2]))
    amr.default(x[idx,], weights=weights, type="data", xlim=xlim, ylim=ylim, ...)
}

# x must be a function that supports vector arguments
amr.function <- function(x, ...) amr.default(x, type="function", ...)

# the workhorse, since the difference between function and dataset is small
amr.default <- function(x, xlim, ylim, weights=NULL, minlevel=4, maxlevel=8, reltol=0.05, log="", type, ...) {
    pts <- matrix(c(c(xlim,rev(xlim)), rep(ylim, each=2), rep(NA, 4)), byrow=FALSE, ncol=3)
    if(type=="function") pts[1:4, 3] <- x(pts[1:4,1], pts[1:4, 2])
    loglist <- strsplit(log,"")[[1]]
    logx <- "x" %in% loglist
    logy <- "y" %in% loglist
    if(type=="data") zlim <- c(0,1)

    pgrad.function <- function(node) {
        vals <- pts[node[["pts_idx"]],3]
        vals <- vals[is.finite(vals)]
        ret <- sd(vals)/(mean(vals)+1)
        # browser()
        return(ret)
    }
    pgrad.data <- function(node) {
        # prob. mass of this rectangle
        ll <- pts[node[["pts_idx"]][1],1:2]
        ur <- pts[node[["pts_idx"]][3],1:2]
        ret <- node[["value"]] * (ur[[1]]-ll[[1]]) * (ur[[2]]-ll[[2]])
        return(ret)
    }
    pgrad <- switch(type, "function"=pgrad.function, "data"=pgrad.data)

    rectval.function <- function(node) {
        # mean of the function value at the corners
        vals <- pts[node[["pts_idx"]],3]
        return(mean(vals[is.finite(vals)]))
    }
    rectval.data <- function(node) {
        # probability mass per area unit. should integrate to 1.
        ll <- pts[node[["pts_idx"]][1],1:2]
        ur <- pts[node[["pts_idx"]][3],1:2]
        idx <- (x[,1]>=ll[[1]]) & (x[,1]<ur[[1]]) & (x[,2]>=ll[[2]]) & (x[,2]<ur[[2]])
        xlen <- (ur[[1]]-ll[[1]])
        ylen <- (ur[[2]]-ll[[2]])
        # if(logx) xlen <- ur[[1]]/ll[[1]] else xlen <- (ur[[1]]-ll[[1]])
        # if(logy) ylen <- ur[[2]]/ll[[2]] else ylen <- (ur[[2]]-ll[[2]])

        ret <- if(is.null(weights)) sum(idx)/(nrow(x)*xlen*ylen) else sum(weights[idx]) / (sum(weights)*xlen*ylen)
        # if(ret>10000) browser()
        # ret <- (sum(idx)*total_area) / (nrow(x) * (ur[[1]]-ll[[1]]) * (ur[[2]]-ll[[2]]))
        zlim[[2]] <<- max(c(zlim[[2]], ret))
        return(ret)
    }
    rectval <- switch(type, "function"=rectval.function, "data"=rectval.data)

    refinement <- function(node) {
        ll <- pts[node[["pts_idx"]][1],1:2]
        ur <- pts[node[["pts_idx"]][3],1:2]
        lvl <- node[["level"]]
        midx <- if(logx) sqrt(ll[1]*ur[1]) else (ll[1]+ur[1])/2
        midy <- if(logy) sqrt(ll[2]*ur[2]) else (ll[2]+ur[2])/2

        add_pts <- matrix(c(
            midx, ll[2], NA,  # lower middle
            ll[1], midy, NA,  # left middle
            midx, midy, NA, # middle
            ur[1], midy, NA,  # right middle
            midx, ur[2], NA), # upper middle
            byrow=TRUE, ncol=3
        )
        if(type=="function") add_pts[,3] <- x(add_pts[,1], add_pts[,2])
        old_nrow <- nrow(pts)
        pts <<- rbind(pts, add_pts)

        node1 <- list(pts_idx=c(node[["pts_idx"]][1], old_nrow+c(1,3,2)), refinement=list(), level=lvl+1)
        node1[["value"]] <- rectval(node1)
        node1[["pgrad"]] <- pgrad(node1)
        node2 <- list(pts_idx=c(old_nrow+1, node[["pts_idx"]][2], old_nrow+c(4,3)), refinement=list(), level=lvl+1)
        node2[["value"]] <- rectval(node2)
        node2[["pgrad"]] <- pgrad(node2)
        node3 <- list(pts_idx=c(old_nrow+c(3,4), node[["pts_idx"]][3], old_nrow+5), refinement=list(), level=lvl+1)
        node3[["value"]] <- rectval(node3)
        node3[["pgrad"]] <- pgrad(node3)
        node4 <- list(pts_idx=c(old_nrow+c(2,3,5), node[["pts_idx"]][4]), refinement=list(), level=lvl+1)
        node4[["value"]] <- rectval(node4)
        node4[["pgrad"]] <- pgrad(node4)

        return(list(node1, node2, node3, node4))
    }

    needs_refinement <- function(node) {
        if(node[["level"]]<minlevel) return(TRUE)
        if(node[["level"]]>=maxlevel) return(FALSE)
        if(!is.numeric(node[["pgrad"]]) | is.nan(node[["pgrad"]])) return(TRUE)
        if(abs(node[["pgrad"]]) > reltol) return(TRUE)
        # if(is.numeric(node[["pgrad"]])) browser()
        # ans <- try(if(node[["pgrad"]] > reltol) return(TRUE))
        # if (inherits(ans, "try-error")) browser()

        return(FALSE)
    }

    refine_recursively <- function(node) {
        if(needs_refinement(node)) {
            node[["refinement"]] <- refinement(node)
            # browser()
            for (i in seq_len(length(node[["refinement"]])))
                node[["refinement"]][[i]] <- refine_recursively(node[["refinement"]][[i]])
        }
        return(node)
    }

    root <- list(pts_idx=seq(1,4), refinement=list(), level=0)
    root[["value"]] <- rectval(root)
    root[["pgrad"]] <- pgrad(root)
    root <- refine_recursively(root)
    if(type=="function")  {
        zlim <- c(min(pts[,3]), max(pts[,3]))
        # browser()
    } else if(type=="data") {
        # browser()
        vals <- amr_vals(root)
        zlim <- c(min(vals), max(vals))
    }

    res <- list(pts=pts, node=root, zlim=zlim)
    class(res) <- "amr"
    return(res)
}

# use classInt? quantile method not helpful due to irregular sampling
# image has a breaks argument that looks cool
plot.amr <- function(x, col=heat.colors(50), border=NA, xlab="x", ylab="y", zmag=1, ...) {
    zmin <- if(is.finite(x[["zlim"]][[1]])) x[["zlim"]][[1]] else -zmag
    zmax <- if(is.finite(x[["zlim"]][[2]])) x[["zlim"]][[2]] else zmag
    colidx <- seq(zmin, zmax, length.out=length(col))

    plot_node <- function(node) {
        ll <- x[["pts"]][node[["pts_idx"]][1],1:2]
        ur <- x[["pts"]][node[["pts_idx"]][3],1:2]
        rect(ll[1], ll[2], ur[1], ur[2], border=border, col=col[findInterval(node[["value"]], colidx)])
        if(length(node[["refinement"]])>0) lapply(node[["refinement"]], plot_node)
    }

    root <- x[["node"]]
    ll <- x[["pts"]][root[["pts_idx"]][1], 1:2]
    ur <- x[["pts"]][root[["pts_idx"]][3], 1:2]
    plot_arg <- list(...)
    plot_arg[["x"]] <- 1; plot_arg[["col"]] <- "white"; plot_arg[["type"]] <- "p"; plot_arg[["xlim"]] <- c(ll[1], ur[1])
    plot_arg[["ylim"]] <- c(ll[2], ur[2]); plot_arg[["xlab"]] <- xlab; plot_arg[["ylab"]] <- ylab
    do.call(plot, plot_arg)
    plot_node(root)
}


# may need a different method (weighted mean instead of sum) for function
integrate_amr <- function(x, xlim, ylim) {
    eval_rect <- function(node) {
        ll <- x[["pts"]][node[["pts_idx"]][1],1:2]
        ur <- x[["pts"]][node[["pts_idx"]][3],1:2]
        if(xlim[1]>ur[1] || xlim[2]<ll[1] || ylim[1]>ur[2] || ylim[2]<ll[2]) return(NULL)
        if(length(node[["refinement"]])>0) return(unlist(lapply(node[["refinement"]], eval_rect)))
        area <- (min(c(ur[1], xlim[2])) - max(c(ll[1], xlim[1]))) * (min(c(ur[2], ylim[2]))-max(c(ll[2], ylim[1])))
        tot_area <- (ur[1]-ll[1]) * (ur[2]-ll[2])
        return(node[["value"]] * area / tot_area) # this is integreation
    }
    sum(eval_rect(x[["node"]])) # this is integration
}

# guess xlab and ylab
# add breaks argument
image.amr <- function(x, col=rev(gray(0:12/12)), border=NA, xlim=NULL, ylim=NULL, log="", xlab="x", ylab="y", zmag=1, length.out=100, breaks=NULL, ...) {
    if(is.null(xlim)) xlim <- c(min(x[["pts"]][,1]), max(x[["pts"]][,1]))
    if(is.null(ylim)) ylim <- c(min(x[["pts"]][,2]), max(x[["pts"]][,2]))

    zmin <- if(is.finite(x[["zlim"]][[1]])) x[["zlim"]][[1]] else -zmag
    zmax <- if(is.finite(x[["zlim"]][[2]])) x[["zlim"]][[2]] else zmag

    if(length(length.out)==1) length.out <- rep(length.out,2)

    loglist <- strsplit(log,"")[[1]]
    logx <- "x" %in% loglist
    logy <- "y" %in% loglist

    xgrid <- if(logx) exp(seq(log(xlim[1]), log(xlim[2]), length.out=length.out[1])) else seq(xlim[1], xlim[2], length.out=length.out[1])
    ygrid <- if(logy) exp(seq(log(ylim[1]), log(ylim[2]), length.out=length.out[2])) else seq(ylim[1], ylim[2], length.out=length.out[2])
    z <- matrix(0, nrow=length.out[1], ncol=length.out[2])

    for (i in seq.int(2,length.out[1]))
        for (j in seq.int(2,length.out[2]))
            z[i,j] <- integrate_amr(x, xlim=xgrid[seq.int(i-1,i)], ylim=ygrid[seq.int(j-1,j)])

    if(!is.null(breaks) && length(breaks)==1)  {
        breaks <- classIntervals(as.vector(z), n=breaks, style="fisher")$brks
        # browser()
        # breaks <- quantile(as.vector(z), seq(0,1, length.out=(breaks+1)))
    }
    # browser()
    image(x=xgrid, y=ygrid, z=z, zlim=c(zmin, zmax), xlim=xlim, ylim=ylim, col=col, xlab=xlab, ylab=ylab, log=log, breaks=breaks, ...)


}