import numpy as np
import torch
from . import common_utils


class ResidualCoder(object):
    def __init__(self, code_size=7):
        super().__init__()
        self.code_size = code_size

    @staticmethod
    def encode_np(boxes, anchors):
        """
        :param boxes: (N, 7 + ?) x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param anchors: (N, 7 + ?)
        :return:
        """
        box_ndim = anchors.shape[-1]
        xa, ya, za, wa, la, ha, ra, *cas = np.split(anchors, box_ndim, axis=-1)
        xg, yg, zg, wg, lg, hg, rg, *cgs = np.split(boxes, box_ndim, axis=-1)

        # need to convert boxes to z-center format
        zg = zg + hg / 2
        za = za + ha / 2

        diagonal = np.sqrt(la ** 2 + wa ** 2)  # 4.3
        xt = (xg - xa) / diagonal
        yt = (yg - ya) / diagonal
        zt = (zg - za) / ha  # 1.6
        lt = np.log(lg / la)
        wt = np.log(wg / wa)
        ht = np.log(hg / ha)
        rt = rg - ra
        cts = [g - a for g, a in zip(cgs, cas)]
        return np.concatenate([xt, yt, zt, wt, lt, ht, rt, *cts], axis=-1)

    @staticmethod
    def decode_np(box_encodings, anchors):
        """
        :param box_encodings: (N, 7 + ?) x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param anchors: (N, 7 + ?)
        :return:
        """
        box_ndim = anchors.shape[-1]
        xa, ya, za, wa, la, ha, ra, *cas = np.split(anchors, box_ndim, axis=-1)
        xt, yt, zt, wt, lt, ht, rt, *cts = np.split(box_encodings, box_ndim, axis=-1)

        # need to convert box_encodings to z-bottom format
        za = za + ha / 2

        diagonal = np.sqrt(la ** 2 + wa ** 2)
        xg = xt * diagonal + xa
        yg = yt * diagonal + ya
        zg = zt * ha + za

        lg = np.exp(lt) * la
        wg = np.exp(wt) * wa
        hg = np.exp(ht) * ha
        rg = rt + ra
        zg = zg - hg / 2
        cgs = [t + a for t, a in zip(cts, cas)]
        return np.concatenate([xg, yg, zg, wg, lg, hg, rg, *cgs], axis=-1)

    @staticmethod
    def encode_torch(boxes, anchors):
        """
        :param boxes: (N, 7 + ?) x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param anchors: (N, 7 + ?)
        :return:
        """
        xa, ya, za, wa, la, ha, ra, *cas = torch.split(anchors, 1, dim=-1)
        xg, yg, zg, wg, lg, hg, rg, *cgs = torch.split(boxes, 1, dim=-1)

        za = za + ha / 2
        zg = zg + hg / 2

        diagonal = torch.sqrt(la ** 2 + wa ** 2)
        xt = (xg - xa) / diagonal
        yt = (yg - ya) / diagonal
        zt = (zg - za) / ha
        lt = torch.log(lg / la)
        wt = torch.log(wg / wa)
        ht = torch.log(hg / ha)
        rt = rg - ra

        cts = [g - a for g, a in zip(cgs, cas)]
        return torch.cat([xt, yt, zt, wt, lt, ht, rt, *cts], dim=-1)

    @staticmethod
    def decode_torch(box_encodings, anchors):
        """
        :param box_encodings: (N, 7 + ?) x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param anchors: (N, 7 + ?)
        :return:
        """
        xa, ya, za, wa, la, ha, ra, *cas = torch.split(anchors, 1, dim=-1)
        xt, yt, zt, wt, lt, ht, rt, *cts = torch.split(box_encodings, 1, dim=-1)

        za = za + ha / 2

        diagonal = torch.sqrt(la ** 2 + wa ** 2)
        xg = xt * diagonal + xa
        yg = yt * diagonal + ya
        zg = zt * ha + za

        lg = torch.exp(lt) * la
        wg = torch.exp(wt) * wa
        hg = torch.exp(ht) * ha
        rg = rt + ra

        zg = zg - hg / 2
        cgs = [t + a for t, a in zip(cts, cas)]
        return torch.cat([xg, yg, zg, wg, lg, hg, rg, *cgs], dim=-1)

    def decode_with_head_direction_torch(self, box_preds, anchors, dir_cls_preds,
                                         num_dir_bins, dir_offset, dir_limit_offset, use_binary_dir_classifier=False):
        """
        :param box_preds: (batch_size, N, 7 + ?), x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param anchors: (batch_size, N, 7 + ?), x, y, z, w, l, h, r, custom values, z is the box center in z-axis
        :param dir_cls_preds: (batch_size, H, W, num_anchors_per_locations*2)
        :return:
        """
        batch_box_preds = self.decode_torch(box_preds, anchors)

        if dir_cls_preds is not None:
            dir_cls_preds = dir_cls_preds.view(box_preds.shape[0], box_preds.shape[1], -1)
            if use_binary_dir_classifier:
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]
                opp_labels = (batch_box_preds[..., -1] > 0) ^ dir_labels.byte()
                batch_box_preds[..., -1] += torch.where(
                    opp_labels,
                    torch.tensor(np.pi).type_as(batch_box_preds),
                    torch.tensor(0.0).type_as(batch_box_preds)
                )
            else:
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]

                period = (2 * np.pi / num_dir_bins)
                dir_rot = common_utils.limit_period_torch(
                    batch_box_preds[..., 6] - dir_offset, dir_limit_offset, period
                )
                batch_box_preds[..., 6] = dir_rot + dir_offset + period * dir_labels.to(batch_box_preds.dtype)

        return batch_box_preds


if __name__ == '__main__':
    pass

[((19, 43, 19, 79), 'numpy.split', 'np.split', (), '', True, 'import numpy as np\n'), ((20, 43, 20, 77), 'numpy.split', 'np.split', (), '', True, 'import numpy as np\n'), ((26, 19, 26, 45), 'numpy.sqrt', 'np.sqrt', ({(26, 27, 26, 44): 'la ** 2 + wa ** 2'}, {}), '(la ** 2 + wa ** 2)', True, 'import numpy as np\n'), ((30, 13, 30, 28), 'numpy.log', 'np.log', ({(30, 20, 30, 27): 'lg / la'}, {}), '(lg / la)', True, 'import numpy as np\n'), ((31, 13, 31, 28), 'numpy.log', 'np.log', ({(31, 20, 31, 27): 'wg / wa'}, {}), '(wg / wa)', True, 'import numpy as np\n'), ((32, 13, 32, 28), 'numpy.log', 'np.log', ({(32, 20, 32, 27): 'hg / ha'}, {}), '(hg / ha)', True, 'import numpy as np\n'), ((35, 15, 35, 74), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((45, 43, 45, 79), 'numpy.split', 'np.split', (), '', True, 'import numpy as np\n'), ((46, 43, 46, 85), 'numpy.split', 'np.split', (), '', True, 'import numpy as np\n'), ((51, 19, 51, 45), 'numpy.sqrt', 'np.sqrt', ({(51, 27, 51, 44): 'la ** 2 + wa ** 2'}, {}), '(la ** 2 + wa ** 2)', True, 'import numpy as np\n'), ((62, 15, 62, 74), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((71, 43, 71, 74), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((72, 43, 72, 72), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((77, 19, 77, 48), 'torch.sqrt', 'torch.sqrt', ({(77, 30, 77, 47): 'la ** 2 + wa ** 2'}, {}), '(la ** 2 + wa ** 2)', False, 'import torch\n'), ((81, 13, 81, 31), 'torch.log', 'torch.log', ({(81, 23, 81, 30): 'lg / la'}, {}), '(lg / la)', False, 'import torch\n'), ((82, 13, 82, 31), 'torch.log', 'torch.log', ({(82, 23, 82, 30): 'wg / wa'}, {}), '(wg / wa)', False, 'import torch\n'), ((83, 13, 83, 31), 'torch.log', 'torch.log', ({(83, 23, 83, 30): 'hg / ha'}, {}), '(hg / ha)', False, 'import torch\n'), ((87, 15, 87, 68), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((96, 43, 96, 74), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((97, 43, 97, 80), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((101, 19, 101, 48), 'torch.sqrt', 'torch.sqrt', ({(101, 30, 101, 47): 'la ** 2 + wa ** 2'}, {}), '(la ** 2 + wa ** 2)', False, 'import torch\n'), ((113, 15, 113, 68), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((56, 13, 56, 23), 'numpy.exp', 'np.exp', ({(56, 20, 56, 22): 'lt'}, {}), '(lt)', True, 'import numpy as np\n'), ((57, 13, 57, 23), 'numpy.exp', 'np.exp', ({(57, 20, 57, 22): 'wt'}, {}), '(wt)', True, 'import numpy as np\n'), ((58, 13, 58, 23), 'numpy.exp', 'np.exp', ({(58, 20, 58, 22): 'ht'}, {}), '(ht)', True, 'import numpy as np\n'), ((106, 13, 106, 26), 'torch.exp', 'torch.exp', ({(106, 23, 106, 25): 'lt'}, {}), '(lt)', False, 'import torch\n'), ((107, 13, 107, 26), 'torch.exp', 'torch.exp', ({(107, 23, 107, 25): 'wt'}, {}), '(wt)', False, 'import torch\n'), ((108, 13, 108, 26), 'torch.exp', 'torch.exp', ({(108, 23, 108, 25): 'ht'}, {}), '(ht)', False, 'import torch\n'), ((128, 29, 128, 61), 'torch.max', 'torch.max', (), '', False, 'import torch\n'), ((136, 29, 136, 61), 'torch.max', 'torch.max', (), '', False, 'import torch\n'), ((132, 20, 132, 39), 'torch.tensor', 'torch.tensor', ({(132, 33, 132, 38): 'np.pi'}, {}), '(np.pi)', False, 'import torch\n'), ((133, 20, 133, 37), 'torch.tensor', 'torch.tensor', ({(133, 33, 133, 36): '(0.0)'}, {}), '(0.0)', False, 'import torch\n')]