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

[((498, 534), 'numpy.split', 'np.split', (['anchors', 'box_ndim'], {'axis': '(-1)'}), '(anchors, box_ndim, axis=-1)\n', (506, 534), True, 'import numpy as np\n'), ((578, 612), 'numpy.split', 'np.split', (['boxes', 'box_ndim'], {'axis': '(-1)'}), '(boxes, box_ndim, axis=-1)\n', (586, 612), True, 'import numpy as np\n'), ((735, 761), 'numpy.sqrt', 'np.sqrt', (['(la ** 2 + wa ** 2)'], {}), '(la ** 2 + wa ** 2)\n', (742, 761), True, 'import numpy as np\n'), ((885, 900), 'numpy.log', 'np.log', (['(lg / la)'], {}), '(lg / la)\n', (891, 900), True, 'import numpy as np\n'), ((914, 929), 'numpy.log', 'np.log', (['(wg / wa)'], {}), '(wg / wa)\n', (920, 929), True, 'import numpy as np\n'), ((943, 958), 'numpy.log', 'np.log', (['(hg / ha)'], {}), '(hg / ha)\n', (949, 958), True, 'import numpy as np\n'), ((1043, 1102), 'numpy.concatenate', 'np.concatenate', (['[xt, yt, zt, wt, lt, ht, rt, *cts]'], {'axis': '(-1)'}), '([xt, yt, zt, wt, lt, ht, rt, *cts], axis=-1)\n', (1057, 1102), True, 'import numpy as np\n'), ((1428, 1464), 'numpy.split', 'np.split', (['anchors', 'box_ndim'], {'axis': '(-1)'}), '(anchors, box_ndim, axis=-1)\n', (1436, 1464), True, 'import numpy as np\n'), ((1508, 1550), 'numpy.split', 'np.split', (['box_encodings', 'box_ndim'], {'axis': '(-1)'}), '(box_encodings, box_ndim, axis=-1)\n', (1516, 1550), True, 'import numpy as np\n'), ((1656, 1682), 'numpy.sqrt', 'np.sqrt', (['(la ** 2 + wa ** 2)'], {}), '(la ** 2 + wa ** 2)\n', (1663, 1682), True, 'import numpy as np\n'), ((1970, 2029), 'numpy.concatenate', 'np.concatenate', (['[xg, yg, zg, wg, lg, hg, rg, *cgs]'], {'axis': '(-1)'}), '([xg, yg, zg, wg, lg, hg, rg, *cgs], axis=-1)\n', (1984, 2029), True, 'import numpy as np\n'), ((2305, 2336), 'torch.split', 'torch.split', (['anchors', '(1)'], {'dim': '(-1)'}), '(anchors, 1, dim=-1)\n', (2316, 2336), False, 'import torch\n'), ((2380, 2409), 'torch.split', 'torch.split', (['boxes', '(1)'], {'dim': '(-1)'}), '(boxes, 1, dim=-1)\n', (2391, 2409), False, 'import torch\n'), ((2481, 2510), 'torch.sqrt', 'torch.sqrt', (['(la ** 2 + wa ** 2)'], {}), '(la ** 2 + wa ** 2)\n', (2491, 2510), False, 'import torch\n'), ((2620, 2638), 'torch.log', 'torch.log', (['(lg / la)'], {}), '(lg / la)\n', (2629, 2638), False, 'import torch\n'), ((2652, 2670), 'torch.log', 'torch.log', (['(wg / wa)'], {}), '(wg / wa)\n', (2661, 2670), False, 'import torch\n'), ((2684, 2702), 'torch.log', 'torch.log', (['(hg / ha)'], {}), '(hg / ha)\n', (2693, 2702), False, 'import torch\n'), ((2788, 2841), 'torch.cat', 'torch.cat', (['[xt, yt, zt, wt, lt, ht, rt, *cts]'], {'dim': '(-1)'}), '([xt, yt, zt, wt, lt, ht, rt, *cts], dim=-1)\n', (2797, 2841), False, 'import torch\n'), ((3133, 3164), 'torch.split', 'torch.split', (['anchors', '(1)'], {'dim': '(-1)'}), '(anchors, 1, dim=-1)\n', (3144, 3164), False, 'import torch\n'), ((3208, 3245), 'torch.split', 'torch.split', (['box_encodings', '(1)'], {'dim': '(-1)'}), '(box_encodings, 1, dim=-1)\n', (3219, 3245), False, 'import torch\n'), ((3292, 3321), 'torch.sqrt', 'torch.sqrt', (['(la ** 2 + wa ** 2)'], {}), '(la ** 2 + wa ** 2)\n', (3302, 3321), False, 'import torch\n'), ((3619, 3672), 'torch.cat', 'torch.cat', (['[xg, yg, zg, wg, lg, hg, rg, *cgs]'], {'dim': '(-1)'}), '([xg, yg, zg, wg, lg, hg, rg, *cgs], dim=-1)\n', (3628, 3672), False, 'import torch\n'), ((1787, 1797), 'numpy.exp', 'np.exp', (['lt'], {}), '(lt)\n', (1793, 1797), True, 'import numpy as np\n'), ((1816, 1826), 'numpy.exp', 'np.exp', (['wt'], {}), '(wt)\n', (1822, 1826), True, 'import numpy as np\n'), ((1845, 1855), 'numpy.exp', 'np.exp', (['ht'], {}), '(ht)\n', (1851, 1855), True, 'import numpy as np\n'), ((3426, 3439), 'torch.exp', 'torch.exp', (['lt'], {}), '(lt)\n', (3435, 3439), False, 'import torch\n'), ((3458, 3471), 'torch.exp', 'torch.exp', (['wt'], {}), '(wt)\n', (3467, 3471), False, 'import torch\n'), ((3490, 3503), 'torch.exp', 'torch.exp', (['ht'], {}), '(ht)\n', (3499, 3503), False, 'import torch\n'), ((4489, 4521), 'torch.max', 'torch.max', (['dir_cls_preds'], {'dim': '(-1)'}), '(dir_cls_preds, dim=-1)\n', (4498, 4521), False, 'import torch\n'), ((4888, 4920), 'torch.max', 'torch.max', (['dir_cls_preds'], {'dim': '(-1)'}), '(dir_cls_preds, dim=-1)\n', (4897, 4920), False, 'import torch\n'), ((4714, 4733), 'torch.tensor', 'torch.tensor', (['np.pi'], {}), '(np.pi)\n', (4726, 4733), False, 'import torch\n'), ((4780, 4797), 'torch.tensor', 'torch.tensor', (['(0.0)'], {}), '(0.0)\n', (4792, 4797), False, 'import torch\n')]