library(ggplot2)
library(gganimate)


linscale.initpopulation <- function(n) {
    return(runif(10))
}

linscale.scale <- function(population, a, b) {
    return(a * population + b)
}

linscale.relfitness <- function(population) {
    return(population / sum(population))
}

linscale.tracegens <- function(popul, a, b, n) {
    lenpopul <- length(popul)

    df <- data.frame( generation = 0
                     ,individual = 1:lenpopul
                     ,fitness = popul
                     ,relfitness = linscale.relfitness(popul)
                    )

    
    
    for(i in 1:n) {
        popul <- linscale.scale(popul, a, b)

        df <- rbind(df, data.frame( generation = i
                                   ,individual = 1:lenpopul
                                   ,fitness = popul
                                   ,relfitness = linscale.relfitness(popul)))
    }

    return(df)
}

linscale.experiment <- function(avals, bvals, n, popul=c()) {
    df <- data.frame( generation = integer()
                     ,individual = integer()
                     ,fitness = double()
                     ,relfitness = double()
                     ,a = double()
                     ,b = double())

    if(length(popul) == 0) {
        popul <- linscale.initpopulation(10)
    }

    for(a in avals) {
        for(b in bvals) {
            dftmp <- linscale.tracegens(popul, a, b, n)
            dftmp["a"] <- a
            dftmp["b"] <- b

            df <- rbind(df, dftmp)
        }    
    }

    return(df)
}

linscale.plot <- function(df) {
    pdf(file="linscale.pdf", onefile=TRUE)
    
    for(g in unique(df$generation)) {
        p <- ggplot(data=df[df$generation == g, ]) +
            geom_col(aes(x=individual, y=relfitness)) +
            facet_grid(b ~ a, labeller = label_both) +
            labs(title=sprintf("generation: %i", g))
        print(p)
    }

    dev.off()
}

linscale.animate <- function(df) {
    anim <- ggplot(data=df) +
        geom_col(aes(x=individual, y=relfitness)) +
        facet_grid(b ~ a, labeller = label_both) +
        labs(title="generation: {closest_state}") +
        transition_states(generation)

    anim_save("linescale.gif", animation=anim)
}

data.population <- linscale.initpopulation(10)