"""
A hybrid method of Nelder-Mead and PSO.
It uses a population of 3m+1 particles (where m = number of parameters) in which the best m+1 particles take a Nelder-Mead simplex step and the remaining 2m particles follow a PSO step.
The initial position sampling procedure is only partially defined. m+1 particles are generated like a NM simplex, but the remaining 2m particles are only defined as being generated randomly. We assume ionBench defaults for these particles.
No initial velocity sampling procedure is defined. We assume ionBench defaults.
For points generated by NM (accepting or shrinking), their velocity (in case they are demotes to the PSO step) is not defined. We assume a new randomly sampled velocity for new accepted points and the particles previous velocity for a simplex shrink step.
A particle's best position is treated similarly. Reset for accept and maintained for shrinking.
The PSO step is not defined. We assume a fairly typical setup with fixed acceleration coefficients at 1.496.
Guidance for the choice of eps is not given so a value of 1e-6 is used.
"""
import numpy as np
import ionbench


# noinspection PyShadowingNames
def run(bm, x0=None, maxIter=1000, eps=1e-6, debug=False):
    """
    Run the hybrid Nelder-Mead and PSO optimiser from Liu et al. 2011.

    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%, clamped to bounds if necessary. If x0=None (the default), then the population will be sampled uniformly between the bounds.
    maxIter : int, optional
        Maximum number of iterations. The default is 1000.
    eps : float
        Convergence criterion, roughly equivalent to a typical ftol criterion. The default is 0.01.
    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 x0 is None:
        x0 = bm.sample()

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

    # noinspection PyShadowingNames
    class Particle(ionbench.utils.particle_optimisers.Particle):
        def __init__(self, position=None, transform=False):
            super().__init__(bm, cost_func, x0)
            if position is None:
                # Initial sampling (using util.particle_optimisers.Particle default)
                self.set_position(x0)
            else:
                # Force to be position (position is in input space, store in transformed [0,1] space)
                # This is only used for the initial population, after that the Nelder Mead and PSO are automatically using the transformed space.
                if transform:
                    self.position = self.transform(position)
                else:
                    self.position = np.copy(position)
            # Clamp to bounds
            self.clamp()

    # class for the simplex
    # noinspection PyShadowingNames
    class Simplex:
        def __init__(self, particles):
            self.particles = particles

        def get_best(self):
            # returns the best point in the simplex
            x_best = self.particles[0]
            for p in self.particles:
                if p.currentCost < x_best.currentCost:  # pragma: no cover
                    x_best = p
            return x_best

        def get_worst(self):
            # returns the worst point in the simplex
            x_worst = self.particles[0]
            for p in self.particles:
                if p.currentCost > x_worst.currentCost:
                    x_worst = p
            return x_worst

        def get_second_worst(self):
            # returns second-worst point in the simplex
            x_worst = self.get_worst()
            x_secondWorst = self.particles[0]
            for p in self.particles:
                if p.currentCost > x_secondWorst.currentCost and p != x_worst:
                    x_secondWorst = p
            return x_secondWorst

        def centroid(self, x_worst):
            # find the centroid of the simplex by averaging the positions, not including x_worst. Don't generate a point, cost not needed
            cen = [0] * bm.n_parameters()
            for p in self.particles:
                if p != x_worst:
                    cen += p.position
            cen /= len(self.particles) - 1
            return cen

        def accept(self, x):
            x_worst = self.get_worst()
            for i in range(len(self.particles)):
                if self.particles[i] == x_worst:
                    self.particles[i].position = x.position
                    self.particles[i].set_cost()
                    return

        def step(self):
            """
            Take a modified nelder mead step with acceptance probability related to the current temperature.
            """
            x_worst = self.get_worst()
            x_best = self.get_best()
            x_secondWorst = self.get_second_worst()
            c = self.centroid(x_worst)
            # attempt reflection
            xr = Particle(position=2 * c - x_worst.position)
            xr.set_cost()
            if x_secondWorst.currentCost > xr.currentCost >= x_best.currentCost:
                if debug:
                    print("Accept reflection")
                self.accept(xr)
                return
            if xr.currentCost < x_best.currentCost:
                # Attempt expand
                xe = Particle(position=2 * xr.position - c)
                xe.set_cost()
                if xe.currentCost < xr.currentCost:
                    if debug:
                        print("Accept expansion")
                    self.accept(xe)
                else:
                    if debug:
                        print("Ignore expansion. Accept reflection")
                    self.accept(xr)
                return
            if xr.currentCost >= x_secondWorst.currentCost:
                # contract
                if debug:
                    print("Attempt contraction")
                if x_worst.currentCost > xr.currentCost >= x_secondWorst.currentCost:
                    # outside contract
                    xc = Particle(position=(xr.position + c) / 2)
                    xc.set_cost()
                    if xc.currentCost <= xr.currentCost:
                        if debug:
                            print("Accept inside contraction")
                        self.accept(xc)
                        return
                else:
                    # inside contract
                    xc = Particle(position=(x_worst.position + c) / 2)
                    xc.set_cost()
                    if xc.currentCost < x_worst.currentCost:
                        if debug:
                            print("Accept inside contraction")
                        self.accept(xc)
                        return
            # shrink
            if debug:
                print("Shrink")
            for i in range(len(self.particles)):
                if self.particles[i] != x_best:
                    self.particles[i].position = (x_best.position + self.particles[i].position) / 2
                    self.particles[i].set_cost()
            return

        def stop(self):
            """
            Check if the simplex has converged.
            """
            costs = np.array([p.currentCost for p in self.particles])
            costs = costs[1:]
            Sc = np.sqrt(np.mean((np.mean(costs) - costs) ** 2))
            return Sc < eps

    L = 0
    assert bm.in_parameter_bounds(x0)
    assert bm.in_rate_bounds(x0)
    # Set population size n
    m = bm.n_parameters()
    n = 3 * m + 1
    # Initial population
    particleList = []
    for i in range(n):
        if i < m + 1:
            if i == 0:
                particleList.append(Particle(position=x0, transform=True))
            else:
                perturb = np.ones(m)
                perturb[i - 1] = 1.05
                particleList.append(Particle(position=x0 * perturb, transform=True))
        else:
            particleList.append(Particle())

    Gcost = [np.inf] * maxIter  # Best cost ever
    Gpos = [None] * maxIter  # Position of best cost ever
    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]}')

        for p in particleList:
            p.set_cost()
            if p.currentCost < Gcost[L]:
                Gcost[L] = p.currentCost
                Gpos[L] = np.copy(p.position)
        if L == 0:
            assert particleList[0].currentCost < 1e5
        # Sort particle list by cost
        costs = [p.currentCost for p in particleList]
        particleList = [p for _, p in sorted(zip(costs, particleList), key=lambda pair: pair[0])]

        # Build simplex and take 1 step
        s = Simplex(particleList[:m + 1])
        s.step()
        particleList[:m + 1] = s.particles

        # Renew velocities
        c1 = 1.496  # Assume they used the fixed value
        c2 = 1.496  # Assume they used the fixed value
        for p in particleList[m + 1:]:
            localAcc = c1 * np.random.rand() * (p.bestPosition - p.position)
            globalAcc = c2 * np.random.rand() * (Gpos[L] - p.position)
            p.velocity = p.velocity + localAcc + globalAcc
        if debug:
            print("Velocities renewed")
        # Move positions
        for p in particleList[m + 1:]:
            p.position += p.velocity
            # Enforce bounds by clamping
            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 bm.is_converged() or s.stop():
            break

    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 -> Liu2011

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

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