import taichi as ti
import math

@ti.data_oriented
class Motionclass:
    def __init__(self, fish, mpm: ti.template(), dt: ti.f32, offset: ti.f32):
        self.translation = fish["translation"]
        self.scale = fish["scale"]
        self.knifePos = fish["knifePos"]	# [m]
        self.knifeSize = fish["knifeSize"]	# [m]
        self.relax =  fish["relax"]		# relaxation time
        self.wallDamping =  fish["wallDamping"] # C of mx'' + cx + kx = f
        self.dt = dt				# [s]
        self.offset = offset			# [m] particle radius


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

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

    @ti.func
    def collision( self, mpm: ti.template(), minX: ti.f32, maxX: ti.f32, minY: ti.f32, maxY: ti.f32, minZ: ti.f32, maxZ: ti.f32 ):
        for i in range(self.fish_particle_num):
            p1 = mpm.x[ self.particleID[i] ]
            # Collision check
            if minX < p1[0] and p1[0] < maxX and minY < p1[1] and p1[1] < maxY and minZ < p1[2] and p1[2] < maxZ:
               flg = 0
               p0 = self.fishX0[i]
               # collision with maxZ
               if p1[2] < maxZ and maxZ <= p0[2]:
                 k = ( maxZ - p0[2] ) / ( p1[2] - p0[2] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minX < pos[0] and pos[0] < maxX and minY < pos[1] and pos[1] < maxY:
                    mpm.x[ self.particleID[i] ][2] = self.fishX0[i][2]		# return back
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][2] = 0.0
                    flg = 1
               # collision with minZ
               if flg == 0 and p0[2] <= minZ and minZ < p1[2]:
                 k = ( minZ - p0[2] ) / ( p1[2] - p0[2] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minX < pos[0] and pos[0] < maxX and minY < pos[1] and pos[1] < maxY:
                    mpm.x[ self.particleID[i] ][2] = self.fishX0[i][2]	# return back
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][2] = 0.0
                    flg = 1
               # collision with maxY
               if flg == 0 and p1[1] < maxY and maxY <= p0[1]:
                 k = ( maxY - p0[1] ) / ( p1[1] - p0[1] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minX < pos[0] and pos[0] < maxX and minZ < pos[2] and pos[2] < maxZ:
                    mpm.x[ self.particleID[i] ][1] = self.fishX0[i][1]	# return back
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][1] = 0.0
                    flg = 1
               # collision with minY
               if flg == 0 and p0[1] <= minY and minY < p1[1]:
                 k = ( minY - p0[1] ) / ( p1[1] - p0[1] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minX < pos[0] and pos[0] < maxX and minZ < pos[2] and pos[2] < maxZ:
                    mpm.x[ self.particleID[i] ][1] = self.fishX0[i][1]	# return back
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][1] = 0.0
                    flg = 1
               # collision with maxX
               if flg == 0 and p1[0] < maxX and maxX <= p0[0]:
                 k = ( maxX - p0[0] ) / ( p1[0] - p0[0] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minY < pos[1] and pos[1] < maxY and minZ < pos[2] and pos[2] < maxZ:
                    mpm.x[ self.particleID[i] ][0] = self.fishX0[i][0]	# return back
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][0] = 0.0
                    flg = 1
               # collision with minX
               if flg == 0 and p0[0] <= minX and minX < p1[0]:
                 k = ( minX - p0[0] ) / ( p1[0] - p0[0] ) 
                 pos = p0 + k * ( p1 - p0 )
                 if minY < pos[1] and pos[1] < maxY and minZ < pos[2] and pos[2] < maxZ:
                    mpm.x[ self.particleID[i] ][0] = self.fishX0[i][0]
                    mpm.v[ self.particleID[i] ] *= ti.exp( - self.wallDamping * self.dt  )
                    mpm.v[ self.particleID[i] ][0] = 0.0
                    flg = 1

    @ti.kernel	# *** One way FSI ***
    def cut( self, mpm: ti.template() ) -> ti.i32:
        minX = self.knifePos[0] - self.knifeSize[0]/2.0 - self.offset
        maxX = self.knifePos[0] + self.knifeSize[0]/2.0 + self.offset
        minY = self.knifePos[1] - self.knifeSize[1]/2.0 - self.offset
        maxY = self.knifePos[1] + self.knifeSize[1]/2.0 + self.offset
        minZ = self.knifePos[2] - self.knifeSize[2]/2.0 - self.offset
        maxZ = self.knifePos[2] + self.knifeSize[2]/2.0 + self.offset
        self.collision( mpm, minX, maxX, minY, maxY, minZ, maxZ )

        flg = 0
        for i in range(self.fish_particle_num):
            self.fishX0[i] = mpm.x[ self.particleID[i] ]

        return flg
