import os
import tempfile
import numpy as np
import imageio
import gradio as gr
import mujoco

from infer_engine import SpacePolicyInference

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
ROBOT_DIR = os.path.join(CURRENT_DIR, "robot")
POLICIES_DIR = os.path.join(CURRENT_DIR, "policies")

def simulate_motion(motion_type, duration, camera_angle, progress=gr.Progress()):
    progress(0, desc="시뮬레이션 환경 초기화 중...")
    
    # Select XML model
    if motion_type == "⚽ 왼발 공차기 (Ball Kick)":
        xml_path = os.path.join(ROBOT_DIR, "scene_ball.xml")
    else:
        xml_path = os.path.join(ROBOT_DIR, "scene.xml")

    model = mujoco.MjModel.from_xml_path(xml_path)
    model.opt.timestep = 0.005
    data = mujoco.MjData(model)

    # Initialize policy
    engine = SpacePolicyInference(model, data, POLICIES_DIR)

    # Reset robot position to standing default
    freejoint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
    qpos_adr = model.jnt_qposadr[freejoint_id]
    data.qpos[qpos_adr + 0] = 0.0
    data.qpos[qpos_adr + 1] = 0.0
    data.qpos[qpos_adr + 2] = 0.125
    data.qpos[qpos_adr + 3:qpos_adr + 7] = [1, 0, 0, 0]
    for i, idx in enumerate(engine.joint_qpos_indices):
        data.qpos[idx] = engine.default_pose[i]
    data.ctrl[:] = engine.default_pose
    mujoco.mj_forward(model, data)

    # Configure behavior
    if motion_type == "🚶 전진 걷기 (Forward Walk)":
        engine.set_vel_cmd(0.25, 0.0, 0.0)
    elif motion_type == "🚶 후진 걷기 (Backward Walk)":
        engine.set_vel_cmd(-0.2, 0.0, 0.0)
    elif motion_type == "🔄 좌회전 (Turn Left)":
        engine.set_vel_cmd(0.1, 0.0, 0.8)
    elif motion_type == "🔄 우회전 (Turn Right)":
        engine.set_vel_cmd(0.1, 0.0, -0.8)
    elif motion_type == "🧍 제자리 서기 (Standing Balance)":
        engine.set_vel_cmd(0.0, 0.0, 0.0)
    elif motion_type == "🤸 앞구르기 (Roulade)":
        engine.trigger_behavior("roulade")
    elif motion_type == "⚽ 왼발 공차기 (Ball Kick)":
        engine.trigger_behavior("kick_left")
    elif motion_type == "🧎 바닥 줍기 (Ground Pick)":
        engine.trigger_ground_pick()
    elif motion_type == "🪑 앉기 (Sit Down)":
        engine.trigger_sit()

    # Offscreen renderer
    width, height = 640, 480
    renderer = mujoco.Renderer(model, height, width)

    # Setup camera
    camera = mujoco.MjvCamera()
    mujoco.mjv_defaultCamera(camera)

    frames = []
    fps = 30
    control_freq = 50 # Hz
    decimation = 4
    total_steps = int(duration * control_freq)
    dt = 1.0 / control_freq

    # Render frequency (every ~1.66 control steps to reach 30 fps)
    sim_time = 0.0
    next_render_time = 0.0

    for step_i in range(total_steps):
        engine.update(dt)
        engine.step()

        for _ in range(decimation):
            mujoco.mj_step(model, data)

        sim_time += dt

        if sim_time >= next_render_time:
            # Update camera to track robot
            robot_pos = data.qpos[qpos_adr:qpos_adr + 3]
            camera.lookat = [robot_pos[0], robot_pos[1], robot_pos[2] + 0.05]

            if camera_angle == "측면 추적 (Side View)":
                camera.distance = 1.1
                camera.elevation = -15
                camera.azimuth = 90
            elif camera_angle == "정면 뷰 (Front View)":
                camera.distance = 1.0
                camera.elevation = -10
                camera.azimuth = 180
            elif camera_angle == "대각선 뷰 (Isometric)":
                camera.distance = 1.2
                camera.elevation = -22
                camera.azimuth = 135
            elif camera_angle == "후방 3인칭 (Behind View)":
                camera.distance = 1.1
                camera.elevation = -12
                camera.azimuth = 0

            renderer.update_scene(data, camera=camera)
            img = renderer.render()
            frames.append(img)
            next_render_time += 1.0 / fps

        if step_i % 15 == 0:
            progress(step_i / total_steps, desc=f"시뮬레이션 진행 중... ({step_i}/{total_steps})")

    progress(0.95, desc="비디오 파일 생성 중...")
    temp_dir = tempfile.mkdtemp()
    video_path = os.path.join(temp_dir, "output.mp4")
    imageio.mimsave(video_path, frames, fps=fps)

    final_pos = data.qpos[qpos_adr:qpos_adr + 3]
    distance = np.linalg.norm(final_pos[:2])

    status_md = f"""
### 📊 시뮬레이션 결과 요약
- **선택된 동작**: `{motion_type}`
- **시뮬레이션 시간**: `{duration:.1f}초` (총 {len(frames)} 프레임 / 30 FPS)
- **최종 이동 거리**: `{distance:.3f} m` (X: `{final_pos[0]:.2f}m`, Y: `{final_pos[1]:.2f}m`, Z: `{final_pos[2]:.2f}m`)
- **로봇 상태**: `정상 안정 상태 유지 완료`
"""
    return video_path, status_md

# Gradio Interface
custom_css = """
#container { max-width: 960px; margin: 0 auto; }
.gr-button-primary { background: linear-gradient(90deg, #F59E0B 0%, #D97706 100%) !important; border: none !important; }
"""

with gr.Blocks(title="🦆 MicroDuck 3D Robot Simulator", css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🦆 MicroDuck 3D 로봇 강화학습 시뮬레이터
        MuJoCo 물리 엔진과 온닉스(ONNX) 심층강화학습(Deep RL) 정책 모델로 구동되는 **MicroDuck** 로봇의 동작을 웹에서 직접 실행하고 관찰해보세요.
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            motion_input = gr.Radio(
                label="🎯 로봇 동작 (Motion Policy)",
                choices=[
                    "🚶 전진 걷기 (Forward Walk)",
                    "🚶 후진 걷기 (Backward Walk)",
                    "🔄 좌회전 (Turn Left)",
                    "🔄 우회전 (Turn Right)",
                    "🧍 제자리 서기 (Standing Balance)",
                    "🤸 앞구르기 (Roulade)",
                    "⚽ 왼발 공차기 (Ball Kick)",
                    "🧎 바닥 줍기 (Ground Pick)",
                    "🪑 앉기 (Sit Down)",
                ],
                value="🚶 전진 걷기 (Forward Walk)"
            )
            
            duration_input = gr.Slider(
                label="⏱️ 시뮬레이션 시간 (초)",
                minimum=1.5,
                maximum=4.5,
                value=2.5,
                step=0.5
            )

            camera_input = gr.Dropdown(
                label="🎥 카메라 앵글",
                choices=["측면 추적 (Side View)", "정면 뷰 (Front View)", "대각선 뷰 (Isometric)", "후방 3인칭 (Behind View)"],
                value="대각선 뷰 (Isometric)"
            )

            run_btn = gr.Button("🎬 3D 시뮬레이션 렌더링 시작", variant="primary", size="lg")

        with gr.Column(scale=6):
            video_output = gr.Video(label="📹 시뮬레이션 렌더링 영상", autoplay=True, loop=True)
            status_output = gr.Markdown("시뮬레이션을 실행하면 결과 통계가 표시됩니다.")

    run_btn.click(
        fn=simulate_motion,
        inputs=[motion_input, duration_input, camera_input],
        outputs=[video_output, status_output]
    )

    gr.Markdown(
        """
        ---
        💡 **MicroDuck 기술 스택**: MuJoCo Physics, ONNX Runtime, Isaac Gym / Genesis RL Trained Policies.
        """
    )

if __name__ == "__main__":
    demo.launch()