import taichi as ti
import numpy as np
import utils
import os
import glob

from moviepy import ImageSequenceClip

from engine.mpm_solver import MPMSolver
from VTU import VTUclass
from motion import Motionclass

# write_to_disk = False

# Parameters
SimulationTime = 0.6 	# [s]
CalculationSpace = 0.6 	# [m]
Resolution = 128
dt = 2e-4
fileName = "fishbody.obj"
VTUinterval = 10		# VTU saving interval 
fish = {
        "translation": [CalculationSpace/2.0, CalculationSpace/2.0, CalculationSpace/2.0+0.05],
        "scale": [1.0, 1.0, 1.0],
        "knifePos": [CalculationSpace/2.0, CalculationSpace/2.0, CalculationSpace/2.0-0.1],
        "knifeSize": [0.002, 0.12, 0.1],
        "wallDamping": 1.0,
        "relax": 0.1
        }

ti.init(arch=ti.cpu)

mpm = MPMSolver(res=(Resolution, Resolution, Resolution), size=CalculationSpace, max_num_particles=2 ** 19, use_ggui=True)

# mpm.add_cube(lower_corner=[0, 0, 0], cube_size=[0.6, 0.6, 0.2], material=MPMSolver.material_water)
mpm.add_CADmodel(fileName, offset=fish["translation"], scale=fish["scale"], material=MPMSolver.material_elastic, size_ratio=0.25)

mpm.set_gravity((0, 0, -9.8))

@ti.kernel
def set_color(ti_color: ti.template(), material_color: ti.types.ndarray(), ti_material: ti.template()):
    for I in ti.grouped(ti_material):
        material_id = ti_material[I]
        color_4d = ti.Vector([0.0, 0.0, 0.0, 1.0])
        for d in ti.static(range(3)):
            color_4d[d] = material_color[material_id, d]
        ti_color[I] = color_4d

res = (1024, 760)
window = ti.ui.Window("Real MPM 3D", res, vsync=True)
canvas = window.get_canvas()
scene = ti.ui.Scene()
camera = ti.ui.make_camera()
camera.position(-0.5, -1, 1)
camera.lookat(0.2, 0.2, 0.2)
camera.up(0, 0, 1)
camera.fov(24)
particles_radius = CalculationSpace / Resolution / 8.0

# ground lines
N_ground = 20
ground_vbo = ti.Vector.field(3, dtype=ti.f32, shape=4*(N_ground+1))
ground_verts = np.zeros((4*(N_ground+1), 3), dtype=np.float32)
for i in range(N_ground+1):
    x = i * CalculationSpace / N_ground
    ground_verts[4*i+0] = [x, 0, 0]
    ground_verts[4*i+1] = [x, CalculationSpace, 0]
    ground_verts[4*i+2] = [0, x, 0]
    ground_verts[4*i+3] = [CalculationSpace, x, 0]

def draw_ground(scene):
    ground_vbo.from_numpy(ground_verts)
    scene.lines(ground_vbo, width=1, color=(0.1,0.1,0.1))

knife_vbo = ti.Vector.field(3, dtype=ti.f32, shape=8)
knife_verts = np.zeros((8, 3), dtype=np.float32)
dx = fish["knifeSize"][0]/2.0
dy = fish["knifeSize"][1]/2.0
dz = fish["knifeSize"][2]/2.0

knife_verts[0] = [-dx+fish["knifePos"][0], -dy+fish["knifePos"][1], -dz+fish["knifePos"][2]]
knife_verts[1] = [ dx+fish["knifePos"][0], -dy+fish["knifePos"][1], -dz+fish["knifePos"][2]]
knife_verts[2] = [ dx+fish["knifePos"][0],  dy+fish["knifePos"][1], -dz+fish["knifePos"][2]]
knife_verts[3] = [-dx+fish["knifePos"][0],  dy+fish["knifePos"][1], -dz+fish["knifePos"][2]]
knife_verts[4] = [-dx+fish["knifePos"][0], -dy+fish["knifePos"][1],  dz+fish["knifePos"][2]]
knife_verts[5] = [ dx+fish["knifePos"][0], -dy+fish["knifePos"][1],  dz+fish["knifePos"][2]]
knife_verts[6] = [ dx+fish["knifePos"][0],  dy+fish["knifePos"][1],  dz+fish["knifePos"][2]]
knife_verts[7] = [-dx+fish["knifePos"][0],  dy+fish["knifePos"][1],  dz+fish["knifePos"][2]]

knife_indices = ti.field(ti.i32, shape=36)
def draw_knife(scene):
    knife_vbo.from_numpy(knife_verts)
    knife_indices.from_numpy(np.array([
              0,1,2, 0,2,3,	# bottom
              4,5,6, 4,6,7, 	# top
              0,1,5, 0,5,4,	# front
              2,3,7, 2,7,6,	# back
              0,3,7, 0,7,4,	# left
              1,2,6, 1,6,5,	# right
              ], dtype=np.int32))
    scene.mesh(knife_vbo, indices=knife_indices, color=(0.9,0.9,0.9))

def render():
    camera.track_user_inputs(window, movement_speed=0.03, hold_key=ti.ui.RMB)
    scene.set_camera(camera)
    scene.ambient_light((0, 0, 0))
    set_color(mpm.color_with_alpha, material_type_colors, mpm.material)
    scene.particles(mpm.x, per_vertex_color=mpm.color_with_alpha, radius=particles_radius)
    scene.point_light(pos=(0.1, 0.1, 0.2), color=(0.7, 0.7, 0.7))
    scene.point_light(pos=(0.1, 0.1, 1.0), color=(0.7, 0.7, 0.7))
    draw_ground(scene)	# ground z = 0 plane
    draw_knife(scene)
    canvas.scene(scene)

def show_options():
    global particles_radius

    window.GUI.begin("Solver Property", 0.05, 0.1, 0.2, 0.10)
    window.GUI.text(f"Current particle number {mpm.n_particles[None]}")
    particles_radius = window.GUI.slider_float("particles radius ", particles_radius, 0, 0.1)
    window.GUI.end()

    window.GUI.begin("Camera", 0.05, 0.3, 0.3, 0.16)
    camera.curr_position[0] = window.GUI.slider_float("camera pos x", camera.curr_position[0], -10, 10)
    camera.curr_position[1] = window.GUI.slider_float("camera pos y", camera.curr_position[1], -10, 10)
    camera.curr_position[2] = window.GUI.slider_float("camera pos z", camera.curr_position[2], -10, 10)

    camera.curr_lookat[0] = window.GUI.slider_float("camera look at x", camera.curr_lookat[0], -10, 10)
    camera.curr_lookat[1] = window.GUI.slider_float("camera look at y", camera.curr_lookat[1], -10, 10)
    camera.curr_lookat[2] = window.GUI.slider_float("camera look at z", camera.curr_lookat[2], -10, 10)

    window.GUI.end()

material_type_colors = np.array([
    [0.1, 0.1, 1.0, 0.8],				# water
    [0.9, 0.9, 1.0, 0.8],				# elastic
    [1.0, 1.0, 1.0, 1.0],				# snow
    [1.0, 1.0, 0.0, 1.0] ] )				# sand

os.makedirs("image", exist_ok=True)
# os.makedirs("VTU", exist_ok=True)

vtu = VTUclass()
motion = Motionclass(fish, mpm, dt, particles_radius)

t = 0.0
cnt = 0
cnt2 = 0

while t<SimulationTime:
    print(f"time={t:.3f}")
    flg = motion.cut(mpm)
    if flg == 1: input()
    mpm.step(dt)

    if cnt%VTUinterval == 0:
       render()
       # show_options()
       window.save_image(f"./image/frame{cnt2:04}.png")  
       window.show()

       # vtu.write_vtuMPMASCII(f"./VTU/SPH{cnt:04}.vtu", mpm )
       cnt2 += 1

    t += dt
    cnt += 1

# for movie 
files = sorted(glob.glob("./image/frame*.png"))
if len(files) > 0:
   clip = ImageSequenceClip(files, fps=60)
   clip.write_videofile("out.mp4", codec="libx264", audio=False)
