import torch
from maskrcnn_benchmark.config import cfg

# transpose
FLIP_LEFT_RIGHT = 0
FLIP_TOP_BOTTOM = 1


class Keypoints(object):
    def __init__(self, keypoints, size, mode=None):
        # FIXME remove check once we have better integration with device
        # in my version this would consistently return a CPU tensor
        device = keypoints.device if isinstance(keypoints, torch.Tensor) else torch.device("cpu")
        keypoints = torch.as_tensor(keypoints, dtype=torch.float32, device=device)
        num_keypoints = keypoints.shape[0]
        if num_keypoints:
            keypoints = keypoints.view(num_keypoints, -1, 3)

        # TODO should I split them?
        # self.visibility = keypoints[..., 2]
        self.keypoints = keypoints  # [..., :2]

        self.size = size
        self.mode = mode
        self.extra_fields = {}

    def crop(self, box):
        raise NotImplementedError()

    def resize(self, size, *args, **kwargs):
        ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(size, self.size))
        ratio_w, ratio_h = ratios
        resized_data = self.keypoints.clone()
        resized_data[..., 0] *= ratio_w
        resized_data[..., 1] *= ratio_h
        keypoints = type(self)(resized_data, size, self.mode)
        for k, v in self.extra_fields.items():
            keypoints.add_field(k, v)
        return keypoints

    def transpose(self, method):
        if method not in (FLIP_LEFT_RIGHT,):
            raise NotImplementedError("Only FLIP_LEFT_RIGHT implemented")

        flip_inds = self.FLIP_INDS
        flipped_data = self.keypoints[:, flip_inds]
        width = self.size[0]
        TO_REMOVE = 1
        # Flip x coordinates
        flipped_data[..., 0] = width - flipped_data[..., 0] - TO_REMOVE

        # Maintain COCO convention that if visibility == 0, then x, y = 0
        inds = flipped_data[..., 2] == 0
        flipped_data[inds] = 0

        keypoints = type(self)(flipped_data, self.size, self.mode)
        for k, v in self.extra_fields.items():
            keypoints.add_field(k, v)
        return keypoints

    def to(self, *args, **kwargs):
        keypoints = type(self)(self.keypoints.to(*args, **kwargs), self.size, self.mode)
        for k, v in self.extra_fields.items():
            if hasattr(v, "to"):
                v = v.to(*args, **kwargs)
            keypoints.add_field(k, v)
        return keypoints

    def __getitem__(self, item):
        keypoints = type(self)(self.keypoints[item], self.size, self.mode)
        for k, v in self.extra_fields.items():
            keypoints.add_field(k, v[item])
        return keypoints

    def add_field(self, field, field_data):
        self.extra_fields[field] = field_data

    def get_field(self, field):
        return self.extra_fields[field]

    def __repr__(self):
        s = self.__class__.__name__ + "("
        s += "num_instances={}, ".format(len(self.keypoints))
        s += "image_width={}, ".format(self.size[0])
        s += "image_height={})".format(self.size[1])
        return s


class PersonKeypoints(Keypoints):
    _NAMES = [
        "nose",
        "left_eye",
        "right_eye",
        "left_ear",
        "right_ear",
        "left_shoulder",
        "right_shoulder",
        "left_elbow",
        "right_elbow",
        "left_wrist",
        "right_wrist",
        "left_hip",
        "right_hip",
        "left_knee",
        "right_knee",
        "left_ankle",
        "right_ankle",
    ]
    _FLIP_MAP = {
        "left_eye": "right_eye",
        "left_ear": "right_ear",
        "left_shoulder": "right_shoulder",
        "left_elbow": "right_elbow",
        "left_wrist": "right_wrist",
        "left_hip": "right_hip",
        "left_knee": "right_knee",
        "left_ankle": "right_ankle",
    }

    def __init__(self, *args, **kwargs):
        super(PersonKeypoints, self).__init__(*args, **kwargs)
        if len(cfg.MODEL.ROI_KEYPOINT_HEAD.KEYPOINT_NAME) > 0:
            self.NAMES = cfg.MODEL.ROI_KEYPOINT_HEAD.KEYPOINT_NAME
            self.FLIP_MAP = {
                l: r for l, r in PersonKeypoints._FLIP_MAP.items() if l in cfg.MODEL.ROI_KEYPOINT_HEAD.KEYPOINT_NAME
            }
        else:
            self.NAMES = PersonKeypoints._NAMES
            self.FLIP_MAP = PersonKeypoints._FLIP_MAP

        self.FLIP_INDS = self._create_flip_indices(self.NAMES, self.FLIP_MAP)
        self.CONNECTIONS = self._kp_connections(self.NAMES)

    def to_coco_format(self):
        coco_result = []
        for i in range(self.keypoints.shape[0]):
            coco_kps = [0] * len(PersonKeypoints._NAMES) * 3
            for ki, name in enumerate(self.NAMES):
                coco_kps[3 * PersonKeypoints._NAMES.index(name)] = self.keypoints[i, ki, 0].item()
                coco_kps[3 * PersonKeypoints._NAMES.index(name) + 1] = self.keypoints[i, ki, 1].item()
                coco_kps[3 * PersonKeypoints._NAMES.index(name) + 2] = self.keypoints[i, ki, 2].item()
            coco_result.append(coco_kps)
        return coco_result

    def _create_flip_indices(self, names, flip_map):
        full_flip_map = flip_map.copy()
        full_flip_map.update({v: k for k, v in flip_map.items()})
        flipped_names = [i if i not in full_flip_map else full_flip_map[i] for i in names]
        flip_indices = [names.index(i) for i in flipped_names]
        return torch.tensor(flip_indices)

    def _kp_connections(self, keypoints):
        CONNECTIONS = [
            ["left_eye", "right_eye"],
            ["left_eye", "nose"],
            ["right_eye", "nose"],
            ["right_eye", "right_ear"],
            ["left_eye", "left_ear"],
            ["right_shoulder", "right_elbow"],
            ["right_elbow", "right_wrist"],
            ["left_shoulder", "left_elbow"],
            ["left_elbow", "left_wrist"],
            ["right_hip", "right_knee"],
            ["right_knee", "right_ankle"],
            ["left_hip", "left_knee"],
            ["left_knee", "left_ankle"],
            ["right_shoulder", "left_shoulder"],
            ["right_hip", "left_hip"],
        ]

        kp_lines = [
            [keypoints.index(conn[0]), keypoints.index(conn[1])]
            for conn in CONNECTIONS
            if conn[0] in self.NAMES and conn[1] in self.NAMES
        ]
        return kp_lines


# TODO make this nicer, this is a direct translation from C2 (but removing the inner loop)
def keypoints_to_heat_map(keypoints, rois, heatmap_size):
    if rois.numel() == 0:
        return rois.new().long(), rois.new().long()
    offset_x = rois[:, 0]
    offset_y = rois[:, 1]
    scale_x = heatmap_size / (rois[:, 2] - rois[:, 0])
    scale_y = heatmap_size / (rois[:, 3] - rois[:, 1])

    offset_x = offset_x[:, None]
    offset_y = offset_y[:, None]
    scale_x = scale_x[:, None]
    scale_y = scale_y[:, None]

    x = keypoints[..., 0]
    y = keypoints[..., 1]

    x_boundary_inds = x == rois[:, 2][:, None]
    y_boundary_inds = y == rois[:, 3][:, None]

    x = (x - offset_x) * scale_x
    x = x.floor().long()
    y = (y - offset_y) * scale_y
    y = y.floor().long()

    x[x_boundary_inds] = heatmap_size - 1
    y[y_boundary_inds] = heatmap_size - 1

    valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size)
    vis = keypoints[..., 2] > 0
    valid = (valid_loc & vis).long()

    lin_ind = y * heatmap_size + x
    heatmaps = lin_ind * valid

    return heatmaps, valid