import numpy as np
import genesis as gs
import cv2
import imageio

########################## init ##########################
gs.init(backend=gs.cpu)

DT = 2e-4
Bodies = 9
Freq = 20
Phase = 180/Bodies
Iteration = int(2/Freq/DT)
PI = 3.141292
angle = 30.0/180.0*PI

########################## create a scene ##########################
# https://genesis-world.readthedocs.io/en/latest/api_reference/scene/scene.html←ここみると分かる
scene = gs.Scene(
    sim_options=gs.options.SimOptions(
        dt       = DT,
        substeps = 10,
    ),
    sph_options=gs.options.SPHOptions(
        lower_bound   = (-0.2, -0.2, 0.0),
        upper_bound   = (0.2, 0.2, 0.5),
        particle_size = 0.005,
    ),
    viewer_options=gs.options.ViewerOptions(
        res=(960, 720),
        camera_pos=(2.0, -1.0, 4.0),
        camera_lookat=(0.0, 0.0, 0.3),
        camera_fov=10,
        max_FPS=60,
    ),
    vis_options=gs.options.VisOptions(
        visualize_sph_boundary = True,
        show_world_frame=False,
        world_frame_size=1.0,
        show_link_frame=False,
        show_cameras=False,
        plane_reflection=True,
        ambient_light=(0.1,0.1,0.1),
    ),
    show_viewer = True,
    renderer=gs.renderers.Rasterizer(),
)

########################## entities ##########################
# https://genesis-world.readthedocs.io/en/latest/api_reference/entity/index.html
plane = scene.add_entity(
    morph=gs.morphs.Plane(),
)

fish = scene.add_entity( gs.morphs.MJCF(file='./fish.xml') )

liquid = scene.add_entity(
    # https://genesis-world.readthedocs.io/en/latest/api_reference/material/sph/liquid.html
    material=gs.materials.SPH.Liquid(
	rho=998.2, 
	stiffness=5000000, 	# 50000.0.
        mu=0.001,	# default is 0.005. 
	# exponent	# 7.0
	gamma=0.072,	# default is 0.01. but why?
        sampler='pbs',
    ),
    morph=gs.morphs.Box(
        pos  = (0.0, 0.0, 0.05),
        size = (0.4, 0.4, 0.1),
    ),
    surface=gs.surfaces.Default(
        color    = (0.4, 0.8, 1.0),
        vis_mode = 'recon',
    ),
)

cam = scene.add_camera(
        res=(960, 720),
        pos=(2.0, -1.0, 4.0),
        lookat=(0.0, 0.0, 0.3),
        fov=10,
        GUI=False,
)

########################## build ##########################
scene.build()

jnt_names = [
    'joint0',
    'joint1',
    'joint2',
    'joint3',
    'joint4',
    'joint5',
    'joint6',
    'joint7',
    'joint8',
]
angles = [
    0.1*angle,
    0.2*angle,
    0.3*angle,
    0.4*angle,
    0.5*angle,
    0.6*angle,
    0.7*angle,
    0.8*angle,
    0.9*angle,
]
dofs_idx = [fish.get_joint(name).dof_idx_local for name in jnt_names]

fish.set_dofs_position(np.array([0.0,
				angles[0]*np.sin(- Phase/180*PI*0) - angles[1]*np.sin(- Phase/180*PI*1),
				angles[1]*np.sin(- Phase/180*PI*1) - angles[2]*np.sin(- Phase/180*PI*2),
				angles[2]*np.sin(- Phase/180*PI*2) - angles[3]*np.sin(- Phase/180*PI*3),
				angles[3]*np.sin(- Phase/180*PI*3) - angles[4]*np.sin(- Phase/180*PI*4),
				angles[4]*np.sin(- Phase/180*PI*4) - angles[5]*np.sin(- Phase/180*PI*5),
				angles[5]*np.sin(- Phase/180*PI*5) - angles[6]*np.sin(- Phase/180*PI*6),
				angles[6]*np.sin(- Phase/180*PI*6) - angles[7]*np.sin(- Phase/180*PI*7),
				angles[7]*np.sin(- Phase/180*PI*7) - angles[8]*np.sin(- Phase/180*PI*8)]),
												dofs_idx)

############ Optional: set control gains ############
# set positional gains
fish.set_dofs_kp(
    kp             = np.array([4500, 4500, 4500, 4500, 4500, 4500, 4500, 4500, 4500]),
    dofs_idx_local = dofs_idx,
)
# set velocity gains
fish.set_dofs_kv(
    kv             = np.array([450, 450, 450, 450, 450, 450, 450, 450, 450]),
    dofs_idx_local = dofs_idx,
)
# set force range for safety
fish.set_dofs_force_range(
    lower          = np.array([-87, -87, -87, -87, -87, -87, -87, -87, -87]),
    upper          = np.array([ 87, 87, 87, 87, 87, 87, 87, 87, 87]),
    dofs_idx_local = dofs_idx,
)

# 7. 動画保存設定
rgb, depth, seg, normal = cam.render(depth=True, segmentation=True, normal=True)
width, height = rgb.shape[1], rgb.shape[0]
fps = 30
video = cv2.VideoWriter('fish.mp4',
                        cv2.VideoWriter_fourcc(*'mp4v'),
                        fps,
                        (width, height))

# 8. 録画ループ
cam.start_recording()

for i in range(Iteration):
    print(str(i)+"/"+str(Iteration))

    fish.set_dofs_velocity(np.array([
	2.0*PI*Freq*(       0.0*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*8) - angles[0]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*0) ),
	2.0*PI*Freq*( angles[0]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*0) - angles[1]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*1) ),
	2.0*PI*Freq*( angles[1]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*1) - angles[2]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*2) ),
	2.0*PI*Freq*( angles[2]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*2) - angles[3]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*3) ),
	2.0*PI*Freq*( angles[3]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*3) - angles[4]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*4) ),
	2.0*PI*Freq*( angles[4]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*4) - angles[5]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*5) ),
	2.0*PI*Freq*( angles[5]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*5) - angles[6]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*6) ),
	2.0*PI*Freq*( angles[6]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*6) - angles[7]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*7) ),
	2.0*PI*Freq*( angles[7]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*7) - angles[8]*np.cos(2.0*PI*Freq * i*DT - Phase/180*PI*8) )]), 
														dofs_idx_local=dofs_idx)

    scene.step()

    cam.set_pose(
            pos=(2.0, -1.0, 4.0), # pos=(3.0 * np.sin(i / 60), 3.0 * np.cos(i / 60), 2.5),
            lookat=(0.0, 0.0, 0.1),
    )
    cam.render()  # 録画モードで内部的にフレーム保存

    # フレームを取得して動画に書き込み
    rgb_frame, _, _, _ = cam.render(depth=False, segmentation=False, normal=False)
    frame_bgr = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR)
    video.write(frame_bgr)

video.release()

input()