from __future__ import division
import numpy as np

from pySDC.Problem import ptype
from pySDC.datatype_classes.particles import particles, fields, acceleration

class penningtrap(ptype):
    """
    Example implementing a single particle in a penning trap

    Attributes:
        nparts: number of particles (needs to be 1 here)
    """

    def __init__(self, cparams, dtype_u, dtype_f):
        """
        Initialization routine

        Args:
            cparams: custom parameters for the example
            dtype_u: particle data type (will be passed parent class)
            dtype_f: acceleration data type (will be passed parent class)
        """

        # these parameters will be used later, so assert their existence
        assert 'omega_B' in cparams # B field frequency
        assert 'omega_E' in cparams # E field frequency
        assert 'u0' in cparams      # initial position and velocity
        assert 'nparts' in cparams  # number of particles
        assert 'sig' in cparams     # smoothing parameter for Coulomb interaction

        # add parameters as attributes for further reference
        for k,v in cparams.items():
            setattr(self,k,v)

        # set number of particles
        self.nparts = cparams['nparts']
        # invoke super init, passing nparts, dtype_u and dtype_f
        super(penningtrap,self).__init__(self.nparts, dtype_u, dtype_f)


    def get_interactions(self,part):
        """
        Routine to compute the particle-particle interaction, assuming q = 1 for all particles

        Args:
            part: the particles
        Returns:
            the internal E field for each particle

        """

        N = self.nparts

        Efield = np.zeros(3*N)

        for i in range(N):
            for j in range(N):
                dist2 = np.linalg.norm(part.pos.values[3*i:3*i+3]-part.pos.values[3*j:3*j+3],2)**2+self.sig**2
                contrib = part.q[j]*(part.pos.values[3*i:3*i+3]-part.pos.values[3*j:3*j+3]) / dist2**(3/2)

                Efield[3*i  ] += contrib[0]
                Efield[3*i+1] += contrib[1]
                Efield[3*i+2] += contrib[2]

        return Efield


    def eval_f(self,part,t):
        """
        Routine to compute the E and B fields (named f for consistency with the original PEPC version)

        Args:
            t: current time (not used here)
            part: the particles
        Returns:
            Fields for the particles (internal and external)
        """

        N = self.nparts

        Emat = np.diag([1,1,-2])
        f = fields(self.nparts)

        f.elec.values = self.get_interactions(part)


        for n in range(N):
            f.elec.values[3*n:3*n+3] += self.omega_E**2 / (part.q[n]/part.m[n]) * np.dot(Emat,part.pos.values[
                                                                                              3*n:3*n+3])
            f.magn.values[3*n:3*n+3] = self.omega_B * np.array([0,0,1])

        return f


    def u_init(self):
        """
        Routine to compute the starting values for the particles

        Returns:
            particle set filled with initial data
        """

        u0 = self.u0
        N = self.nparts

        u = particles(N)

        if u0[2][0] is not 1 or u0[3][0] is not 1:
            print('Error: so far only q = m = 1 is implemented (I think)')
            exit()