import BioMetaheuristics.TestFunctions: Benchmark, evaluate

"""
    PSO <: PopulationBase

`PSO` is the type associated with the implementation for the
Particle Swarm Optimization with momentum. See [`Algorithms`](@ref implementations-docs)
for more information.
"""
struct PSO <: PopulationBase end

# To enable dispatch based on the type
function optimize(f, population, k_max, rng, ::PSO; kwargs...)
    return PSO(f, population, k_max, rng; kwargs...)
end

"""
    PSO(f::Function, population, k_max::Int, rng;
        w=0.9, c1=2.0, c2=2.0
    ) -> OptimizationResults
    PSO(f::Benchmark, population, k_max::Int, rng;
        w=0.9, c1=2.0, c2=2.0
    ) -> OptimizationResults

Method that implements `PSO` for a function `f` of type `Function`
or of type `Benchmark`.

Returns an `OptimizationResults` type with information relevant to the
run executed, see [`OptimizationResults`](@ref).

# Arguments

- `population`: can be any `AbstractArray` that contains [`Particle`](@ref)
instances, but it is expected to be generated by [`Population`](@ref).
- `k_max`: number of maximum iterations until "convergence" of the algorithm.
- `rng`: An object of type `AbstractRNG`.

# Keyword arguments

_It is recommended to use the default values provided._

- `w`: value that controls how much of the initial velocity is retained, i.e. an inertia term.
This value decays linearly over each iteration until it reaches
the default miminum value of 0.4.
- `c1`: balance between the influence of the individual's knowledge, i.e. the best inidividual solution so far.
- `c2`: balance between the influence of the population's knowledge, i.e. the best global solution so far.
- `seed`: an integer to be used as the seed for the pseudo random number generators.
If `nothing` is passed (the default), then a random seed will be taken from the system.

# Examples
```julia
using BioMetaheuristics
using Random

rng = MersenneTwister(1261234)
# Define the Sphere function
f_sphere(x) = sum(x.^2)

# Implement PSO for a 3-dimensional Sphere function, with
# 10000 iterations and 30 particles in the population.
val = PSO(f_sphere, Population(30, 3, -15.0, 15.0), 10_000, rng)
```
"""
function PSO(f::Function, population, k_max::Int, rng; w=0.9, c1=2.0, c2=2.0)

    val = _pso!(f, population, k_max, rng; w=w, c1=c1, c2=c2)

    optim_res = OptimizationResults(val, f(val), "PSO", k_max)

    return optim_res
end

function PSO(f::Benchmark, population, k_max::Int, rng; w=0.9, c1=2.0, c2=2.0)
    val = _pso!(
        x -> evaluate(f, x),
        population,
        k_max,
        rng;
        w=w,
        c1=c1,
        c2=c2
    )

    optim_res = OptimizationResults(val, evaluate(f, val), "PSO", k_max)

    return optim_res
end

function _pso!(f, population, k_max::Int, rng; w=0.9, c1=2.0, c2=2.0)
    # Obtain weight decay rate
    η = _weight_decay(w, k_max)

    # Evaluate initial costs
    dimension = length(population[1].x)

    x_best = similar(population[1].x_best) # Best candidate
    y_best = Inf # Best miminum
    @inbounds for particle in population
        y = f(particle.x)
        if y < y_best
            copyto!(x_best, particle.x)
            y_best = y
        end
    end

    # PSO main loop
    @inbounds for _ in 1:k_max
        _update!(f, population, w, c1, c2, dimension, x_best, y_best, rng)
        # Make the inertia weight decay over time
        w -= η
    end

    # Sort the solutions and return the first one
    return population[1].x_best
end

function _update!(f, population, w, c1, c2, n, x_best, y_best, rng)
    # Create container for random numbers
    rngs = zeros(2, n)

    for P in population
        rand!(rng, rngs)
        # Evaluate velocity
        P.v = (w * P.v) + (c1 * rngs[1, :] .* (P.x_best .- P.x))
        P.v .+= c2 * rngs[2, :] .* (x_best .- P.x)
        # Update position
        P.x .+= P.v
        # Apply boundary values to positions and velocities
        _clip_positions_velocities!(P)
        # Update values if they give lower cost
        y = f(P.x)
        if y < y_best
            copyto!(x_best, P.x)
            y_best = y
        end
        if y < f(P.x_best)
            copyto!(P.x_best, P.x)
        end
    end

    return nothing
end

"""
    Compute the corresponding weight decay depending the maximum
number of iterations and the initial value for it.
"""
@inline function _weight_decay(initial, itr_max)
    # Following the references, the minimum is 0.4
    stop = 0.4
    step_size = (initial - stop) / itr_max

    return step_size
end
