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]
    # Matriz de aceleracion
    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_newton(x_, m, minimize=True, G=0.0001, keep_best=False, weight_method='Linear', t=1.0, max_mass=100.0): 
    """ The procedure that generates the new positions of the particles is based on how bodies are attracted to each other
        through an attractive force (universal law of gravitation). 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 mass based on the solution provided by the particle.
            2.1 The better the solution provided by the particle, the greater the assigned mass.
            2.2 Assignment methods -> Linear, Exponential, Quadratic (further discussion on better assignment methods can be explored).
            2.2.1 In simple terms, a range of solutions is established and adjusted to a charge range between 0 to max_mass according to some assignment method.
        3. Let Fij be the force produced between particle i and j -> Fij = G*mi*mj/(r^2) where r is the distance between the particles, mi and mj are the masses of i and j.
            3.1 Thus, the resulting force on particle i is FRi = Σ{j=1} G*mi*mj/(rij^2).
            3.2 Additionally, by Newton's Law, F = m*a -> a = F/m. Thus, the velocities of each particle are calculated as vf = vi + a*t -> vf = vi + (F/m)*t.
            3.3 Finally, the positions are updated as follows: x' = x + vi*t + (a*t^2)/2 -> x' = x + vi*t + (F/m)*(t^2)/2.
            3.4 Applying the assumption vi=0 -> x' = x + (F/m)*(t^2)/2.

    Args:
        x_ (array(size=(Npar, Nvar)): current particle positions
        m (array(size=(Npar, 1))): current mass of each particle
        minimize (bool, optional): solver objective. Defaults to True.
        G (float, optional): gravitational constant. Defaults to 0.0001.
        keep_best (bool, optional): It will save the best value obtained in each iteration. Defaults to False.
        weight_method (str, optional): method to reassign mass. Defaults to 'Linear'. Options=['Linear', 'Quadratic', 'Exponential'].
        t (float, optional): time. Defaults to 1.0.
        max_mass (float, optional): upper bound of mass to assing. Defaults to 100.0.

    Returns:
        F_total, Ac, new_pos: returns the force obtained on each particle, their acceleration, 
                              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 F_=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 a mass magnitude based on the solution provided by the particle.
    m=np.array(m)
    min_value = m[0]
    max_value = m[-1]

    if minimize:
        m = -1.0*(m-max_value-1.0)
        max_value = m[0]

    if weight_method=="Linear":
        # We adjust the mass according to the following range [0, max_mass].
        m = (m/max_value)*max_mass
        reverse_ind = [len(m)-i-1 for i in range(len(m))] 
        m = m[reverse_ind]
        
    elif weight_method=="Quadratic": 
        # We adjust the mass according to the following range [0, max_mass].
        m_2=m**2
        m=(m_2/np.max(m_2))*max_mass
        reverse_ind = [len(m)-i-1 for i in range(len(m))] 
        m = m[reverse_ind]
        
    elif weight_method=="Exponential": 
        # We adjust the mass according to the following range [0, max_mass].
        m_exp=np.array([np.exp(i) for i in m])
        m=(m_exp/np.max(m_exp))*max_mass
        reverse_ind = [len(m)-i-1 for i in range(len(m))] 
        m = m[reverse_ind]
        
    m = np.nan_to_num(m, nan=0.0001)
    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 attraction force F_ is calculated.            
            d_2=r_2[ind_r][ind_c]
            if d_2==0:
                F_=0.0
            else:
                m1=m[ind_r]
                m2=m[ind_c]
                F_=float(G*m1*m2/d_2)
            F[ind_r][ind_c]=F_*d[ind_r][ind_c]
    F_total = np.sum(F, axis=1)
    # Note: Between two particles, the attractive forces will be the same, but not their accelerations.
    # Remember that F = M * A -> m1 * a1 = m2 * a2, so if m1 > m2 -> a2 > a1 -> Particles with greater mass have lower acceleration.Ac=acc(F_total, m)
    Ac=acc(F_total, m)
    Ac = np.nan_to_num(Ac, nan=0.0)
    
    # Finally, Xf = xi + vi * t + 0.5 * acc * (t^2) (Final Position), but if the particles always start at rest (vi=0) -> xf = 0.5 * acc * (t^2).
    x__ = 0.5*Ac*(t**2)
    new_pos=x_ + x__
    
    if keep_best==True:
        best_ind = np.argmin(m)
        Ac[best_ind]=0.0
        new_pos=x_ + 0.5*Ac*(t**2)
    
    # 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, new_pos