import torch

from lerobot.datasets import LeRobotDataset
from lerobot.policies import make_pre_post_processors
from lerobot.policies.your_policy import YourPolicy
from lerobot.processor import RobotProcessorPipeline, PolicyProcessorPipeline
dataset = LeRobotDataset("hf_user/dataset", episodes=[0])
sample = dataset[10]

model = YourPolicy.from_pretrained(
    "hf_user/model",
)
model.eval()
model.to("cuda")
preprocessor, postprocessor = make_pre_post_processors(model.config, pretrained_path="hf_user/model", dataset_stats=dataset.meta.stats)

preprocessed_sample = preprocessor(sample)
action = model.select_action(preprocessed_sample)
postprocessed_action = postprocessor(action)

from lerobot.processor import ProcessorStep, EnvTransition

class MyProcessorStep(ProcessorStep):
    """Example processor step - inherit and implement abstract methods."""

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        """Transform the transition - REQUIRED abstract method."""
        # Your processing logic here
        return transition

    def transform_features(self, features):
        """Declare how this step transforms feature shapes/types - REQUIRED abstract method."""
        return features  # Most processors return features unchanged

from lerobot.processor import RobotProcessorPipeline, PolicyProcessorPipeline

# For robot hardware (unbatched data)
robot_processor = RobotProcessorPipeline[RobotAction, RobotAction](
    steps=[step1, step2, step3],
    name="robot_pipeline"
)

# For model training/inference (batched data)
policy_processor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
    steps=[step1, step2, step3],
    name="policy_pipeline"
)

# Robot data structures: dict[str, Any] for observations and actions
robot_obs: dict[str, Any] = {
    "joint_1": 0.5,           # Individual joint values
    "joint_2": -0.3,
    "camera_0": image_array   # Raw camera data
}

robot_action: dict[str, Any] = {
    "joint_1": 0.2,          # Target joint positions
    "joint_2": 0.1,
    "gripper": 0.8
}

# Policy data structures: batch dicts and tensors
policy_batch: dict[str, Any] = {
    "observation.state": torch.tensor([[0.5, -0.3]]),      # Batched states
    "observation.images.camera0": torch.tensor(...),        # Batched images
    "action": torch.tensor([[0.2, 0.1, 0.8]])              # Batched actions
}

policy_action: torch.Tensor = torch.tensor([[0.2, 0.1, 0.8]])  # Model output tensor

# Training data preprocessing (optimized order for GPU performance)
training_preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
    steps=[
        RenameObservationsProcessorStep(rename_map={}),     # Standardize keys
        AddBatchDimensionProcessorStep(),                   # Add batch dims
        TokenizerProcessorStep(tokenizer_name="...", ...),  # Tokenize language
        DeviceProcessorStep(device="cuda"),                 # Move to GPU first
        NormalizerProcessorStep(features=..., stats=...),   # Normalize on GPU
    ]
)

# Model output postprocessing
training_postprocessor = PolicyProcessorPipeline[torch.Tensor, torch.Tensor](
    steps=[
        DeviceProcessorStep(device="cpu"),                  # Move to CPU
        UnnormalizerProcessorStep(features=..., stats=...), # Denormalize
    ],
    to_transition=policy_action_to_transition,
    to_output=transition_to_policy_action,
)

# Real deployment: Robot sensors → Model → Robot commands
with torch.no_grad():
    while not done:
        raw_obs = robot.get_observation()  # dict[str, Any]

        # Add your robot observation to policy observation processor

        policy_input = policy_preprocessor(raw_obs)  # Batched dict

        policy_output = policy.select_action(policy_input)  # Policy tensor

        policy_action = policy_postprocessor(policy_output)

        # Add your robot action to policy action processor

        robot.send_action(policy_action)

# Example: A processor that adds velocity to observations
class VelocityProcessor(ObservationProcessorStep):
    def observation(self, obs):
        new_obs = obs.copy()
        if "observation.state" in obs:
            # concatenate computed velocity field to the state
            new_obs["observation.state"] = self._compute_velocity(obs["observation.state"])
        return new_obs

    def transform_features(self, features):
        """Declare the new velocity field we're adding."""
        state_feature = features[PipelineFeatureType.OBSERVATION].get("observation.state")
        if state_feature:
            double_shape = (state_feature.shape[0] * 2,) if state_feature.shape else (2,)
            features[PipelineFeatureType.OBSERVATION]["observation.state"] = PolicyFeature(
                type=FeatureType.STATE, shape=double_shape
            )
        return features

from lerobot.datasets import aggregate_pipeline_dataset_features

# Start with robot's raw features
initial_features = create_initial_features(
    observation=robot.observation_features,  # {"joint_1.pos": float, "camera_0": (480,640,3)}
    action=robot.action_features            # {"joint_1.pos": float, "gripper.pos": float}
)

# Apply processor pipeline to compute final features
final_features = aggregate_pipeline_dataset_features(
    pipeline=my_processor_pipeline,
    initial_features=initial_features,
    use_videos=True
)

# Use for dataset creation
dataset = LeRobotDataset.create(
    repo_id="my_dataset",
    features=final_features,  # Knows exactly what data to expect
    ...
)