import taichi as ti
import math

@ti.data_oriented
class Motionclass:
    def __init__(self, fish):
        for i, rb in enumerate(fish):
            self.fishCenter = tuple( rb["translation"] )
            self.fishLength = float( rb["fishLength"] )
            self.fishWidth = float( rb["fishWidth"] )
            self.wingspan = float( rb["wingspan"] )
            self.wingchord = float( rb["wingchord"] )
            self.LEPos = float( rb["LEPos"] )			# leading edge x pos = distance from fish center to leadigedge
            self.pitchAngle = math.radians( float(rb["pitchAngle"]) )
            self.waveLength = float( rb["waveLength"] )
            self.fishVel = tuple( rb["velocity"] )
            self.fishAmp =  float( rb["amplitude"] )
            self.fishFreq =  float( rb["freq"] )
            self.relax =  float( rb["relax"] )		# relaxation time
            self.startPos = tuple( rb["startPosition"] )
            self.swimPos = tuple( rb["swimmingPosition"] )


    @ti.kernel	# *** One way FSI ***
    def carangiform( self, object_id: ti.template(), fx: ti.template(), fx0: ti.template(), fv: ti.template(), n: ti.i32, t: ti.f32, dt: ti.f32 ):
        PI = ti.math.pi
        center = ti.Vector( self.fishCenter )
        amp = self.fishAmp
        length = self.fishLength
        waveL = self.waveLength
        wingL = self.wingspan/2.0 - self.fishWidth/2.0
        freq = self.fishFreq
        cycle = 1.0/freq
        startPos = ti.Vector( self.startPos )
        swimPos = ti.Vector( self.swimPos )
        relax = self.relax
        width = self.fishWidth
        pitch = self.pitchAngle * ti.min( 1.0, t/relax )
        Ry = ti.Matrix([ [ ti.cos(pitch), 0.0, ti.sin(pitch)],
                         [ 0.0,           1.0, 0.0          ],
                         [-ti.sin(pitch), 0.0, ti.cos(pitch)] ])
        
        for i in range(n):
            if object_id[i] == 1:	# if fish=rigid body: 0 means fluid, Caution: i is renumbered.
                bp0  = fx0[i] - center	# fish coordinate system
                bp1  = fx0[i] - center	# fish coordinate system
                disY = ti.abs( bp0.y ) - width / 2.0
                # ********  current position ***********
                dy = amp / length * ( bp0.x + length/2.0 ) * ti.math.sin( 2.0 * PI * ( ( bp0.x + length/2.0 ) / waveL - freq * t ) )
                if 0.0 < disY:
                    dy *= ti.math.pow( 1.0 - disY/wingL, 2.0 ) 
                bp0.y += ( dy * ti.min( 1.0, t/relax ) )
                # ************* future position *************
                dy = amp / length * ( bp0.x + length/2.0 ) * ti.math.sin( 2.0 * PI * ( ( bp0.x + length/2.0 ) / waveL - freq * (t+dt) ) )
                if 0.0 < disY:
                    dy *= ti.math.pow( 1.0 - disY/wingL, 2.0 ) 
                bp1.y += ( dy * ti.min( 1.0, t/relax ) )	# not accurate

                # wing deployment            
                if ( 0.0 < disY ) and ( (-0.1*self.wingchord) < (bp0.x + self.LEPos) ) and ( (bp0.x + self.LEPos) < (1.1*self.wingchord) ):
                  wp0 = bp0					# fish coordinate system
                  wp1 = bp1					# fish coordinate system
                  if 0.0 < wp0.y: 	# right wing coordinate system
                      wp0 += ti.Vector([self.LEPos, -width/2.0, 0.0])	# shift for right wing
                      wp1 += ti.Vector([self.LEPos, -width/2.0, 0.0])	# shift for right wing
                  else:			# left wing coordinate system
                      wp0 += ti.Vector([self.LEPos,  width/2.0, 0.0])	# shift for left wing
                      wp1 += ti.Vector([self.LEPos,  width/2.0, 0.0])	# shift for left wing
                  lag0 = 0.8 * ti.atan2( wp0.y, wp0.x ) 		# ti.atan2(y, x), ti.atan2(0, 0)->0.0
                  if t < relax:
                     lag0 *= ( ( 1.0 - ti.math.cos( PI * t/relax ) ) / 2.0 )
                  elif t < (relax+2*cycle):
                     lag0 *= 1.0
                  elif t < (relax+3*cycle):
                     lag0 *= ( ( 1.0 + ti.math.cos( PI * (t-relax-2*cycle)/cycle ) ) / 2.0 )
                  else:
                     lag0 = 0.0
                  lag1 = 0.8 * ti.atan2( wp1.y, wp1.x )  			# ti.atan2(y, x), ti.atan2(0, 0)->0.0
                  if (t+dt) < relax:
                     lag1 *= ( ( 1.0 - ti.math.cos( PI * (t+dt)/relax ) ) / 2.0 )
                  elif (t+dt) < (relax+2*cycle):
                     lag1 *= 1.0
                  elif (t+dt) < (relax+3*cycle):
                     lag1 *= ( ( 1.0 + ti.math.cos( PI * (t+dt-relax-2*cycle)/cycle ) ) / 2.0 )
                  else:
                     lag1 = 0.0
                  Rz0 = ti.Matrix([ [ti.cos(-lag0), -ti.sin(-lag0), 0.0],
                                    [ti.sin(-lag0),  ti.cos(-lag0), 0.0],
                                    [0.0,           0.0,          1.0] ])
                  Rz1 = ti.Matrix([ [ti.cos(-lag1), -ti.sin(-lag1), 0.0],
                                    [ti.sin(-lag1),  ti.cos(-lag1), 0.0],
                                    [0.0,           0.0,          1.0] ])
                  bp0 = Rz0 @ wp0 
                  bp1 = Rz1 @ wp1 
                  if 0.0 < wp0.y:
                      bp0 -= ti.Vector([self.LEPos, -width/2.0, 0.0])
                      bp1 -= ti.Vector([self.LEPos, -width/2.0, 0.0])
                  else:
                      bp0 -= ti.Vector([self.LEPos,  width/2.0, 0.0])
                      bp1 -= ti.Vector([self.LEPos,  width/2.0, 0.0])

                # posture
                fx[i] = Ry @ bp0 + center
                Fx    = Ry @ bp1 + center
		
                # Locomotion
                if t < relax:
                  fx[i] += ( startPos - center ) * ti.min( 1.0, (1.0 - ti.cos(PI*t/relax))/2.0 )
                  Fx    += ( startPos - center ) * ti.min( 1.0, (1.0 - ti.cos(PI*t/relax))/2.0 )	# not accurate
                elif t < (relax+2*cycle): 
                  fx[i] += ( ( startPos - center ) + ( swimPos - startPos ) * ( 1.0 - ti.cos(PI*(t-relax)/(2.0*cycle)))/2.0 )
                  Fx    += ( ( startPos - center ) + ( swimPos - startPos ) * ( 1.0 - ti.cos(PI*(t-relax)/(2.0*cycle)))/2.0 )	# not accurate
                else:
                  fx[i] += ( swimPos - center )
                  Fx    += ( swimPos - center )	# not accurate



                # ***** velocity ******** accuracy is very low because of "float"
                fv[i] = ( Fx - fx[i] ) / dt
             
