from numpy import ndarray, array, sqrt, sum, sign
from random import random

cdef extern from "math.h":

    double c_sqrt "sqrt"(double x)

cdef extern from "arrayobject.h":

    ctypedef int intp

    ctypedef extern class numpy.ndarray [object PyArrayObject]:
        cdef char *data
        cdef int nd
        cdef intp *dimensions
        cdef intp *strides
        cdef int flags

cdef void LJ_potential(ndarray p1p, ndarray p2p, ndarray rp):
    cdef double *p1 = <double *>p1p.data
    cdef double *p2 = <double *>p2p.data
    cdef double l_vec[3]
    cdef int i
    for i in range(3):
        l_vec[i] = p2[i] - p1[i]
    cdef double l
    l = c_sqrt(l_vec[0]*l_vec[0] + l_vec[1]*l_vec[1] + l_vec[2]*l_vec[2])
    cdef double *r = <double *>rp.data
    l = (1/l**6 - 1/l**12) / l
    for i in range(3):
        r[i] += l_vec[i] * l

cdef advance_one(ndarray r_oldp, ndarray r_currp, ndarray vp, ndarray fp, double h):
    """
    Advances just one sphere.

    r_old, r_curr .... old and current position (array([0.343, 0.432, 0.534]))
    v ... current speed
    f .... force of all the other spheres (array([0.343, 0.432, 0.534]))
    """
    cdef double *r_old = <double *>r_oldp.data
    cdef double *r_curr = <double *>r_currp.data
    cdef double *v = <double *>vp.data
    cdef double *f = <double *>fp.data
    cdef int i
    cdef double limit = 1.0
    cdef double f_l
    f_l = c_sqrt(f[0]*f[0] + f[1]*f[1] + f[2]*f[2])
    if f_l > limit:
        for i in range(3):
            f[i] = f[i] * limit/f_l
    cdef ndarray r_newp = array([0., 0., 0.])
    cdef double *r_new = <double *>r_newp.data
    for i in range(3):
        r_old[i] = r_curr[i] - v[i]*h
        r_new[i] = 2*r_curr[i] - r_old[i]+h**2*f[i]
    v_new = (r_newp - r_currp)/h
    return r_newp, v_new

def temperature(speeds, h):
    m = 1.
    T = sum([m*v**2/2 for v in speeds])
    return T

def advance(r, v, f, h):
    """
    Advances the positions r[t] to r[t+h] and returns it.

    f ... the forces of other spheres
    v ... all speeds
    h ... time step
    """
    r_old = r[-2]
    r_curr = r[-1]
    new = [advance_one(old, curr, vcurr, ff, h) for old, curr, vcurr, ff in
            zip(r_old, r_curr, v[-1], f)]
    r_new = []
    v_new = []
    for r, v in new:
        r_new.append(r)
        v_new.append(v)
    return r_new, v_new

def assign_speeds(speeds, sigma=0.1):
    from random import gauss
    mu = 0.
    new = []
    for v in speeds:
        new.append(array([gauss(mu, sigma), gauss(mu, sigma), gauss(mu, sigma)]))
    return new

cdef double a(double x, double min, double max):
    cdef double d = max - min
    while x < min:
        x += d
    while x > max:
        x -= d
    return x

def fit_in_the_box(pos, n):
    "Applies periodic boundary conditions to pos"
    # calculate the size of the box so that rho = 0.8
    rho = 0.8
    b = (n/rho)**(1./3) /2
    #xmin, xmax, ymin, ymax, zmin, zmax:
    box = [-b, b, -b, b, -b, b]
    r = []
    for p in pos:
        r.append(array([
            a(p[0], box[0], box[1]),
            a(p[1], box[2], box[3]),
            a(p[2], box[4], box[5])
            ]))
    return r

cdef double LJ_potential_scalar(ndarray p1p, ndarray p2p):
    cdef double *p1 = <double *>p1p.data
    cdef double *p2 = <double *>p2p.data
    cdef double l_vec[3]
    cdef int i
    for i in range(3):
        l_vec[i] = p2[i] - p1[i]
    cdef double l
    l = c_sqrt(l_vec[0]*l_vec[0] + l_vec[1]*l_vec[1] + l_vec[2]*l_vec[2])
    l = (1/l**6 - 1/l**12)
    return l

def get_f(ndarray p, spheres_pos):
    "Calculates the force exerted on 'p' by all the other spheres"
    cdef double fi = 0
    cdef ndarray tmp
    for tmp in spheres_pos:
        if not tmp is p: #we don't want to count "p"
            fi += LJ_potential_scalar(p, tmp)
    return fi

def get_f2(ndarray p, spheres_pos):
    "Calculates the force exerted on 'p' by all the other spheres"
    cdef ndarray fi = array([0., 0., 0.])
    cdef ndarray tmp
    for tmp in spheres_pos:
        if not tmp is p: #we don't want to count "p"
            LJ_potential(p, tmp, fi)
    return fi

def calculate_potential(pos):
    pot = 0.
    for x in pos:
        pot += get_f(x, pos)
    return pot/len(pos)