import copy
import dataclasses as dc

import numpy as np
from pydrake.common.eigen_geometry import Quaternion
from pydrake.math import RigidTransform, RotationMatrix

from robot_gym.multiarm_spaces import (
    MultiarmObservation,
    PosesAndGrippers,
    RestorePosesAndGrippersConfig,
)


def _quaternion_shortest_path(q_A, q_B):
    # Returns q_B s.t. q_A and q_B have minimum dot product.
    assert isinstance(q_A, Quaternion)
    assert isinstance(q_B, Quaternion)
    if np.dot(q_A.wxyz(), q_B.wxyz()) < 0:
        return Quaternion(-q_B.wxyz())
    else:
        return q_B


def se3_interp(s, X_AB_start, X_AB_end):
    p_TP_start = X_AB_start.translation()
    p_TP_end = X_AB_end.translation()
    p_TP = p_TP_start + s * (p_TP_end - p_TP_start)
    q_TP_start = X_AB_start.rotation().ToQuaternion()
    q_TP_end = X_AB_end.rotation().ToQuaternion()
    q_TP_end = _quaternion_shortest_path(q_TP_start, q_TP_end)
    q_TP = q_TP_start.slerp(s, q_TP_end)
    X_AB = RigidTransform(q_TP, p_TP)
    return X_AB


def _interp_vec(s, x1, x2):
    x = x1 + s * (x2 - x1)
    return x


def _assert_nominal_action(action):
    assert action.joint_position is None
    assert action.joint_velocity is None
    assert action.joint_torque is None
    assert action.joint_torque_external is None
    assert action.wrench is None
    assert action.external_wrench is None


def interp_action(s, action_1, action_2):
    # TODO(eric.cousineau): Enable these checks.
    # assert _assert_nominal_action(action_1)
    # assert _assert_nominal_action(action_2)
    action = copy.deepcopy(action_1)
    for key in action.poses.keys():
        action.poses[key] = se3_interp(
            s, action_1.poses[key], action_2.poses[key]
        )
    for key in action.grippers.keys():
        action.grippers[key] = _interp_vec(
            s, action_1.grippers[key], action_2.grippers[key]
        )
    return action


def action_get_stop_info(action):
    debug = action.debugging_output
    if debug is None:
        return None
    stop_info = {}
    is_success = debug.get("is_success", None)
    if is_success is not None:
        stop_info["is_success"] = is_success
    is_retry = debug.get("is_retry", None)
    if is_retry is not None:
        stop_info["is_retry"] = is_retry
    stop_reason = debug.get("stop_reason", None)
    if stop_reason is not None:
        stop_info["stop_reason"] = stop_reason
    return stop_info


def _normalize(x, *, tol=1e-10):
    x = np.asarray(x)
    n = np.linalg.norm(x)
    assert n >= tol
    return x / n


def rotation_6d_to_matrix(d6: np.ndarray) -> np.ndarray:
    """
    NumPy version of diffusion_policy rotation_6d_to_matrix.
    """
    # Ensure we use high precision.
    d6 = d6.astype(np.float64)
    a1, a2 = d6[:3], d6[3:]
    b1 = _normalize(a1)
    b2 = a2 - np.dot(b1, a2) * b1
    b2 = _normalize(b2)
    b3 = np.cross(b1, b2)
    return np.stack((b1, b2, b3))


def matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray:
    """
    NumPy version of diffusion_policy matrix_to_rotation_6d.
    """
    return matrix[:2, :].copy().reshape((6,))


def _rigid_transform_to_xyz_rot_6d(X):
    """
    Converts rigid transform to dictionary of {"xyz": ..., "rot_6d": ...}

    See matrix_to_rotation_6d for more info about "rot_6d".
    """
    assert isinstance(X, RigidTransform)
    xyz = X.translation()
    rot_6d = matrix_to_rotation_6d(X.rotation().matrix())
    return {
        "xyz": xyz,
        "rot_6d": rot_6d,
    }


def _to_np_value(v):
    if isinstance(v, RigidTransform):
        return _rigid_transform_to_xyz_rot_6d(v)
    else:
        v = np.asarray(v)
        if v.shape == ():
            v = v.reshape(1)
        return v


def _flatten_dict(config, prefix="", delim="."):
    """
    Transforms a dict of the form:

        {"top": {"mid": {"bottom": 0.25}}}

    to the following form:

        {"top.mid.bottom": 0.25}
    """
    assert isinstance(config, dict)
    out = dict()
    for k, v in config.items():
        assert isinstance(k, str), repr(k)
        assert "." not in k, repr(k)
        if isinstance(v, dict):
            v = _flatten_dict(v, f"{prefix}{k}{delim}", delim=delim)
            for ki, vi in v.items():
                assert ki not in config
            out.update(v)
        else:
            out[f"{prefix}{k}"] = v
    return out


def _dict_to_flat_pure_np_dict(x):
    """
    Takes a nested dictionary `x` and ensures that all values contained within
    it are purely NumPy values with the correct coordinate representation.
    Additionally flattens the dictionary, e.g.
        {"pose": {"xyz": ..., "rot_6d": ...}}
    becomes
        {"pose__xyz": ..., "pose__rot_6d": ...}
    """
    keys_prev = list(x.keys())
    while True:
        x = _flatten_dict(x, delim="__")
        # Remove any values that are `None`.
        x = {k: v for k, v in x.items() if v is not None}
        # Now sort through.
        for k, v in x.items():
            x[k] = _to_np_value(v)
        keys = list(x.keys())
        if keys == keys_prev:
            break
        else:
            keys_prev = keys
    return x


def _flatten_np_dict_to_vector(x):
    x = np.concatenate([np.asarray(x).reshape(-1) for x in x.values()])
    return x


def multiarm_action_to_dp_vector(action: PosesAndGrippers) -> np.ndarray:
    raw_dict = {}
    if action.poses is not None:
        raw_dict["poses"] = action.poses
    if action.grippers is not None:
        raw_dict["grippers"] = action.grippers
    assert len(raw_dict) > 0
    flat_dict = _dict_to_flat_pure_np_dict(raw_dict)
    vector = _flatten_np_dict_to_vector(flat_dict)
    return vector


def dp_vector_to_multiarm_action(
    config: RestorePosesAndGrippersConfig, vector: np.ndarray
) -> PosesAndGrippers:
    index = 0
    # Poses.
    poses = {}
    for model_name in config.model_names:
        xyz = vector[index : index + 3]
        index += 3
        rot_6d = vector[index : index + 6]
        index += 6
        rot = rotation_6d_to_matrix(rot_6d)
        pose = RigidTransform(R=RotationMatrix(rot), p=xyz)
        poses[model_name] = pose
    # Grippers.
    grippers = {}
    for gripper_name in config.gripper_names:
        gripper_width = vector[index].item()
        index += 1
        grippers[gripper_name] = gripper_width
    # Done.
    assert index == len(vector)
    return PosesAndGrippers(poses=poses, grippers=grippers)


def multiarm_observation_to_dp_dict(
    obs: MultiarmObservation,
    *,
    rgb_only=False,
):
    """
    Converts MultiarmObservation to a dictionary for inference.
    """
    raw_dict = {}
    # Camera images.
    for camera_name, image_set in obs.visuo.items():
        if image_set.rgb is not None:
            raw_dict[camera_name] = image_set.rgb.array
        if image_set.depth is not None:
            raw_dict[f"{camera_name}_depth"] = image_set.depth.array
        if image_set.label is not None:
            raw_dict[f"{camera_name}_label"] = image_set.label.array
    # Proprioception.
    robot_dict = {"robot": dc.asdict(obs.robot)}
    robot_dict_flat = _dict_to_flat_pure_np_dict(robot_dict)
    # Root-level.
    if obs.timestamp_packaged is not None:
        robot_dict_flat["timestamp_packaged"] = _to_np_value(
            obs.timestamp_packaged
        )
    # Language.
    if obs.language_instruction is not None:
        raw_dict["language_instruction"] = obs.language_instruction

    key_overlap = set(raw_dict.keys()) & set(robot_dict_flat.keys())
    assert len(key_overlap) == 0, key_overlap
    raw_dict.update(robot_dict_flat)
    assert len(raw_dict) > 0
    return raw_dict