library(ggplot2)
library(gganimate)
library(patchwork)

gradient <- function(f, x, d){
    return((f(x + d) - f(x - d)) / (2*d))
}

gradient.ascent.move <- function(f, x, d, mu){
    return(x + mu * gradient(f, x, d))
}

func.g <- function(x) x
func.k <- function(x) sin(x)
func.h <- function(x) x * sin(x)
func.l <- function(x) 2 + cos(x) + sin(2*x)

gradient.ascent.iterate <- function(f, x, d, mu, n){
    if(n == 1) {
        return(gradient.ascent.move(f, x, d, mu))
    }

    return(gradient.ascent.niter(f
                                 ,gradient.descent.move(f, x, d, mu)
                                 ,d
                                 ,mu
                                 ,n-1
                                ))
}

gradient.ascent.iterverb <- function(f, x, d, mu, n, xs=numeric()){
    next_x <- gradient.ascent.move(f, x, d, mu)
    xs[length(xs)+1] <- next_x

    if(n == 1) {
        return(xs)
    }

    return(gradient.ascent.iterverb(f, next_x, d, mu, n-1, xs))
}

trace.ascent <- function(f, x, d, eta, n, xs) {
    df_dc = data.frame(x=numeric()
                       ,y=numeric()
                       ,i=integer()
                       ,start_x=character()
                       ,eta=numeric())

    for(start in x) {
        for(e in eta) {
            first_it <- TRUE
            if(first_it == TRUE) {
                df_dc <- rbind(df_dc, data.frame(x=c(start)
                                                 ,y=c(f(start))
                                                 ,i=c(0)
                                                 ,start_x=c(as.character(start))
                                                 ,eta=c(e)
                                                ))
                first_it <- FALSE
            }

            xf <- gradient.ascent.iterverb(f, start, d, e, n)
    
            df_dc <- rbind(df_dc, data.frame(x=xf
                                             ,y=f(xf)
                                             ,i=1:length(xf)
                                             ,start_x=rep(as.character(start), length(xf))
                                             ,eta=e)
                                            )
        }
    }

    return(df_dc)
}

plot.ascent <- function(f, x, d, eta, n, xs) {
    df_dc = trace.ascent(f, x, d, eta, n, xs)

    func_str = deparse(substitute(f))

    df_f <- data.frame(x=xs, y=f(xs))
    p1 <- ggplot(df_f, aes(x=x, y=y)) +
            geom_line() +
            geom_point(aes(colour=start_x
                           ,size=i
                          )
                       ,data=df_dc) +
            labs(size="iteration"
                 ,alpha="iteration"
                 ,color="start x"
                 ,y=sprintf("%s(x)", func_str)) +
            facet_grid(eta ~ ., labeller=label_both)

    p2 <- ggplot(df_dc, aes(x=i, y=y)) +
          geom_line(aes(colour=start_x), show.legend=FALSE) +
          labs(x="iteration"
               ,y=sprintf("%s(x)", func_str)) +
                  facet_grid(eta ~ ., labeller=label_both)

    p <- (p1 | p2) + 
            plot_annotation(title=sprintf("function: %s", func_str)) +
            plot_layout(guides="collect"
                        ,widths=10
                        ,heights=2) 

    return(p)
}

animate.ascent <- function(f, x, d, eta, n, xs) {
    df_dc = trace.ascent(f, x, d, eta, n, xs)

    func_str <- deparse(substitute(f))

    df_f <- data.frame(x=xs, y=f(xs))

    p <- ggplot(df_f, aes(x=x, y=y)) +
            geom_line() +
            geom_point(aes(colour=start_x), size=2.5, data=df_dc) +
            labs(color="start x", y=sprintf("%s(x)", func_str)) +
            facet_grid(eta ~ ., labeller=label_both) +
            ggtitle(sprintf("function: %s", func_str))

    anim <- p + transition_reveal(i)

    return(anim)
}