import numpy as np

def normaliazed_direction(ori, des):
    dir_ = des-ori
    nor_dir_ = dir_/np.sqrt(np.sum(dir_**2))
    nor_dir_=np.nan_to_num(nor_dir_)
    return nor_dir_

def distance(ori, des):
    dis_ = des-ori
    return dis_

def acc(F, m):
    Npar=F.shape[0]
    Nvar=F.shape[1]
    # Acceleration Matrix
    Ac = np.full((Npar, Nvar), 0.0)
    for ind_r, row in enumerate(Ac):
        for ind_c, col in enumerate(row):
            Ac[ind_r][ind_c]=F[ind_r][ind_c]/m[ind_r]
    return Ac

def F_coulomb(x, v, q, P, minimize=True, k=0.0001, weight_method='Linear', t=1.0, max_q=100.0): 
    """ The procedure that generates the new positions of the particles is based on how particles with positive and negative charges move
        in electric fields. Thus, the procedure will be as follows:
        
        1. n particles are created at random positions within the search space.
        2. Each particle is assigned a positive charge type.
        3. The magnitude of each particle's charge is directly related to the objective function.
        3.1 Assignment Methods -> Linear, Exponential, Quadratic (better assignment methods can be explored)
        3.2 Simply put, a solution range is established and adjusted to a charge range between 0 to max_q according to some assignment method.
        4. Each particle holds a velocity. These start at rest but will change over the iterations of the algorithm.
        5. The P particles with the best values obtained from the iteration will remain at rest and emit an electric field of opposite sign (negative) attracting the rest.
        5.1 Let E = k*Q/r^2 be the magnetic field at a point located at a distance r from the source.
        5.2 Let Fe be the electric force between the magnetic field and the particle -> Fe = E*q0 where q0 is the particle's charge.
        5.3 Let Fij be the force produced between particle i and j -> Fij = k*qi*qj/(r^2) where r is the distance between the particles, qi and qj are the charges of i and j.
        5.4 Thus, the resultant forces of particle i are FRi = sum_{p=1}^P k*Qp*Qi/(rip^2) + sum_{j=P+1}^n k*Qi*Qj/(rij^2) 
        5.5 Moreover, by Newton's Law, F=m*a -> Fe=m*a. Thus, the velocities of each particle are calculated vf=vi+a*t -> vf = vi + (FRi/m)*t
        5.6 The positions are updated as follows x = vi*t + (a*t^2)/2 -> x = vi*t + (FRi/m)*(t^2)/2
        
        Finally, if the particles are too close, it could cause a problem. If rij^2 << 0, then the electric force would be too strong, leading
        the particles to abruptly separate from each other. To avoid this behavior, a small distance (0.00000001) is established. If the distance
        between two particles is smaller than that distance, it will be considered that both particles have collided. Then, if more than half of 
        the particles have collided, we will apply a random function to generate new solutions over a reduced search space.
    
    Args:
        x (array(size=(Npar, Nvar)): current particle positions
        v (array(size=(Npar, Nvar)): current particle velocity
        q (array(size=(Npar, 1))): current electric charge of each particle
        P (int): quantity of electric fields (the best P particles become electric fields)
        minimize (bool, optional): solver objective. Defaults to True.
        k (float, optional): electric constant. Defaults to 0.0001.
        weight_method (str, optional): method to reassign electric charges. Defaults to 'Linear'. Options=['Linear', 'Quadratic', 'Exponential'].
        t (float, optional): time. Defaults to 1.0.
        max_q (float, optional): upper bound of electric charge to assing. Defaults to 100.0.

    Returns:
        F_total, Ac, vf, new_pos: returns the force obtained on each particle, their acceleration, their velocities
                                  and their new positions
    """

    Npar=x.shape[0] 
    Nvar=x.shape[1]

    # Distance Matrix
    dis = []
    for ind_r, row in enumerate(x):
        for ind_c, col in enumerate(x):
            dis.append(distance(x[ind_r], x[ind_c]))
    dis=np.array(dis).reshape((Npar,Npar,Nvar))

    # Direction Matrix
    d = []
    for ind_r, row in enumerate(x):
        for ind_c, col in enumerate(x):
            d.append(normaliazed_direction(x[ind_r], x[ind_c]))
    d=np.array(d).reshape((Npar,Npar,Nvar))
    
    colisioned_part = []
    r_2 = np.zeros((Npar, Npar))
    for ind_r, row in enumerate(dis):
        for ind_c, col in enumerate(dis):
            value = dis[ind_r][ind_c]
            value_2 = value**2
            value_sum = np.sum(value_2)
            if value_sum < 0.00000001: # Particles have practically collided. Notice later that fe=0.0 ahead.
                r_2[ind_r][ind_c] = 0.0
                if ind_r != ind_c:
                    colisioned_part.append(ind_c)
            else:
                r_2[ind_r][ind_c] = value_sum
    colisioned_part_ = np.unique(np.array(colisioned_part))


    # Each particle is assigned an electric charge magnitude based on the solution provided by the particle.
    q=np.array(q)
    min_value = q[0]
    max_value = q[-1]
    if minimize:
        q = -1.0*(q-max_value-1.0)
        max_value = q[0]

    if weight_method=="Linear":
        # We adjust the charges according to the following range [0, max_q].
        q = (q/max_value)*max_q
        reverse_ind = [len(q)-i-1 for i in range(len(q))] # It is inverted to give more charge to the particles that are further away.
        q = q[reverse_ind]
        
    elif weight_method=="Quadratic":
        # We adjust the charges according to the following range [0, max_q].
        q_2=q**2
        q=(q_2/np.max(q_2))*max_q
        reverse_ind = [len(q)-i-1 for i in range(len(q))] # It is inverted to give more charge to the particles that are further away.
        q = q[reverse_ind]
        
    elif weight_method=="Exponential":
        # We adjust the charges according to the following range [0, max_q].
        q_exp=np.array([np.exp(i) for i in q])
        q=(q_exp/np.max(q_exp))*max_q
        reverse_ind = [len(q)-i-1 for i in range(len(q))] # It is inverted to give more charge to the particles that are further away.
        q = q[reverse_ind]
        
    Npar=d.shape[0]
    Nvar=d.shape[2]
    F=np.full((Npar, Npar, Nvar), 0.0)
    for ind_r, row in enumerate(dis):
        for ind_c, col in enumerate(row):
            # The magnitude of the electric force Fe is calculated.            
            d_2=r_2[ind_r][ind_c]
            if d_2==0:
                Fe=0.0
            else:
                q1=q[ind_r]
                q2=q[ind_c]
                Fe=float(k*q1*q2/d_2)
                if ind_r >= P and ind_c >= P: # Repulsive forces are generated between particles of the same sign.
                    F[ind_r][ind_c]=-1.0*Fe*d[ind_r][ind_c]
                else: # There is attraction between particles and electric fields.
                    F[ind_r][ind_c]=Fe*d[ind_r][ind_c]
    F[:P, :P] = 0.0
    F_total = np.sum(F, axis=1)
    F_total[:P]=0.0
    # Remember that F=ma -> m1*a1 = m2*a2. So, if m1>m2 -> a2>a1 -> Particles with greater mass have less acceleration.
    # For this method, the weight of the particle is not important, so we will set them all to be equal to 1.0.
    m=np.ones(Npar)
    # vf = acc + vi*t
    Ac=acc(F_total, m)
    vf = Ac + t*v
    Ac[:P]=0.0
    vf[:P]=0.0 # The velocity of the magnetic fields is set to 0.0.
    # Finally, Xf = xi + vi*t + 0.5*acc*(t^2) (Final Position)
    x__=v*t + 0.5*Ac*(t**2)
    new_pos = x + x__
    
    # If more than half of the particles have collided, those that have collided will move randomly within the reduced search space.
    max_crash_part=Npar/2
    if len(colisioned_part_)>max_crash_part:
        lo = np.min(new_pos, axis=0)
        up = np.max(new_pos, axis=0)
        for i in colisioned_part_:
            new_pos[i] = np.random.uniform(low=lo, high=up, size=(1, Nvar))

    return F_total, Ac, vf, new_pos