import mujoco
import numpy as np

# Load model from XML
model = mujoco.MjModel.from_xml_path("robot.xml")
data = mujoco.MjData(model)

# Simulation loop
for _ in range(1000):
    # Set control
    data.ctrl[:] = np.random.randn(model.nu)
    
    # Step simulation
    mujoco.mj_step(model, data)
    
    # Get state
    position = data.qpos.copy()
    velocity = data.qvel.copy()
    
    print(f"Position: {position}, Velocity: {velocity}")

def get_contact_info(model, data):
    """Extract contact information"""
    contacts = []
    
    for i in range(data.ncon):
        contact = data.contact[i]
        
        # Get contacting bodies
        geom1 = model.geom(contact.geom1).name
        geom2 = model.geom(contact.geom2).name
        
        # Contact force
        force = np.zeros(6)
        mujoco.mj_contactForce(model, data, i, force)
        
        contacts.append({
            "geom1": geom1,
            "geom2": geom2,
            "position": contact.pos.copy(),
            "force": force[:3],
            "torque": force[3:]
        })
    
    return contacts

def inverse_kinematics(model, data, target_pos, site_name, max_iter=100):
    """Solve IK for a site to reach target position"""
    site_id = model.site(site_name).id
    
    for _ in range(max_iter):
        # Current site position
        current_pos = data.site_xpos[site_id]
        
        # Error
        error = target_pos - current_pos
        if np.linalg.norm(error) < 1e-4:
            break
        
        # Compute Jacobian
        jacp = np.zeros((3, model.nv))
        mujoco.mj_jacSite(model, data, jacp, None, site_id)
        
        # Pseudoinverse solution
        dq = np.linalg.pinv(jacp) @ error
        
        # Update joint positions
        data.qpos[:] += dq * 0.1
        mujoco.mj_forward(model, data)
    
    return data.qpos.copy()

import mujoco.viewer

# Interactive viewer
with mujoco.viewer.launch_passive(model, data) as viewer:
    while viewer.is_running():
        mujoco.mj_step(model, data)
        viewer.sync()

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

def render_frame(data):
    renderer.update_scene(data)
    return renderer.render()

import gymnasium as gym

# Create environment
env = gym.make("Humanoid-v4", render_mode="rgb_array")

# Collect data for world model
def collect_trajectory(env, policy, num_steps=1000):
    trajectory = {
        "observations": [],
        "actions": [],
        "rewards": [],
        "frames": []
    }
    
    obs, info = env.reset()
    
    for _ in range(num_steps):
        action = policy(obs)
        next_obs, reward, terminated, truncated, info = env.step(action)
        frame = env.render()
        
        trajectory["observations"].append(obs)
        trajectory["actions"].append(action)
        trajectory["rewards"].append(reward)
        trajectory["frames"].append(frame)
        
        obs = next_obs
        
        if terminated or truncated:
            obs, info = env.reset()
    
    return trajectory

class MuJoCoWorldModelEnv:
    """Environment wrapper for world model training"""
    
    def __init__(self, model_path):
        self.model = mujoco.MjModel.from_xml_path(model_path)
        self.data = mujoco.MjData(self.model)
        self.renderer = mujoco.Renderer(self.model)
    
    def get_observation(self):
        return {
            "qpos": self.data.qpos.copy(),
            "qvel": self.data.qvel.copy(),
            "image": self.render()
        }
    
    def step(self, action):
        self.data.ctrl[:] = action
        mujoco.mj_step(self.model, self.data)
        return self.get_observation()
    
    def render(self):
        self.renderer.update_scene(self.data)
        return self.renderer.render()
    
    def get_state(self):
        """Get full simulation state for world model"""
        return {
            "qpos": self.data.qpos.copy(),
            "qvel": self.data.qvel.copy(),
            "ctrl": self.data.ctrl.copy(),
            "time": self.data.time
        }
    
    def set_state(self, state):
        """Set simulation state"""
        self.data.qpos[:] = state["qpos"]
        self.data.qvel[:] = state["qvel"]
        self.data.ctrl[:] = state["ctrl"]
        self.data.time = state["time"]
        mujoco.mj_forward(self.model, self.data)