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

gs.init(backend=gs.cpu)

Simulation_T = 4.0
DT = 1e-3

###### Condition #### 
scene = gs.Scene(
    show_viewer=True,
    sim_options = gs.options.SimOptions(
        dt     = DT,
	substeps = 4,
        gravity= (0.0, -9.81, 0.0),  # 重力ベクトル
	floor_height = 0.0,
#	use_hydroelastic_contact = False, 
    ),
    viewer_options=gs.options.ViewerOptions(
        res=(960, 720),
        camera_pos=(-1.0, 2.0, 4.0),
        camera_lookat=(0.0, 0.0, 0.0),
        camera_up=(0.0, 1.0, 0.0),
        camera_fov=10,
        max_FPS=60,
    ),
    renderer=gs.renderers.Rasterizer(),
)
"""
plane = scene.add_entity(
   gs.morphs.Plane(
        pos   = (0.0, 0.0, 0.0),
        euler = (-90.0,   # roll: 90°
                 0.0,    # pitch
                 0.0),   # yaw
   )
)
"""
floor = scene.add_entity(
    morph = gs.morphs.Box(size=(2, 0.1, 2), pos=(0, -0.05, 0), fixed=True, collision=True),  # 床を厚めに
    material =  gs.materials.Rigid(
        rho = 1000.0,
        friction = 0.8   # 床の摩擦
    ),
    surface=gs.surfaces.Default(
        color=(0.1, 0.4, 0.8, 1.0),
        vis_mode='visual',
    )
)

###### OBJECT ######
cobotta = scene.add_entity(
    material = gs.materials.Rigid(
#        rho = 1000.0,
        friction = 0.8    # 剛体側の摩擦係数
    ),
    morph = gs.morphs.MJCF(
        file = "./cobotta.xml",
        scale = 1.0,
        pos = (0.0, 0.0, 0.0),
        euler = (0.0, 0.0, 0.0),
        visualization = True,
        collision = True,
    ),
)

box = scene.add_entity(
    material = gs.materials.Rigid(
            rho = 1000.0,
            friction = 0.8    # 剛体側の摩擦係数
    ),
    morph = gs.morphs.Box(size=(0.028,0.05,0.06), pos=(-0.042, 0.035, 0.33)),
)

puchipuchi = scene.add_entity(
    material=gs.materials.PBD.Cloth(
        rho = 0.1, # [kg/m^2],
        static_friction = 0.6,
        kinetic_friction = 0.5,
	stretch_compliance = 0.01, stretch_relaxation = 0.3,
	bending_compliance = 0.1, bending_relaxation = 0.1,
	air_resistance = 0.001,
    ),
    morph=gs.morphs.Mesh(
        file='./cloth.obj',
        scale=1.0,
	pos=(0, 0.005, 0.4),
        euler=(90, 0, 0),	# (np.pi/2.0, 0, 0.0), x, y, z rotation [deg]
	decimate=False,
    ),
    surface=gs.surfaces.Default(
        color=(0.4, 0.4, 0.4, 1.0),
        vis_mode='visual',
    )
)

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

scene.build()


##### Initail Condition #######
jnt_names = [
    'joint1',
    'joint2',
    'joint3',
    'joint4',
    'joint5',
    'joint6',
    'joint7',
    'joint8',
]

dofs_idx = [cobotta.get_joint(name).dof_idx_local for name in jnt_names]

# https://genesis-world.readthedocs.io/en/latest/user_guide/getting_started/control_your_robot.html
# set positional gains
cobotta.set_dofs_kp(
#    kp             = np.array([4500, 4500, 4500, 3000, 2500, 1000, 50, 50]),
    kp             = np.array([2000, 2000, 2000, 2000, 2000, 500, 100, 100]),
    dofs_idx_local = dofs_idx,
)
# set velocity gains
cobotta.set_dofs_kv(
    kv             = np.array([450, 450, 450, 300, 250, 200, 10, 10]),
    dofs_idx_local = dofs_idx,
)
# set force range for safety
cobotta.set_dofs_force_range(
    lower          = np.array([-100, -100, -100, -100, -100, -100, -10, -10]),
    upper          = np.array([ 100,  100,  100,  100,  100,  100,  10,  10]),
    dofs_idx_local = dofs_idx,
)

joints = [ 0/180*np.pi, 100/180*np.pi, -8/180*np.pi, 0.0, 0.0, 0.0, 0.007, -0.007 ] #[ -45/180*np.pi, np.pi/2, 35/180*np.pi, 0.0, 0.0, 0.0, 0.007, -0.007 ]
cobotta.set_dofs_position( np.array([joints[0], joints[1], joints[2], joints[3], joints[4], joints[5], joints[6], joints[7] ]), 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('demo.mp4',
                        cv2.VideoWriter_fourcc(*'mp4v'),
                        fps,
                        (width, height))

# Recording
cam.start_recording()

######### Simulation ######
iteration = int(Simulation_T/DT)
for i in range(iteration):
    print(str(i)+"/"+str(iteration))

    if i < 1.0/DT:
       cobotta.control_dofs_position(np.array([0/180*np.pi, 100/180*np.pi, -8/180*np.pi+10/180*np.pi*i/(1.0/DT), 0.0, 0.0, 0.0, 0.007, -0.007]), dofs_idx)
    elif i < 1.5/DT:
       cobotta.control_dofs_position(np.array([0/180*np.pi, 100/180*np.pi, 2/180*np.pi, 0.0, 0.0, 0.0, 0.007-0.017*(i-1.0/DT)/(1.5/DT-1.0/DT), -0.007+0.017*(i-1.0/DT)/(1.5/DT-1.0/DT)]), dofs_idx)
    else:
       cobotta.control_dofs_position(np.array([0/180*np.pi*(i-1.5/DT)*DT, 100/180*np.pi, 2/180*np.pi+1.0/180*np.pi*(i-1.5/DT)*DT, 0.0, 0.0, 0.5*np.pi*(i-1.5/DT)*DT, -0.01, 0.01]), dofs_idx)

    scene.step()
    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)

#    input()

video.release()

input()