"""
@package pmsco.optimizers.swarm
particle swarm optimization handler.

the module starts multiple MSC calculations and optimizes the model parameters
according to the particle swarm optimization algorithm.

Particle swarm optimization adapted from
D. A. Duncan et al., Surface Science 606, 278 (2012)

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-18 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import logging
import numpy as np
import pmsco.optimizers.population as population
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)


class SwarmPopulation(population.Population):
    """
    particle swarm population.
    """

    ## @var friends
    # number of other particles that each particle consults for the global best fit.
    # default = 3.

    ## @var momentum
    # momentum of the particle.
    # default = 0.689343.

    ## @var attract_local
    # preference for returning to the local best fit
    # default = 1.92694.

    ## @var attract_global
    # preference for heading towards the global best fit.
    # default = 1.92694

    def __init__(self):
        """
        initialize the population object.

        """
        super(SwarmPopulation, self).__init__()

        self.friends = 3
        self.momentum = 0.689343
        self.attract_local = 1.92694
        self.attract_global = 1.92694
        self.position_constrain_mode = 'default'
        self.velocity_constrain_mode = 'default'

    def advance_population(self):
        """
        advance the population by one step.

        this method just calls advance_particle() for each particle of the population.
        if generation is lower than zero, the method increases the generation number but does not advance the particles.

        @return: None
        """
        if not self._hold_once:
            self.generation += 1
            for index, __ in enumerate(self.pos):
                self.advance_particle(index)

        super(SwarmPopulation, self).advance_population()

    def advance_particle(self, index):
        """
        advance a particle by one step.

        @param index: index of the particle in the population.
        """

        # note: the following two identifiers are views,
        # assignment will modify the original array
        pos = self.pos[index]
        vel = self.vel[index]
        # best fit that this individual has seen
        xl = self.best[index]
        # best fit that a group of others have seen
        xg = self.best_friend(index)

        for key in self.model_start:
            # update velocity
            dxl = xl[key] - pos[key]
            dxg = xg[key] - pos[key]
            pv = np.random.random()
            pl = np.random.random()
            pg = np.random.random()
            vel[key] = (self.momentum * pv * vel[key] +
                self.attract_local * pl * dxl +
                self.attract_global * pg * dxg)
            pos[key], vel[key], self.model_min[key], self.model_max[key] = \
                self.constrain_velocity(pos[key], vel[key], self.model_min[key], self.model_max[key],
                                        self.velocity_constrain_mode)
            # update position
            pos[key] += vel[key]
            pos[key], vel[key], self.model_min[key], self.model_max[key] = \
                self.constrain_position(pos[key], vel[key], self.model_min[key], self.model_max[key],
                                        self.position_constrain_mode)

        self.update_particle_info(index)

    # noinspection PyUnusedLocal
    def best_friend(self, index):
        """
        select the best fit out of a random set of particles

        returns the "best friend"
        """
        friends = np.random.choice(self.best, self.friends, replace=False)
        index = np.argmin(friends['_rfac'])
        return friends[index]


class ParticleSwarmHandler(population.PopulationHandler):
    """
    model handler which implements the particle swarm optimization algorithm.

    """

    def __init__(self):
        super(ParticleSwarmHandler, self).__init__()
        self._pop = SwarmPopulation()