"""
This module provides the perturbed PSO from Chen et al. 2012.
The implementation of the algorithm is clear, so we don't need to make any assumptions.
The position sampling does not allow the use of an initial guess x0, so we do not use it with ionBench. Instead, using the ionBench default.
The initial velocity sampling can be used, however.
If the groups are not specified, the case for the presented results, then each parameter is in its own group.
"""
import numpy as np
import ionbench
import itertools


# noinspection PyShadowingNames
def run(bm, x0=None, groups=None, n=6, c1=1.4, c2=1.4, qmax=4, maxIter=1000, w=0.6, debug=False):
    """
    Runs the perturbed particle swarm optimisation algorithm from Chen et al. 2012. If the benchmarker is bounded, the solver will search in the interval [lb,ub], otherwise the solver will search in the interval [0,2*default]

    Parameters
    ----------
    bm : Benchmarker
        A benchmarker to evaluate the performance of the optimisation algorithm.
    x0 : list, optional
        Initial parameter guess. Population is generated by randomly perturbing this initial guess +-50%, then applying the appropriate bounds. If x0=None (the default), then the population will be sampled uniformly between the bounds.
    groups : list, optional
        Groupings of parameters to use in the algorithm. Default is None, in which case every parameter will be in its own group.
    n : int, optional
        Number of particles. The default is 6.
    c1 : float, optional
        Scale of the acceleration towards a particle's best positions. The default is 1.4.
    c2 : float, optional
        Scale of the acceleration towards the best point seen across all particles. The default is 1.4.
    qmax : int, optional
        Maximum number of iterations without improvement before perturbations. The default is 4.
    maxIter : int, optional
        Maximum number of iterations. The default is 1000.
    w : float, optional
        Initial value for the inertia of the particles. The default is 0.6.
    debug : bool, optional
        If True, debug information will be printed, reporting that status of the optimisation each generation. The default is False.

    Returns
    -------
    xbest : list
        The best parameters identified.

    """

    if not bm.parametersBounded:  # pragma: no cover
        raise RuntimeError('This optimiser requires bounds.')

    cost_func = ionbench.utils.cache.get_cached_cost(bm)

    if x0 is None:
        x0 = bm.sample()

    # noinspection PyShadowingNames
    class Particle(ionbench.utils.particle_optimisers.Particle):
        def __init__(self):
            super().__init__(bm, cost_func, x0)

        def set_velocity(self):
            self.velocity = 0.1 * np.random.rand(bm.n_parameters())

    if groups is None:
        groups = [[i] for i in range(bm.n_parameters())]

    q = 0  # Number of generations without improvement
    # Generate patterns
    patterns = []  # All combinations of groups
    for i in range(len(groups)):
        for j in itertools.combinations(range(len(groups)), i + 1):
            newPattern = []
            for k in j:
                newPattern += groups[k]
            patterns += [newPattern]
    N = len(patterns)  # Number of patterns

    particleList = []
    for i in range(n):
        particleList.append(Particle())

    Gcost = [np.inf] * maxIter  # Best cost ever
    Gpos = [None] * maxIter  # Position of best cost ever
    alpha = [None] * maxIter
    L = None
    for L in range(maxIter):
        if L > 0:
            Gcost[L] = Gcost[L - 1]
            Gpos[L] = Gpos[L - 1]

        if debug:
            print('-------------')
            print(f'Beginning population: {L}')
            print(f'Best cost so far: {Gcost[L]}')
            print(f'Found at position: {Gpos[L]}')

        foundImprovement = False
        for p in particleList:
            p.set_cost()
            if p.currentCost < Gcost[L]:
                Gcost[L] = p.currentCost
                Gpos[L] = np.copy(p.position)
                foundImprovement = True

        if foundImprovement:
            q = 0
            if debug:
                print("Found improvement")
                print(f'Best cost is now:  {Gcost[L]}')
                print(f'Found at position: {Gpos[L]}')
        else:
            q += 1
            if debug:
                print("Didn't find an improvement")
                print(f'Current value of q: {q}')

        if bm.is_converged():
            break

        if q > 5 * qmax:  # pragma: no cover
            # Abort
            print("Too many iterations without improvement")
            print(f'Final cost of {Gcost[L]} found at:')
            print(Gpos[L])
            break

        if q >= qmax:
            if debug:
                print("q exceeds qmax so perturbing")
                print(f'q: {q}, qmax: {qmax}')
            # steps 5.1 and 5.2
            newParticleList = []
            bestNewCost = np.inf
            bestNewPosition = None
            for i in range(N):
                newParticle = Particle()
                newParticle.position = np.copy(Gpos[L])
                for j in patterns[i]:
                    newParticle.position[j] *= 1 + (np.random.rand() - 0.5) / 40
                newParticle.clamp()
                newParticle.set_cost()
                if newParticle.currentCost < bestNewCost:
                    bestNewCost = newParticle.currentCost
                    bestNewPosition = np.copy(newParticle.position)
                newParticleList.append(newParticle)
            if debug:
                print("Perturbed particles")
                print(f'Best new cost is: {bestNewCost} found at {bestNewPosition}')
            if bestNewCost <= Gcost[L]:
                if debug:
                    print("Cost improved by perturbing")
                    print(f'Original best cost: {Gcost[L]}, new best cost: {bestNewCost}')
                q = 0
                Gcost[L] = bestNewCost
                Gpos[L] = np.copy(bestNewPosition)
                # Find worst cost and position
                worstCost = -np.inf
                worstPosition = None
                for p in particleList:
                    if p.currentCost > worstCost:
                        worstCost = p.currentCost
                        worstPosition = p.position
                # Found worst cost and position, now to find that particle and set it to the best position
                if debug:
                    print("Adjusting worst particle")
                for p in particleList:
                    if p.currentCost == worstCost and all(p.position == worstPosition):
                        if debug:
                            print("Worst particle found")
                            print(f'Worst particle cost: {worstCost}, worst particle position: {worstPosition}')
                        p.bestCost = bestNewCost
                        p.bestPosition = np.copy(bestNewPosition)
                        if debug:
                            print(f'New best cost for worst particle: {p.bestCost}, position: {p.position}')
            else:
                if debug:
                    print("Cost wasn't improved by perturbing")

        # Step 6
        alpha[L] = 1 / n * np.sum(np.abs([p.currentCost - Gcost[L] for p in particleList]))
        if L >= 2:
            # Adapt inertia weight w from (17)
            w = np.exp(-alpha[L - 1] / alpha[L - 2])
        if debug:
            print("Updating inertia")
            print(f'w: {w}')
            print(f'alpha: {alpha[L]}')

        # Renew velocities according to (16)
        for p in particleList:
            localAcc = c1 * np.random.rand() * (p.bestPosition - p.position)
            globalAcc = c2 * np.random.rand() * (Gpos[L] - p.position)
            p.velocity = w * p.velocity + localAcc + globalAcc
        if debug:
            print("Velocities renewed")
        # Move positions according to (15) while maintaining [0,1] bounds
        for p in particleList:
            p.position += p.velocity
            p.clamp()

        if debug:
            print("Positions renewed")
            print(f'Finished population {L}')
            print(f'Best cost so far: {Gcost[L]}')
            print(f'Found at position: {Gpos[L]}')

    if L >= maxIter-1:
        bm.set_max_iter_flag()

    bm.evaluate()
    return Particle().untransform(Gpos[L])


# noinspection PyUnusedLocal,PyShadowingNames
def get_modification(modNum=1):
    """
    modNum = 1 -> Chen2012

    Returns
    -------
    mod : modification
        Modification corresponding to inputted modNum. Default is modNum = 1, so Chen2012.

    """
    mod = ionbench.modification.Chen2012()
    return mod
