using Distributions, Statistics, LinearAlgebra
using DataStructures: CircularBuffer

struct AFS <: AbstractOptimizer
    population_size::Int
    δ_init::Real
    λ::Real
    θ::Real
    m::Int
    η::Real
    step::Real
    mutation::Function

    AFS(; populationSize::Int = 100,
        δ_init::Real = 1.0,
        λ::Real = 1e-3,
        θ::Real = 1.0,
        m::Int = 5,
        η::Real = 1e-3,
        step = 1.0,
        mutation::Function = uniform(1.0),
        ) =
        new(populationSize, δ_init, λ, θ, m, η, step, mutation)
end

population_size(method::AFS) = method.population_size
default_options(method::AFS) = (iterations = 1000, abstol = 1e-32)
summary(m::AFS) = "($(m.δ_init)/$(m.λ)/$(m.θ)/$(m.η)/$(m.m)/$(m.step))"
show(io::IO, m::AFS) = print(io, summary(m))

mutable struct AFSState{T, IT} <: AbstractOptimizerState
    n::Int
    δ::Real
    board::Vector
    fitness::Vector{T}
    fittest::IT
    best_fitnesses::CircularBuffer{T}
    i::Int
end

value(s::AFSState) = minimum(s.fitness)
minimizer(s::AFSState) = s.fittest

function get_norms(population)
    return norm.(population)
end

function initial_state(method::AFS, options, objfun, population)
    T = typeof(value(objfun))
    n = length(first(population))
    fitness = fill(maxintfloat(T), method.population_size)
    return AFSState(n, copy(method.δ_init), deepcopy(population), fitness, population[1], CircularBuffer{T}(method.m), 0)
end

function get_centroid(School::Vector{Vector{T}}) where T<:AbstractFloat
    l = length(School)
    sum(School) ./ l
end

function scan_visual_field(i::Int, field::T, School::Vector{Vector{IT}}) where {T<:AbstractFloat, IT<:AbstractFloat}
    V = @view School[setdiff(1:length(School), i)]
    findall([norm(School[i] - V[j]) for j = 1:length(V)] .> field)
end

function random_move!(S::Vector{Vector{T}}, board::Vector{Vector{T}}, n::Int, field::IT, fish::Int, f::Function) where {T<:AbstractFloat, IT<:AbstractFloat}
    board[fish] = S[fish] + f(zeros(n)) .* field
end

function search!(S::Vector{Vector{T}}, board::Vector{Vector{T}}, n::Int, fish::Int, V::Vector, step::IT, f::Function) where {T<:AbstractFloat, IT<:AbstractFloat}
    j = rand(1:length(V))
    d = S[V[j]] - S[fish]
    board[fish] = S[fish] + step * f(zeros(n)) .* d ./ abs.(d)
end

function swarm!(S::Vector{Vector{T}}, board::Vector{Vector{T}}, n::Int, fish::Int, centroid::Vector{T}, step::IT, f::Function) where {T<:AbstractFloat, IT<:AbstractFloat}
    d = centroid - S[fish]
    board[fish] = S[fish] + step * f(zeros(n)) .* d ./ abs.(d)
end

function chase!(S::Vector{Vector{T}}, board::Vector{Vector{T}}, n::Int, fish::Int, neighbor::Int, step::IT, f::Function) where {T<:AbstractFloat, IT<:AbstractFloat}
    d = S[neighbor] - S[fish]
    board[fish] = S[fish] + step * f(zeros(n)) .* d ./ abs.(d)
end

function leap!(S::Vector{Vector{T}}, fish::Int, n::Int, max_norm::IT, f::Function) where {T<:AbstractFloat, IT<:AbstractFloat}
    S[fish] = S[fish] + f(zeros(n)) .* (max_norm * 2 .- S[fish])
end

function eval_fitness(constraints, objfun, x)
    value(objfun, x) + penalty(constraints, x)
end

function update_state!(objfun, constraints, state, population::AbstractVector{IT}, method::AFS, itr) where {IT}
    state.i += 1
    max_norm = first(findmax(get_norms(population)))
    v = state.δ * max_norm
    for fish = 1:method.population_size
        V = scan_visual_field(fish, v, population)
        if isempty(V)
            random_move!(population, state.board, state.n, v, fish, method.mutation)
        elseif length(V) / method.population_size > method.θ
            search!(population, state.board, state.n, fish, V, method.step, method.mutation)
        else
            c = get_centroid(population[V])
            state.fitness[fish] = eval_fitness(constraints, objfun, population[fish])
            if eval_fitness(constraints, objfun, c) < state.fitness[fish]
                swarm!(population, state.board, state.n, fish, c, method.step, method.mutation)
            else
                search!(population, state.board, state.n, fish, V, method.step, method.mutation)
            end
            best_neighbor = findmin(state.fitness[V])
            if first(best_neighbor) < state.fitness[fish]
                chase!(population, state.board, state.n, fish, V[last(best_neighbor)], method.step, method.mutation)
            else
                search!(population, state.board, state.n, fish, V, method.step, method.mutation)
            end
        end
    end

    for fish = 1:method.population_size
        o = apply!(constraints, state.board[fish])
        y_fit = value(objfun, o) + penalty(constraints, o)
        if y_fit < state.fitness[fish]
            population[fish] = o
            state.fitness[fish] = y_fit
        end
    end

    minfit = findmin(state.fitness)
    state.fittest = population[last(minfit)]
    pushfirst!(state.best_fitnesses, first(minfit))

    if state.i >= method.m
        stagnation = last(state.best_fitnesses) - first(state.best_fitnesses)
        if stagnation < method.η
            j = rand(1:method.population_size)
            leap!(population, j, state.n, max_norm, method.mutation)
            apply!(constraints, population[j])
        end
        if state.best_fitnesses[end] <= state.best_fitnesses[end-1]
            state.δ *= exp(-method.λ)
        end
    end
    
    return false

end