import torch
import numpy as np
from diffusion_policy.model.diffusion_policy import DiffusionPolicy

# Load pretrained model
def load_model(device="cuda" if torch.cuda.is_available() else "cpu"):
    model_path = "diffusion-policy/pretrained/6dof_robot"  # Replace with actual model path
    model = DiffusionPolicy.load_from_checkpoint(model_path, map_location=device)
    model.to(device)
    model.eval()
    return model

# Generate trajectory for a 6-DOF robot
def generate_trajectory(model, start_state, goal_state, num_steps=50, device="cuda"):
    trajectory = torch.zeros((num_steps, 6), device=device)  # Store trajectory
    state = torch.tensor(start_state, dtype=torch.float32, device=device)

    with torch.no_grad():
        for t in range(num_steps):
            action = model.predict(state)  # Predict next action
            state = state + action  # Simulate movement
            trajectory[t] = state

    return trajectory.cpu().numpy()  # Convert to NumPy for plotting