import numpy as np
from scipy import constants
from my_pytools.my_numpy.array import atleast
from abc import ABCMeta, abstractmethod

class integrator:
    """Integrator base class. Inherit from this to define custom predictor/corrector integrators"""
    __metaclass__ = ABCMeta

    def __init__(self):
        pass

    def initialize(self, particles):
        """initialize the integrator with a particle system"""
        self.particles = particles
        self.dt = particles.dt
        self.mass = particles.mass
        self.Iz = particles.Iz
        
    @abstractmethod
    def predict(self): pass

    @abstractmethod
    def correct(self): pass


class verlet(integrator):
    """Verlet integrator (or kick-drift-kick)"""

    def __init__(self):
        pass

    def initialize(self, particles):
        super().initialize(particles)
        self.prevF = np.zeros_like(particles.F)
        self.prevT = np.zeros_like(particles.T)
    
    def predict(self):
        self.particles.velocity += (self.prevF/self.mass[:,np.newaxis])*self.dt/2
        self.particles.angular_velocity += (self.prevT/self.Iz[:,np.newaxis])*self.dt/2

        self.particles.position += self.particles.velocity*self.dt
        self.particles.angle += self.particles.angular_velocity[...,2]*self.dt

    def correct(self):
        self.particles.velocity += (self.particles.F/self.mass[:,np.newaxis])*self.dt/2
        self.particles.angular_velocity += (self.particles.T/self.Iz[:,np.newaxis])*self.dt/2

        self.prevF = np.copy(self.particles.F)
        self.prevT = np.copy(self.particles.T)


#TODO implement interactions=True
class BAOAB(integrator):
    """BAOAB integrator -- a highly stable solver for the Langevin equation"""

    def __init__(self, radii, temperature, viscosity, interactions=False):
        """Initialize the fluid properties
                radii[N]          particle radii array
                temperature       fluid temperature
                viscosity         fluid viscosity
                interactions      include hydrodynamic interactions
        """
        self.radii = atleast(radii, dim=1, length=particles.Nparticles, dtype=float)
        self.temperature = temperature
        self.viscosity = viscosity

    def initialize(self, particles):
        super().initialize(particles)
        self.prevF = np.zeros_like(particles.F)
        self.prevT = np.zeros_like(particles.T)

        alpha_T = 6*np.pi*self.viscosity*self.radii
        alpha_R = 8*np.pi*self.viscosity*self.radii**3
        self.gamma_T = alpha_T/self.mass
        self.gamma_R = alpha_R/self.Iz
        self.gamma_T = self.gamma_T[:,np.newaxis]
        self.gamma_R = self.gamma_R[:,np.newaxis]

    def predict(self):
        # B
        self.particles.velocity += (self.prevF/self.mass[:,np.newaxis])*self.dt/2
        self.particles.angular_velocity += (self.prevT/self.Iz[:,np.newaxis])*self.dt/2

        # A
        self.particles.position += self.particles.velocity*self.dt/2                 
        self.particles.angle += self.particles.angular_velocity*self.dt/2

        # O
        self.particles.velocity = np.exp(-self.gamma_T*self.dt)*self.particles.velocity \
                + np.sqrt(constants.k*self.temperature*(1 - np.exp(-2*self.gamma_T*self.dt))/self.mass[:,np.newaxis])*np.random.normal(size=np.shape(self.prevF))

        self.particles.angular_velocity = np.exp(-self.gamma_R*self.dt)*self.particles.angular_velocity \
          + np.sqrt(constants.k*self.temperature*(1 - np.exp(-2*self.gamma_R*self.dt))/self.Iz[:,np.newaxis])*np.random.normal(size=np.shape(self.prevT))

        # A
        self.particles.position += self.dt/2*self.particles.velocity
        self.particles.angle += self.dt/2*self.particles.angular_velocity

    def correct(self):
        # B
        self.particles.velocity += (self.particles.F/self.mass[:,np.newaxis])*self.dt/2
        self.particles.angular_velocity += (self.particles.T/self.Iz[:,np.newaxis])*self.dt/2

        self.prevF = np.copy(self.particles.F)
        self.prevT = np.copy(self.particles.T)
