import os
import argparse
import taichi as ti
import numpy as np
from config_builder import SimConfig
from particle_systemkikut import ParticleSystem
from moviepy import ImageSequenceClip
import glob
from VTU import VTUclass

# ti.init(arch=ti.gpu, device_memory_fraction=0.5)

ti.init(arch=ti.cpu)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='SPH Taichi')
    parser.add_argument('--scene_file', default='', help='scene file')
    args = parser.parse_args()
    scene_path = args.scene_file
    config = SimConfig(scene_file_path=scene_path)
    vtu = VTUclass()
    scene_name = scene_path.split("/")[-1].split(".")[0]
    substeps = config.get_cfg("numberOfStepsPerRenderUpdate")
    output_frames = config.get_cfg("exportFrame")
    dt = config.get_cfg("timeStepSize")
    simulationTime = config.get_cfg("simulationTime")
    series_prefix = "{}_output/particle_object_{}.ply".format(scene_name, "{}")
    if output_frames:
        os.makedirs("image", exist_ok=True)
    os.makedirs("VTU", exist_ok=True)

    ps = ParticleSystem(config, GGUI=True)	# particle variables
    solver = ps.build_solver()			# 4 means DFSPH
    solver.initialize()
    solver.viscosity = 0.001			# water = 0.001 [Pa s]. default is 0.01

    # Invisible objects
    invisible_objects = config.get_cfg("invisibleObjects")
    if not invisible_objects:
        invisible_objects = []

    if output_frames:
        window = ti.ui.Window('SPH', (800, 600), show_window = True, vsync=False)
        scene = ti.ui.Scene()
        camera = ti.ui.Camera()
        camera.position(2.0, 2.0, 4.0)
        camera.up(0.0, 0.0, 1.0)
        camera.lookat( 0.0, 0.0, -0.5)
        camera.fov(10)
        scene.set_camera(camera)

        canvas = window.get_canvas()
        movement_speed = 0.02
        background_color = (0, 0, 0)  # 0xFFFFFF
        particle_color = (1, 1, 1)


    # Draw the lines for domain
    x_max, y_max, z_max = config.get_cfg("domainEnd")
    box_anchors = ti.Vector.field(3, dtype=ti.f32, shape = 8)
    box_anchors[0] = ti.Vector([0.0, 0.0, 0.0])
    box_anchors[1] = ti.Vector([0.0, y_max, 0.0])
    box_anchors[2] = ti.Vector([x_max, 0.0, 0.0])
    box_anchors[3] = ti.Vector([x_max, y_max, 0.0])

    box_anchors[4] = ti.Vector([0.0, 0.0, z_max])
    box_anchors[5] = ti.Vector([0.0, y_max, z_max])
    box_anchors[6] = ti.Vector([x_max, 0.0, z_max])
    box_anchors[7] = ti.Vector([x_max, y_max, z_max])

    box_lines_indices = ti.field(dtype=ti.i32, shape=(2 * 12))	# dtype=ti.i32 means "int"

    for i, val in enumerate([0, 1, 0, 2, 1, 3, 2, 3, 4, 5, 4, 6, 5, 7, 6, 7, 0, 4, 1, 5, 2, 6, 3, 7]):
        box_lines_indices[i] = int(val)

    cnt = 0
    t = 0

    while t < simulationTime:
        print(f"time={t:.3f}")
        for i in range(substeps):	# solver
            solver.step()
            t += dt;
        ps.copy_to_vis_buffer(invisible_objects=invisible_objects)

        # visualize for window   
        if output_frames:
            camera.track_user_inputs(window, movement_speed=movement_speed, hold_key=ti.ui.LMB)
            scene.set_camera(camera)
            scene.point_light((2.0, 2.0, 2.0), color=(1.0, 1.0, 1.0))
            scene.particles(ps.x_vis_buffer, radius=ps.particle_radius, per_vertex_color=ps.color_vis_buffer)
            scene.lines(box_anchors, indices=box_lines_indices, color = (0.99, 0.68, 0.28), width = 1.0)
            canvas.scene(scene)
            window.save_image(f"./image/frame{cnt:06}.png")      
            window.show()

        # for paraview visualization
        pos, density, vel, pressure = vtu.extract_pos_density_vel_pressure(ps)
        object_id = vtu._to_numpy(ps.object_id)
        vtu.write_vtu_with_material(f"./VTU/SPH{cnt:06}.vtu", pos, density=density, velocity=vel, pressure=pressure, object_id=object_id)

        cnt += 1

    # collect frames and write movie
    if output_frames:
        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)
