"""
    ParticleSwarm(n_particles = 3,
                  w = 1.0,
                  c1 = 2.0,
                  c2 = 2.0,
                  prob_shift = 0.25,
                  rng = Random.GLOBAL_RNG)

Instantiate a particle swarm optimization tuning strategy. A swarm is initiated
by sampling hyperparameters with their customizable priors, and new models are
generated by referencing each member's and the swarm's best models so far.

### Supported Ranges

A single one-dimensional range or vector of one-dimensional ranges can be
specified. `ParamRange` objects are constructed using the `range` method. If not
paired with a prior, then one is fitted and truncated if bounded, as follows:

| Range Types             | Default Distribution                        |
|:----------------------- |:------------------------------------------- |
| `NominalRange`          | `Dirichlet([1, 1, ..., 1])`                 |
| Bounded `NumericRange`  | `Uniform(lower, upper)`                     |
| Positive `NumericRange` | `Gamma(α=(origin/unit)^2, θ=unit^2/origin`) |
| Other `NumericRange`    | `Normal(origin, unit)`                      |

Specifically, in `ParticleSwarm`, the `range` field of a `TunedModel` instance
can be:

- a single one-dimensional range (`ParamRange` object) `r`

- a pair of the form `(r, d)`, with `r` as above and where `d` is:

    - a Dirichlet distribution with the same number of categories as `r.values`
      (for `NominalRange` `r`)

    - any `Distributions.UnivariateDistribution` *instance* (for `NumericRange`
      `r`)

    - one of the distribution *types* in the table below, for automatic fitting
      using `Distributions.fit(d, r)` to a distribution whose support always
      lies between `r.lower` and `r.upper` (for `NumericRange` `r`)

- any vector of objects of the above form

| Range Types             | Distribution Types                                                                           |
|:----------------------- |:-------------------------------------------------------------------------------------------- |
| Bounded `NumericRange`  | `Arcsine`, `Uniform`, `Biweight`, `Cosine`, `Epanechnikov`, `SymTriangularDist`, `Triweight` |
| Positive `NumericRange` | `Gamma`, `InverseGaussian`, `Poisson`                                                        |
| Any `NumericRange`      | `Normal`, `Logistic`, `LogNormal`, `Cauchy`, `Gumbel`, `Laplace`                             |

### Examples

    using Distributions

    range1 = range(model, :hyper1, lower=0, upper=1)

    range2 = [(range(model, :hyper1, lower=1, upper=10), Arcsine),
              range(model, :hyper2, lower=2, upper=Inf, unit=1, origin=3),
              (range(model, :hyper2, lower=2, upper=4), Normal(0, 3)),
              range(model, :hyper3, values=[:ball, :tree])]

### Algorithm

Hyperparameter ranges are sampled and concatenated into position vectors for
each swarm particle. Velocity is initiated to be zeros, and in each iteration,
every particle's position is updated to approach its personal best and the
swarm's best models so far with the equations:

\$vₖ₊₁ = w⋅vₖ + c₁⋅rand()⋅(pbest - xₖ) + c₂⋅rand()⋅(gbest - xₖ)\$

\$xₖ₊₁ = xₖ + vₖ₊₁\$

New models are then generated for evaluation by mutating the fields of a deep
copy of `model`. If the corresponding range has a specified `scale` function,
then the transformation is applied before the hyperparameter is returned. If
`scale` is a symbol (eg, `:log`), it is ignored.

### Discrete Hyperparameter Handling

Since particle swarm is an optimization method for continuous problems, integer
and nominal hyperparameters require special handling: they are converted to
continuous values, and transformed back to their original domains at each step
for evaluation.

For integer `NumericRange`s, a continuous distribution is fitted to generate
initial values for the swarm. They are then rounded when each particle is mapped
to the corresponding candidate model.

`NominalRange`s on the other hand are represented as categorical distributions
over their values. Hence, we use Dirichlet prior distributions to initialize a
probability vector for each particle, defaulting to the uniform distribution
Dirichlet([1, 1, ..., 1]). The same velocity and position updates apply, but
probability values are further clamped in the range [0, 1] and normalized to sum
up to 1. When a better model is found, we replicate both its probability vector
and sampled value by shifting unchosen categories' weights towards the selected
one for pbest and gbest models:

\$pᵢ = (1 - prob_shift) * pᵢ\$

\$pₛ = pₛ + prob_shift\$

where pₛ is the probability of the sampled hyperparameter value. For more
information, refer to "A New Discrete Particle Swarm Optimization Algorithm" by
Strasser, Goodman, Sheppard, and Butcher.
"""
mutable struct ParticleSwarm{R<:AbstractRNG} <: AbstractParticleSwarm
    n_particles::Int
    w::Float64
    c1::Float64
    c2::Float64
    prob_shift::Float64
    rng::R
    # TODO: topology
end

function ParticleSwarm(;
    n_particles=3,
    w=1.0,
    c1=2.0,
    c2=2.0,
    prob_shift=0.25,
    rng::R=Random.GLOBAL_RNG
) where {R}
    swarm = ParticleSwarm{R}(n_particles, w, c1, c2, prob_shift, rng)
    message = MLJTuning.clean!(swarm)
    isempty(message) || @warn message
    return swarm
end

function MLJTuning.clean!(tuning::ParticleSwarm)
    warning = ""
    if tuning.n_particles < 3
        warning *= "ParticleSwarm requires at least 3 particles. Resetting n_particles=3. "
        tuning.n_particles = 3
    end
    if tuning.w < 0
        warning *= "ParticleSwarm requires w ≥ 0. Resetting w=1. "
        tuning.w = 1
    end
    if tuning.c1 < 0
        warning *= "ParticleSwarm requires c1 ≥ 0. Resetting c1=2. "
        tuning.c1 = 2
    end
    if tuning.c2 < 0
        warning *= "ParticleSwarm requires c2 ≥ 0. Resetting c2=2. "
        tuning.c2 = 2
    end
    if !(0 ≤ tuning.prob_shift < 1)
        warning *= "ParticleSwarm requires 0 ≤ prob_shift < 1. Resetting prob_shift=0.25. "
        tuning.prob_shift = 0.25
    end
    return warning
end

function MLJTuning.setup(tuning::ParticleSwarm, model, ranges, n, verbosity)
    return initialize(ranges, tuning)
end

function MLJTuning.models(
    tuning::ParticleSwarm,
    model,
    history,
    state::ParticleSwarmState{T},
    n_remaining,
    verbosity
) where {T}
    n_particles = tuning.n_particles
    if !isnothing(history)
        sign = MLJTuning.signature(history[1].measure[1])
        measurements = Vector{T}(undef, n_particles)
        map(history[end-n_particles+1:end]) do h
            measurements[h.metadata] = sign * h.measurement[1]
        end
        pbest!(state, measurements, tuning)
        gbest!(state)
        move!(tuning.rng, state, T(tuning.w), T(tuning.c1), T(tuning.c2))
    end
    retrieve!(state, tuning)
    fields = getproperty.(state.ranges, :field)
    new_models = map(1:n_particles) do i
        clone = deepcopy(model)
        for (field, param) in zip(fields, getindex.(state.parameters, i))
            recursive_setproperty!(clone, field, param)
        end
        (clone, i)
    end
    return new_models, state
end

function MLJTuning.tuning_report(tuning::ParticleSwarm, history, state)
    fields = getproperty.(state.ranges, :field)
    scales = MLJBase.scale.(state.ranges)
    return (; plotting = MLJTuning.plotting_report(fields, scales, history))
end