# Example for a forcefield implementation in Pyrex.
# Use Pyrex 0.9.3 or later.


#
# Get all the required declarations
#
include 'forcefield.pxi'

#
# The force field term implementation.
# The rules:
#
# - The class must inherit from EnergyTerm.
#
# - EnergyTerm.__init__() must be called with the arguments
#   shown here. The third argument is the name of the EnergyTerm
#   object, the fourth a tuple of the names of all the terms it
#   implements (one object can implement several terms).
#   The assignment to self.eval_func is essential, without it
#   any energy evaluation will crash.
#
# - The function "evaluate" must have exactly the parameter
#   list given in this example.
#
cdef class HarmonicOscillatorTerm(EnergyTerm):

    cdef int atom_index
    cdef double ref_x, ref_y, ref_z
    cdef double k

    def __init__(self, universe, atom, reference, force_constant):
        cdef ArrayType ref_array
        EnergyTerm.__init__(self, universe,
                            "harmonic_oscillator", ("harmonic_oscillator",))
        self.eval_func = <void *>HarmonicOscillatorTerm.evaluate
        self.atom_index = atom.index
        ref_array = reference.array
        self.ref_x = (<double *>ref_array.data)[0]
        self.ref_y = (<double *>ref_array.data)[1]
        self.ref_z = (<double *>ref_array.data)[2]
        self.k = force_constant

        print self.test1()

    cdef test1(self):
        a=0
        for i in[1,2,3]:
           a=a+i
        return a


    # The function evaluate is called for every single energy
    # evaluation and should therefore be optimized for speed.
    # Its first argument is the global energy evaluator object,
    # which is needed only for parallelized energy terms.
    # The second argument is a C structure that contains all the
    # input data, in particular the particle configuration.
    # The third argument is a C structure that contains the
    # energy term fields and gradient arrays for storing the results.
    # For details, see MMTK_forcefield.pxi.
    cdef void evaluate(self, EnergyEvaluator eval,
                       energy_spec *input, energy_data *energy):
        cdef vector3 *coordinates, *gradients
        cdef double *fc
        cdef double dx, dy, dz
        cdef int n, offset
        coordinates = <vector3 *>input.coordinates.data
        dx = coordinates[self.atom_index][0] - self.ref_x
        dy = coordinates[self.atom_index][1] - self.ref_y
        dz = coordinates[self.atom_index][2] - self.ref_z

        print self.test1()


        energy.energy_terms[self.index] = 0.5*self.k*(dx*dx + dy*dy + dz*dz)
        energy.virial_available = 0 # we don't do virials here
        if energy.gradients != NULL:
            gradients = <vector3 *>(<PyArrayObject *> energy.gradients).data
            gradients[self.atom_index][0] = gradients[self.atom_index][0] + self.k*dx
            gradients[self.atom_index][1] = gradients[self.atom_index][1] + self.k*dy
            gradients[self.atom_index][2] = gradients[self.atom_index][2] + self.k*dz
        if energy.force_constants != NULL:
            fc = <double *>(<PyArrayObject *> energy.force_constants).data
            n = (<PyArrayObject *> energy.force_constants).dimensions[0]
            offset = (9*n+3)*self.atom_index
            fc[offset + 3*n*0 + 0] = fc[offset + 3*n*0 + 0] + self.k
            fc[offset + 3*n*1 + 1] = fc[offset + 3*n*1 + 1] + self.k
            fc[offset + 3*n*2 + 2] = fc[offset + 3*n*2 + 2] + self.k