using StatsBase
#=
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.10: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.10: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 = 1000
    # First initialize the population
    population = initialize_population(populationSize, chromosomeSize, bounds)

    # Then loop for the required iterations
    bestScores = Array{Float64}(0)
    i = 0
    while i < maxIterations
        # Score and sort the population
        population = score_sort_population(population, fitness_function)
        push!(bestScores,population[1][end])
        population = create_next_generation(population, recombRate)
        population = mutate(population, mutateRate, bounds)
        i += 1
    end
    println(bestScores)
    println(population)
end

function ackley(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,ackley,bounds)