# Covariance Matrix Adaptation Evolution Strategy
# ==================
#
# Implementation: (μ/μ_I,λ)-CMA-ES
#
# μ is the number of parents
# λ is the number of offspring.
#
@with_kw struct CMAES <: Optimizer
    N::Int
    initPopulation::Individual = ones(N)
    initStrategy::Strategy = strategy(τ = sqrt(N), τ_c = N^2, τ_σ = sqrt(N))
    μ::Integer = 1
    λ::Integer = 1
end

function optimize(objfun::Function, opt::CMAES;   
                    iterations::Integer = 1_000,
                    tol::Float64 = 1e-10,
                    verbose = false)
    @unpack N,initPopulation,initStrategy,μ,λ = opt
    @assert μ < λ "Offspring population must be larger then parent population"

    # Initialize parent population
    individual = getIndividual(initPopulation, N)
    population = fill(individual, μ)
    offspring = Array{typeof(individual)}(undef, λ)
    fitpop = fill(Inf, μ)
    fitoff = fill(Inf, λ)

    parent = copy(individual)
    C = Diagonal{Float64}(I, N)
    s = zeros(N)
    s_σ = zeros(N)
    σ = 1.0
    E = zeros(N, λ)
    W = zeros(N, λ)

    # Generation cycle
    count = 0
    while true
        SqrtC = (C + C') / 2.0
        try
            SqrtC = cholesky(SqrtC).U
        catch
            println("Break on Cholesky")
            break
        end

        for i in 1:λ
            # offspring are generated by transforming standard normally distributed random vectors using a transformation matrix
            E[:,i] = randn(N)
            W[:,i] = σ * (SqrtC * E[:,i])
            offspring[i] = parent + W[:,i]   # (L1)
            fitoff[i] = objfun(offspring[i]) # Evaluate fitness
        end

        # Select new parent population
        idx = sortperm(fitoff)[1:μ]
        for i in 1:μ
            population[i] = offspring[idx[i]]
            fitpop[i] = fitoff[idx[i]]
        end

        w = vec(mean(W[:,idx], dims=2))
        ɛ = vec(mean(E[:,idx], dims=2))
        parent += w            #  forming recombinant perent for next generation (L2)
        s = (1.0 - 1.0/initStrategy[:τ])*s +
            (sqrt(μ/initStrategy[:τ] * (2.0 - 1.0/initStrategy[:τ]))/σ)*w      # (L3)
        C = (1.0 - 1.0/initStrategy[:τ_c]).*C + (s./initStrategy[:τ_c])*s'     # (L4)
        s_σ = (1.0 - 1.0/initStrategy[:τ_σ])*s_σ +
            sqrt(μ/initStrategy[:τ_σ]*(2.0 - 1.0/initStrategy[:τ_σ]))*ɛ        # (L5)
        σ = σ*exp(((s_σ'*s_σ)[1] - N)/(2*N*sqrt(N)))                           # (L6)

        # termination condition
        count += 1
        if count == iterations || σ <= tol
            break
        end
        verbose && println("BEST: $(fitpop[1]): $(σ)")
    end

    return population[1], fitpop[1], count
end
