import taichi as ti
import math

@ti.data_oriented
class Motionclass:
    def __init__(self, bucket, mpm: ti.template(), dt: ti.f32):
        self.bucketCenter = bucket["translation"]
        self.scale = bucket["scale"]
        self.cycle = bucket["cycle"]		# [s]
        self.sideStep = bucket["sideStep"]	# [m] y-direction offset intercal 
        self.xLocomotion =  bucket["xLocomotion"]	# [m] 
        self.zLocomotion =  bucket["zLocomotion"]	# [m]
        self.relax =  bucket["relax"]		# relaxation time
        self.dt = dt

        self.bucket_particle_num = 0;		# count elastic material number
        for i in range(mpm.x.shape[0]):
            if mpm.material[i] == mpm.material_elastic:
               self.bucket_particle_num += 1;

        self.bucketX0 = ti.Vector.field(3, ti.f32, shape=self.bucket_particle_num)
        self.particleID = ti.field(dtype=ti.i32, shape=self.bucket_particle_num)
        j = 0
        for i in range(mpm.x.shape[0]):
            if mpm.material[i] == mpm.material_elastic:
               self.bucketX0[j] = mpm.x[i]
               self.particleID[j] = i
               j += 1

    @ti.kernel	# *** One way FSI ***
    def excavation( self, mpm: ti.template(), tt: ti.f32 ):
        PI = ti.math.pi
        center = self.bucketCenter 
        xradius = self.xLocomotion
        zradius = self.zLocomotion 
        cycle = self.cycle	# [s]
        subcycle = cycle/4.0	# [s]
        sideStep = self.sideStep	# [m]

        cycleID = int(tt/cycle)
        offsetY = float(cycleID) * sideStep
        t = tt - cycleID* cycle    
 
        # posture control
        pitch0 =-PI / 4.0
        if t < subcycle: 			# start point to excavation point
            pitch0 = -PI/4.0
        elif t < ( 2.0*subcycle ): 		# excavation and move to loading point
            pitch0 = -PI/4.0 + PI/2.0 * (t-subcycle)/subcycle
        elif t < ( 3.0*subcycle ):		# loading
            pitch0 = PI/4.0 - 3.0*PI/4.0 * (t-2.0*subcycle)/subcycle
        else:					# back to start point
            pitch0 = -PI/2.0 + 2.0 * PI/4.0 * (t-3.0*subcycle)/subcycle
            if -PI/4.0 < pitch0: pitch0 = -PI/4.0     

        pitch1 =-PI / 4.0
        if (t+self.dt) < subcycle: 		# start point to excavation point
            pitch1 = -PI/4.0
        elif (t+self.dt) < ( 2.0*subcycle ): 	# excavation and move to loading point
            pitch1 = -PI/4.0 + PI/2.0 * ((t+self.dt)-subcycle)/subcycle
        elif (t+self.dt) < ( 3.0*subcycle ):	# loading
            pitch1 = PI/4.0 - 3.0*PI/4.0 * ((t+self.dt)-2.0*subcycle)/subcycle
        else:					# back to start point
            pitch1 = -PI/2.0 + 2.0 * PI/4 * ((t+self.dt)-3.0*subcycle)/subcycle
            if  -PI/4.0 < pitch1: pitch1 = -PI/4.0     

        Ry0 = ti.Matrix([ [ ti.cos(pitch0), 0.0, ti.sin(pitch0)],
                          [ 0.0,            1.0, 0.0          ],
                          [-ti.sin(pitch0), 0.0, ti.cos(pitch0)] ])
        Ry1 = ti.Matrix([ [ ti.cos(pitch1), 0.0, ti.sin(pitch1)],
                          [ 0.0,            1.0, 0.0          ],
                          [-ti.sin(pitch1), 0.0, ti.cos(pitch1)] ])

        # locomotion control
        displacement0 = ti.Vector([0.0,0.0,0.0])
        theta0 = 0.0
        if t < subcycle:			# start point to excavation point
           theta0 =-PI/2.0 * t/subcycle
           displacement0 = [-xradius + xradius * ti.math.cos(theta0), offsetY, zradius * ti.math.sin(theta0)]
        elif t < ( 2.0*subcycle ): 		# loading
           theta0 =-PI/2.0 - PI/2.0 * (t-subcycle)/subcycle
           displacement0 = [-xradius + xradius * ti.math.cos(theta0), offsetY, zradius * ti.math.sin(theta0)]
        elif t < ( 3.0*subcycle ): 		# loading
           displacement0 = [-2.0*xradius, offsetY, 0.0]
        else:
           displacement0 = [-2.0*xradius+2.0*xradius*(t-3.0*subcycle)/subcycle, offsetY+sideStep*(t-3.0*subcycle)/subcycle, 0.0]

        displacement1 = ti.Vector([0.0,0.0,0.0])
        theta1 = 0.0
        if (t+self.dt) < subcycle:			# start point to excavation point
           theta1 =-PI/2.0 * (t+self.dt)/subcycle
           displacement1 = [-xradius + xradius * ti.math.cos(theta1), offsetY, zradius * ti.math.sin(theta1)]
        elif (t+self.dt) < ( 2.0*subcycle ): 		# loading
           theta1 =-PI/2.0 - PI/2.0 * ((t+self.dt)-subcycle)/subcycle
           displacement1 = [-xradius + xradius * ti.math.cos(theta1), offsetY, zradius * ti.math.sin(theta1)]
        elif (t+self.dt) < ( 3.0*subcycle ): 		# loading
           displacement1 = [-2.0*xradius, offsetY, 0.0]
        else:
           displacement1 = [-2.0*xradius+2.0*xradius*((t+self.dt)-3.0*subcycle)/subcycle, offsetY+sideStep*((t+self.dt)-3.0*subcycle)/subcycle, 0.0]

        for i in range(self.bucket_particle_num):
            pos = ( self.bucketX0[i] - center )
            p0 = Ry0 @ pos + displacement0 + center
            p1 = Ry1 @ pos + displacement1 + center
            mpm.x[ self.particleID[i] ] = p0
            mpm.v[ self.particleID[i] ] = ( p1 - p0 ) / self.dt
