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 = 60 	# [s]
CalculationSpace = 4 	# [m]
Resolution = 64
dt = 2e-4
fileName = "001bucket.obj"
VTUinterval = 100		# VTU saving interval 
bucket = {
        "translation": [3.2, 1, 1.4],
        "scale": [10.0, 10.0, 10.0],
        "cycle": 30.0,		# [s]
        "sideStep": 1.0,	# [m] y-direction offset intercal 
        "xLocomotion": 1.1,	# [m]
        "zLocomotion": 0.9,	# [m]
        "relax": 0.1
        }

ti.init(arch=ti.cpu)

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

mpm.add_cube(lower_corner=[0, 0, 0], cube_size=[4, 3.2, 0.5], material=MPMSolver.material_sand)
mpm.add_CADmodel(fileName, offset=bucket["translation"], scale=bucket["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(-2, -10, 10)
camera.lookat(1, 1, 1)
camera.up(0, 0, 1)
camera.fov(16)
particles_radius = CalculationSpace / Resolution / 4.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))

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.5, 1.5, 0.5), color=(0.7, 0.7, 0.7))
    scene.point_light(pos=(-0.5, 1.5, 1.5), color=(0.7, 0.7, 0.7))
    draw_ground(scene)	# ground z = 0 plane
    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],
    [236.0 / 255.0, 84.0 / 255.0, 59.0 / 255.0, 1.0],
    [1.0, 1.0, 1.0, 1.0],
    [1.0, 1.0, 0.0, 1.0] ] )

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

vtu = VTUclass()
motion = Motionclass(bucket, mpm, dt)

t = 0.0
cnt = 0
cnt2 = 0

while t<SimulationTime:
    print(f"time={t:.3f}")
    motion.excavation(mpm, t)
    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)
