using StatsBase
using PyPlot
#=
Parents <— {randomly generated population}
While not (termination criterion)
    Calculate the fitness of each parent in the population
    Children <- 0
    While | Children | < | Parents |
        Use fitnesses to probabilistically select a pair of parents for mating
        Mate the parents to create children c\ and c<i
        Children <— Children U {ci,C2}
    Loop
    Randomly mutate some of the children
    Parents <— Children
Next generation

Use continuous GA, not binary

include array of independent variables along with associated range

=#

function initialize_population(populationSize::Int, chromosomeSize::Int, bounds::Array{Tuple{Float64, Float64},1})
    """
    Initializes the population based on the given population size and chromosome size
    Values will be from lb to ub (lower to upper bound)
    """

    # Check to see if we dont have enough bounds or have too many
    if length(bounds) != chromosomeSize
        println("Length of bounds ($(length(bounds))) not equal to specified chromosome size ($chromosomeSize)")
        exit()
    end

    # Init population
    population = Array{Array{Float64}}(0)

    # Foreach member in population
    for i = 1:populationSize
        # Init the member
        row = []
        # Create member elements based on bounds
        for bound in bounds
            push!(row, rand(bound[1]:0.01:bound[2]))
        end
        # Add another element to the end to keep score
        push!(row, 1.0)

        # hcat row, append to 2D population matrix
        population = push!(population, row)
    end
    return population

end

function score_sort_population(population::Array{Array{Float64}}, fitness_function::Function)
    """
    Args
        population: The population to score. Must be 2D array of Float64 elements
        fitness_function: Function to score the population
                          Must be able to take array as argument
                          For example f(x,y) should be f(A) where A[1]=x, A[2]=y
    Returns scored population. Score is the last element of each chromosome (row)
    in population
    """
    # Enumerate through population, set last element of each chromosome to score
    for (index,member) in enumerate(population)
        population[index][end] = fitness_function(member)
    end
    sort!(population, by = x -> x[end])
    return population
end

function create_next_generation(population::Array{Array{Float64}}, recombRate::Float64)
    """
    Creates children given a population and recombination rate
    Uses an elite selection of 0.10 (rounded up)
    Creates random chance
    TODO: Implement ranking
    """

    popSize = length(population)
    memberSize = length(population[1])
    # Init children array
    children = Array{Array{Float64}}(0)
    # Elite selection
    for i in 1:Int(ceil(popSize*0.10))
        push!(children, population[i])
    end
    
    # Generate rest of population
    while length(children) < popSize
        # Two random children
        choices = sample(1:popSize, 2, replace=false)
        # Check for crossover
        if rand() < recombRate
            # Choose xover point
            # Use -2 since last one is score, and last actual element cant be crossedover since that means there is none
            xover = rand(1:memberSize-2)
            child1 = cat(1,population[choices[1]][1:xover],population[choices[2]][xover+1:end])
            child2 = cat(1,population[choices[2]][1:xover],population[choices[1]][xover+1:end])
        else
            # Else no crossover, just copy
            child1 = population[choices[1]]
            child2 = population[choices[2]]
        end
        # Push child 1 and 2 to children array
        push!(children, child1, child2)
    end

    # We might have one extra due to rounding in the elite selection, let's check
    if length(children) != popSize
        pop!(children)
    end
    return children
end

function mutate(population::Array{Array{Float64}}, mutateRate::Float64, bounds::Array{Tuple{Float64, Float64},1})
    """
    Iterates through each chromosome and mutates if needed
    """

    for member in population
        for i = 1:length(member)-1
            if rand() < mutateRate
                member[i] = rand(bounds[i][1]:0.01:bounds[i][2])
            end
        end
    end

    return population
end

function GA(populationSize::Int, chromosomeSize::Int, fitness_function::Function, bounds::Array{Tuple{Float64, Float64},1})
    """
    Args
        populationSize: Total number of chromosomes
        chromosomeSize: How long each chromosome should be (variables in fitness function)
        fitness_function: Function to determine fitness of each solution
                          Should take an array as an arg
                          Example: f(x,y,z) = x+y+z
                           should be f(X) = X[1]+X[2]+X[3]
                          in order to allow GA to take in any function with any
                          amount of args to the fitness function
        bounds: 1D array of tuples, each tuple is the (lower, upper) bounds
                for each variable. Both lower and upper need to be Float64 types
                Length should match chromosomeSize
                e.g: [(1.0,2.0), (3.5,4.5)]
    """

    # Set recombination and mutation rate, lower and upper bound
    recombRate = 0.7
    mutateRate = 0.05
    maxIterations = 100
    runs = 20
    # First initialize the population

    # Then loop for the required iterations
    # Initialize iteration counter, run counter, and array to hold generations of best population
    i = 0
    run = 0
    bestGenerations = [[[0.0,0.0,Inf]]]
    generations = Array{Array{Array{Float64}}}(0)

    # For each run
    while run < runs
        nextPopulation = initialize_population(populationSize, chromosomeSize, bounds)
        # Iterate through max amount of generations
        while i < maxIterations
            # Score and sort the population
            population = score_sort_population(nextPopulation, fitness_function)
            # Push current population to generations array
            push!(generations, deepcopy(population))
            # Create children
            nextPopulation = create_next_generation(population, recombRate)
            # Mutate children
            nextPopulation = mutate(nextPopulation, mutateRate, bounds)
            i += 1
        end
        #=
        At the end of the run, if the last generations best solution
        # is better than that of the best of all runs so far
        set new best run so we can graph later
        =#
        if generations[end][1][end] < bestGenerations[end][1][end]
            bestGenerations = deepcopy(generations)
        end
        clear!(:generations)
        clear!(:population)
        run += 1
    end

    # Reporting and plotting
    # Report best solution and value
    bestVariables = bestGenerations[end][1][1:end-1]
    bestScore = bestGenerations[end][1][end]
    println("Best solution\nVariables: $bestVariables\nScore: $bestScore")

    # Generate best fitness vs. # generation
    # Init score array
    y = Array{Float64}(0)
    # Get scores, push to array
    for g in bestGenerations
        push!(y,g[1][end])
    end
    # Plot and label
    PyPlot.plot(y)
    xlabel("Generation #")
    ylabel("Fitness")
    # Save and close
    PyPlot.savefig("BestFitness.png")
    PyPlot.close()

    # Contour
    # Check to make sure we only have two variables, otherwise exit
    if length(bestGenerations[1][1]) > 3
        println("Plotting the 4th dimension currently disabled to avoid tearing the space-time continuum")
        quit()
    end

    
    # Pick X and Y range for contour plot based on bounds
    xrange = bounds[1][1]:0.01:bounds[1][2]
    yrange = bounds[2][1]:0.01:bounds[2][2]
    size = length(xrange)
    # Init z array for values of fitness function
    z = zeros(size,size)
    for i in 1:size
        for j in 1:size
            # Generate a z value for each x and y value
            z[i,j] = fitness_function([xrange[i],yrange[j]])
        end
    end
    # If we are at 1st gen or evry 10th, plot
    for (index, gen) in enumerate(bestGenerations)
        if (index == 1) || (mod(index,10) == 0)
            # Re-init x and y for the generation
            x = Array{Float64}(0)
            y = Array{Float64}(0)
            # Push x and y to separate arrays
            for m in gen
                push!(x,m[1])
                push!(y,m[2])
            end
            # Plt the contour plot
            contour(xrange,yrange,z)
            # Then plot the scatter underneath
            scatter(x,y)
            title("Generation $index")
            savefig("contour_gen$index.png")
            close()
        end
    end

end

function fitness_function(X)
    return e - 20*exp(-0.2*sqrt((X[1]^2 + X[2]^2)/2)) - exp((cos(2*pi*X[1]) + cos(2*pi*X[2]))/2)
end

bounds = [(-2.0,2.0), (-2.0,2.0)]
GA(25,2,fitness_function,bounds)