library(reshape)
library(tidyverse)
library(ggforce)

normalize_landscape <- function(landscape) {
    min_val <- min(landscape)
    max_val <- max(landscape)
    range_val <- max_val - min_val 

    return((landscape - min_val) /  range_val)
}

plot_landscape <- function(landscape) {
    p <- ggplot(data=melt(landscape), aes(x=X1, y=X2, z=value)) +
        geom_contour_filled()

    return(p)
}

init_population <- function(landscape, n) {

    population <- list()
    dims <- length(dim(landscape))

    for(i in 1:n) {
        coords <- round(runif(dims, 0, 1) * dim(landscape))
        sigmas <- rnorm(dims)

        population[[i]] <- matrix(c(coords, sigmas), ncol=2) 
    }

    return(population)
}

next_generation <- function(landscape, population) {
    return(map(population, function(indiv) {
        return(select_indiv(landscape, list(indiv, create_child(indiv))))
    })) 
} 

select_indiv <- function(landscape, indivs) {
    return(reduce(indivs, function(a, b) {
        if(eval_indiv(landscape, a) >= eval_indiv(landscape, b)) {
            return(a)
        }

        return(b)
    })) 
}

eval_indiv <- function(landscape, indiv) {
    dims <- dim(landscape)

    x <- indiv[1,1]
    y <- indiv[2, 1]

    if(x > dims[1] || y > dims[2] || x < 1 || y < 1) {
      return(-1)
    }

    return(landscape[indiv[1,1], indiv[2,1]])
}

create_child <- function(parent) {
    new_sigmas <- mutate_sigmas(parent[,2])
    new_coords <- mutate_coords(parent[,1], new_sigmas)

    return(matrix(c(new_coords, new_sigmas), ncol=2))
}

mutate_sigmas <- function(sigmas) {
    global_rate <- 1 / sqrt(2 * length(sigmas))
    local_rate <- 1 / (2 * sqrt(length(sigmas)))
    
    global_step <- global_rate * rnorm(1)

    return(map_dbl(sigmas, function(s) s * exp(global_step + local_rate * rnorm(1))))
}

mutate_coords <- function(coords, sigmas) {
    return(imap_dbl(coords, function(x, i) x + sigmas[i] * rnorm(1)))
}

experiment <- function(landscape, population, gens) {
    df <- population_to_df(landscape, population, 0)

    for(g in 1:gens) {
        population <- next_generation(landscape, population)

        df <- rbind(df, population_to_df(landscape, population, g))
    }

    return(df)
}

population_to_df <- function(landscape, population, gen) {
    df <- reduce(imap(population, function(indv, i) indiv_to_df(landscape, indv, i)), rbind)
    df["generation"] <- gen

    return(df)
}

indiv_to_df <- function(landscape, indiv, index) {
    return(data.frame(x=indiv[1,1], y=indiv[2,1],
                      sx=indiv[1, 2], sy=indiv[2,2],
                      individual=index, value=eval_indiv(landscape, indiv)))
}

plot_generation <- function(landscape, df) {
    p <-ggplot(data=df) +
        geom_contour_filled(data=melt(landscape), aes(x=X1, y=X2, z=value)) +
        geom_point(aes(x=x, y=y)) +
        geom_ellipse(aes(x0=x, y0=y, a=sx, b=sy, angle=0))

    return(p)
}

plot_experiment <- function(landscape, df, filename) {
    pdf(file=filename, onefile=TRUE)

    for(g in unique(df$generation)) {
        tmp_df <- df[df$generation == g,]

        p <- plot <- plot_generation(landscape, tmp_df)

        print(p)
    }

    dev.off()
}