# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from typing import Dict, List, Optional, Tuple

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import torch

import theseus as th

# ----------------------------------------------------------------------------------- #
# ------------------------------------ Data loading --------------------------------- #
# ----------------------------------------------------------------------------------- #


class TactilePushingDataset:
    def __init__(
        self,
        data_fname: str,
        sdf_fname: str,
        episode_length: int,
        batch_size: int,
        max_episodes: int,
        max_steps: int,
        device: th.DeviceType,
        split_episodes: bool = False,
        data_mode: str = "all",
        val_ratio: float = 0.1,
        seed: int = 1234567,
    ):
        assert data_mode in ["all", "train", "val"]

        batch_size = min(batch_size, max_episodes)
        data = TactilePushingDataset._load_dataset_from_file(
            data_fname, episode_length, max_episodes, device, split_episodes
        )
        (
            self.sdf_data_tensor,
            self.sdf_cell_size,
            self.sdf_origin,
        ) = TactilePushingDataset._load_tactile_sdf_from_file(sdf_fname, device)

        # keys: img_feats, eff_poses, obj_poses, contact_episode, contact_flag
        num_episodes = data["obj_poses"].shape[0]
        if data_mode == "all":
            idx = np.arange(num_episodes)
        else:
            rng = np.random.default_rng(seed)
            order = rng.permutation(num_episodes)
            stop = max(int(np.ceil(num_episodes * val_ratio)), 2)
            idx = order[:stop] if data_mode == "val" else order[stop:]

        self.img_feats = data["img_feats"][idx]  # type: ignore
        self.eff_poses = data["eff_poses"][idx]  # type: ignore
        self.obj_poses = data["obj_poses"][idx]  # type: ignore
        self.contact_episode = data["contact_episode"][idx]  # type: ignore
        self.contact_flag = data["contact_flag"][idx]  # type: ignore
        # Check sizes of the attributes assigned above
        self.dataset_size: int = -1
        for key in data:
            if self.dataset_size == -1:
                self.dataset_size = getattr(self, key).shape[0]
            else:
                assert self.dataset_size == getattr(self, key).shape[0]
        print(f"Dataset for mode '{data_mode}' has size {self.dataset_size}.")

        # obj_poses is shape (num_episodes, episode_length, 3)
        self.time_steps = np.minimum(max_steps, self.obj_poses.shape[1])
        self.batch_size = batch_size
        self.num_batches = (self.dataset_size - 1) // self.batch_size + 1

    @staticmethod
    def _load_dataset_from_file(
        filename: str,
        episode_length: int,
        max_episodes: int,
        device: th.DeviceType,
        split_episodes: bool = False,
    ) -> Dict[str, torch.Tensor]:
        # Load all episode data
        with open(filename) as f:
            import json

            data_from_file = json.load(f)

        dataset_all = {}
        dataset_all["obj_poses"] = torch.tensor(
            data_from_file["obj_poses_2d"], device=device
        )
        dataset_all["eff_poses"] = torch.tensor(
            data_from_file["ee_poses_2d"], device=device
        )
        dataset_all["img_feats"] = torch.tensor(
            data_from_file["img_feats"], device=device
        )
        dataset_all["contact_episode"] = torch.tensor(
            data_from_file["contact_episode"], device=device
        )
        dataset_all["contact_flag"] = torch.tensor(
            data_from_file["contact_flag"], device=device
        )

        # Read all episodes and filter those with length less than desired
        episode_indices = [
            idx.item() for idx in dataset_all["contact_episode"].unique()
        ]
        data: Dict[str, List[torch.Tensor]] = dict(
            [(k, []) for k in dataset_all.keys()]
        )

        for episode in episode_indices:
            if len(data["obj_poses"]) >= max_episodes:
                break
            ds_idxs = torch.nonzero(dataset_all["contact_episode"] == episode).squeeze()
            if len(ds_idxs) < episode_length:
                continue
            for key, val in dataset_all.items():
                if split_episodes:
                    tensors = TactilePushingDataset._get_tensor_splits(
                        val[ds_idxs], episode_length
                    )
                else:
                    ds_idxs = ds_idxs[:episode_length]
                    tensors = [val[ds_idxs]]
                for tensor in tensors:
                    data[key].append(tensor)

        # Stack all episode data into single tensors
        data_tensors = {}
        for key in data:
            data_tensors[key] = torch.stack(data[key])
        num_ep_read, len_ep_read = data_tensors[key].shape[:2]
        print(f"Read {num_ep_read} episodes of length {len_ep_read}.")
        return data_tensors

    @staticmethod
    def _get_tensor_splits(
        tensor: torch.Tensor, episode_length: int
    ) -> List[torch.Tensor]:
        squeeze = False
        if tensor.ndim == 1:
            squeeze = True
            tensor = tensor.view(-1, 1)
        length, dof = tensor.shape
        num_splits = length // episode_length
        mod_tensor = tensor[: num_splits * episode_length]
        reshaped_tensor = mod_tensor.view(num_splits, -1, dof)
        if squeeze:
            reshaped_tensor = reshaped_tensor.squeeze(2)
        return [t for t in reshaped_tensor]

    @staticmethod
    def _load_tactile_sdf_from_file(
        filename: str, device: th.DeviceType
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        with open(filename) as f:
            import json

            sdf_from_file = json.load(f)
        sdf_data_vec = sdf_from_file["grid_data"]

        sdf_data_mat = np.zeros(
            (sdf_from_file["grid_size_y"], sdf_from_file["grid_size_x"])
        )
        for i in range(sdf_data_mat.shape[0]):
            for j in range(sdf_data_mat.shape[1]):
                sdf_data_mat[i, j] = sdf_data_vec[i][j]

        sdf_data_tensor = torch.tensor(sdf_data_mat, device=device).unsqueeze(0)
        cell_size = torch.tensor([sdf_from_file["grid_res"]], device=device).unsqueeze(
            0
        )
        origin = torch.tensor(
            [sdf_from_file["grid_origin_x"], sdf_from_file["grid_origin_y"]],
            device=device,
        ).unsqueeze(0)

        return sdf_data_tensor, cell_size, origin

    def get_batch(self, batch_idx: int) -> Dict[str, torch.Tensor]:
        assert batch_idx < self.num_batches
        start = batch_idx * self.batch_size
        end = min(start + self.batch_size, self.dataset_size)
        batch = {}
        batch["img_feats"] = self.img_feats[start:end, : self.time_steps]
        batch["eff_poses"] = self.eff_poses[start:end, : self.time_steps]
        batch["obj_poses"] = self.obj_poses[start:end, : self.time_steps]
        batch["obj_poses_gt"] = self.obj_poses[start:end, : self.time_steps, :].clone()
        batch["eff_poses_gt"] = self.eff_poses[start:end, : self.time_steps, :].clone()
        batch["obj_start_pose"] = self.obj_poses[start:end, 0]
        for i in range(self.time_steps):
            batch[f"motion_capture_{i}"] = self.eff_poses[start:end, i]
        return batch


# ----------------------------------------------------------------------------------- #
# ---------------------------------- Visualization ---------------------------------- #
# ----------------------------------------------------------------------------------- #
def _draw_tactile_effector(poses: np.ndarray, label: str = "groundtruth"):
    linestyle = "-"
    color = "tab:orange"
    if label == "groundtruth":
        linestyle = "--"
        color = "tab:gray"

    # plot contact point and normal
    plt.plot(poses[-1][0], poses[-1][1], "k*", linestyle=linestyle)
    ori = poses[-1][2]
    (dx, dy) = (0.03 * -np.sin(ori), 0.03 * np.cos(ori))
    plt.arrow(
        poses[-1][0],
        poses[-1][1],
        dx,
        dy,
        linewidth=2,
        head_width=0.001,
        color=color,
        head_length=0.01,
        fc=color,
        ec=color,
    )

    eff_radius = 0.0075
    circle = mpatches.Circle(
        (poses[-1][0], poses[-1][1]), color=color, radius=eff_radius
    )
    plt.gca().add_patch(circle)


def _draw_tactile_object(
    poses: np.ndarray, rect_len_x: float, rect_len_y: float, label: str = "optimizer"
):
    linestyle = "-"
    color = "tab:orange"
    if label == "groundtruth":
        linestyle = "--"
        color = "tab:gray"

    plt.plot(
        poses[:, 0],
        poses[:, 1],
        color=color,
        linestyle=linestyle,
        label=label,
        linewidth=2,
        alpha=0.9,
    )

    # shape: rect
    yaw = poses[-1][2]
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    offset = np.matmul(R, np.array([[0.5 * rect_len_x], [0.5 * rect_len_y]]))
    xb, yb = poses[-1][0] - offset[0], poses[-1][1] - offset[1]
    rect = mpatches.Rectangle(
        (xb, yb),
        rect_len_x,
        rect_len_y,
        angle=(np.rad2deg(yaw)),
        facecolor="None",
        edgecolor=color,
        linestyle=linestyle,
        linewidth=2,
    )
    plt.gca().add_patch(rect)


def visualize_tactile_push2d(
    obj_poses: torch.Tensor,
    eff_poses: torch.Tensor,
    obj_poses_gt: torch.Tensor,
    eff_poses_gt: torch.Tensor,
    rect_len_x: float,
    rect_len_y: float,
    save_fname: Optional[str] = None,
):
    plt.cla()
    plt.xlim((-2, 2))
    plt.ylim((-2, 2))
    plt.gca().axis("equal")
    plt.axis("off")

    _draw_tactile_object(
        obj_poses.cpu().detach().numpy(), rect_len_x, rect_len_y, label="optimizer"
    )
    _draw_tactile_effector(eff_poses.cpu().detach().numpy(), label="optimizer")
    _draw_tactile_object(
        obj_poses_gt.cpu().detach().numpy(), rect_len_x, rect_len_y, label="groundtruth"
    )
    _draw_tactile_effector(eff_poses_gt.cpu().detach().numpy(), label="groundtruth")

    plt.show()
    plt.pause(1e-9)

    if save_fname is not None:
        plt.savefig(save_fname)
