import numpy as np

from rlkit.torch import pytorch_util as ptu

from rorlkit.torch.data import MotionGraphDataset
from rorlkit.torch.data_management.rosbag_data import batch_features_to_pose_arrays, get_identity_feature


class MotionGCNEvaluator(object):
    def __init__(
            self,
            model,
            test_dataset_info,
            save_path,
            batch_size=1,
    ):
        model.to(ptu.device)
        self.model = model
        self.database = MotionGraphDataset(**test_dataset_info)
        self.save_path = save_path
        self.batch_size = batch_size

    def make_rollout(
            self,
            unit_idx,
            n_rollout,
            turbulence=None
    ):
        self.model.eval()
        rollout = []
        # Add the starting point for better visualization
        base_feature = get_identity_feature()
        left_hand_feature, right_hand_feature = self.database.get_motion_features(0, unit_idx)
        if turbulence is not None:
            left_hand_feature += turbulence
            right_hand_feature += turbulence
        pose_arrays = batch_features_to_pose_arrays([base_feature, left_hand_feature, right_hand_feature])
        # print(pose_arrays)
        rollout.append(pose_arrays)

        graph_curr, _ = self.database.get_motion_graphs(0, unit_idx, shift=turbulence)
        for i in range(n_rollout):
            graph_next = self.model(graph_curr)
            graph_next_np = graph_next.squeeze().detach().cpu().numpy()
            # print(graph_curr, '\n', graph_next_np)
            pose_arrays = batch_features_to_pose_arrays(graph_next_np)
            # print(pose_arrays)
            rollout.append(pose_arrays)
            graph_curr = self.database.feature_array_to_graph(graph_next_np)
        np.save(self.save_path, rollout)
