import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F

# ==================================================================================================

base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/"
det_target_size = (320, 320)
pose_target_size = (384, 288)

# ==================================================================================================


class Letterbox(nn.Module):
    def __init__(self, target_size, fill_value=128):
        """Resize and pad image while keeping aspect ratio"""
        super(Letterbox, self).__init__()

        self.target_size = target_size
        self.fill_value = fill_value

    def calc_params_and_crop(self, ishape, bbox=None):
        ih0, iw0 = ishape[1], ishape[2]
        th, tw = self.target_size

        if bbox is not None:
            bbox = bbox[0].float()
            x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]

            # Slightly increase bbox size
            factor = 1.25
            w = x2 - x1
            h = y2 - y1
            x1 -= w * (factor - 1) / 2
            x2 += w * (factor - 1) / 2
            y1 -= h * (factor - 1) / 2
            y2 += h * (factor - 1) / 2

            zero = torch.tensor(0)
            x1 = torch.max(x1, zero).to(torch.int64)
            y1 = torch.max(y1, zero).to(torch.int64)
            x2 = torch.min(x2, iw0).to(torch.int64)
            y2 = torch.min(y2, ih0).to(torch.int64)
            bbox = torch.stack((x1, y1, x2, y2), dim=0).unsqueeze(0)

            ih = y2 - y1
            iw = x2 - x1
        else:
            ih, iw = ih0, iw0

        scale = torch.min(tw / iw, th / ih)
        nw = torch.round(iw * scale)
        nh = torch.round(ih * scale)

        pad_w = tw - nw
        pad_h = th - nh
        pad_left = pad_w // 2
        pad_top = pad_h // 2
        pad_right = pad_w - pad_left
        pad_bottom = pad_h - pad_top
        paddings = (pad_left, pad_right, pad_top, pad_bottom)

        return paddings, scale, (nw, nh), bbox

    def forward(self, img, bbox=None):
        paddings, _, (nw, nh), bbox = self.calc_params_and_crop(img.shape, bbox)

        # Optional: Crop the image
        if bbox is not None:
            x1, y1, x2, y2 = bbox[0, 0], bbox[0, 1], bbox[0, 2], bbox[0, 3]
            img = img.to(torch.float32)
            img = img[:, y1:y2, x1:x2, :]

        # Resize the image
        img = img.to(torch.float32)
        img = F.interpolate(
            img.permute(0, 3, 1, 2),
            size=(nh, nw),
            mode="bilinear",
            align_corners=False,
        )
        img = img.permute(0, 2, 3, 1)
        img = img.round()

        # Pad the image
        img = F.pad(
            img.permute(0, 3, 1, 2),
            pad=paddings,
            mode="constant",
            value=self.fill_value,
        )
        img = img.permute(0, 2, 3, 1)
        canvas = img

        return canvas


# ==================================================================================================


class DetPreprocess(nn.Module):
    def __init__(self, target_size, fill_value=114):
        super(DetPreprocess, self).__init__()
        self.letterbox = Letterbox(target_size, fill_value)

    def forward(self, img):
        # img: torch.Tensor of shape [batch, H, W, C], dtype=torch.uint8
        img = self.letterbox(img, None)
        return img


# ==================================================================================================


class DetPostprocess(nn.Module):
    def __init__(self, target_size):
        super(DetPostprocess, self).__init__()

        self.target_size = target_size
        self.letterbox = Letterbox(target_size)

    def forward(self, img, boxes):
        paddings, scale, _, _ = self.letterbox.calc_params_and_crop(img.shape, None)

        boxes = boxes.float()
        boxes[:, :, 0] -= paddings[0]
        boxes[:, :, 2] -= paddings[0]
        boxes[:, :, 1] -= paddings[2]
        boxes[:, :, 3] -= paddings[2]

        zero = torch.tensor(0)
        boxes = torch.max(boxes, zero)

        th, tw = self.target_size
        pad_w = paddings[0] + paddings[1]
        pad_h = paddings[2] + paddings[3]
        max_w = tw - pad_w - 1
        max_h = th - pad_h - 1
        b0 = boxes[:, :, 0]
        b1 = boxes[:, :, 1]
        b2 = boxes[:, :, 2]
        b3 = boxes[:, :, 3]
        b0 = torch.min(b0, max_w)
        b1 = torch.min(b1, max_h)
        b2 = torch.min(b2, max_w)
        b3 = torch.min(b3, max_h)
        boxes[:, :, 0] = b0
        boxes[:, :, 1] = b1
        boxes[:, :, 2] = b2
        boxes[:, :, 3] = b3

        boxes[:, :, 0:4] /= scale
        return boxes


# ==================================================================================================


class PosePreprocess(nn.Module):
    def __init__(self, target_size, fill_value=114):
        super(PosePreprocess, self).__init__()
        self.letterbox = Letterbox(target_size, fill_value)

    def forward(self, img, bbox):
        # img: torch.Tensor of shape [1, H, W, C], dtype=torch.uint8
        # bbox: torch.Tensor of shape [1, 4], dtype=torch.float32
        img = self.letterbox(img, bbox)
        return img


# ==================================================================================================


class PosePostprocess(nn.Module):
    def __init__(self, target_size):
        super(PosePostprocess, self).__init__()

        self.target_size = target_size
        self.letterbox = Letterbox(target_size)

    def forward(self, img, bbox, keypoints):
        paddings, scale, _, bbox = self.letterbox.calc_params_and_crop(img.shape, bbox)

        kp = keypoints.float()
        kp[:, :, 0] -= paddings[0]
        kp[:, :, 1] -= paddings[2]

        zero = torch.tensor(0)
        kp = torch.max(kp, zero)

        th, tw = self.target_size
        pad_w = paddings[0] + paddings[1]
        pad_h = paddings[2] + paddings[3]
        max_w = tw - pad_w - 1
        max_h = th - pad_h - 1
        k0 = kp[:, :, 0]
        k1 = kp[:, :, 1]
        k0 = torch.min(k0, max_w)
        k1 = torch.min(k1, max_h)
        kp[:, :, 0] = k0
        kp[:, :, 1] = k1

        kp[:, :, 0:2] /= scale

        kp[:, :, 0] += bbox[0, 0]
        kp[:, :, 1] += bbox[0, 1]
        return kp


# ==================================================================================================


def main():

    img_path = "/RapidPoseTriangulation/scripts/../data/h1/54138969-img_003201.jpg"
    image = cv2.imread(img_path, 3)

    # Initialize the DetPreprocess module
    preprocess_model = DetPreprocess(target_size=det_target_size)
    det_dummy_input_a0 = torch.from_numpy(image).unsqueeze(0)

    # Export to ONNX
    torch.onnx.export(
        preprocess_model,
        det_dummy_input_a0,
        base_path + "det_preprocess.onnx",
        opset_version=11,
        input_names=["input_image"],
        output_names=["preprocessed_image"],
        dynamic_axes={
            "input_image": {0: "batch_size", 1: "height", 2: "width"},
            "preprocessed_image": {0: "batch_size"},
        },
    )

    # Initialize the DetPostprocess module
    postprocess_model = DetPostprocess(target_size=det_target_size)
    det_dummy_input_b0 = torch.from_numpy(image).unsqueeze(0)
    det_dummy_input_b1 = torch.rand(1, 10, 5)

    # Export to ONNX
    torch.onnx.export(
        postprocess_model,
        (det_dummy_input_b0, det_dummy_input_b1),
        base_path + "det_postprocess.onnx",
        opset_version=11,
        input_names=["input_image", "boxes"],
        output_names=["output_boxes"],
        dynamic_axes={
            "input_image": {0: "batch_size", 1: "height", 2: "width"},
            "boxes": {0: "batch_size", 1: "num_boxes"},
            "output_boxes": {0: "batch_size", 1: "num_boxes"},
        },
    )

    # Initialize the PosePreprocess module
    preprocess_model = PosePreprocess(target_size=pose_target_size)
    det_dummy_input_c0 = torch.from_numpy(image).unsqueeze(0)
    det_dummy_input_c1 = torch.tensor([[10, 10, 90, 40]]).to(torch.int32)

    # Export to ONNX
    torch.onnx.export(
        preprocess_model,
        (det_dummy_input_c0, det_dummy_input_c1),
        base_path + "pose_preprocess.onnx",
        opset_version=11,
        input_names=["input_image", "bbox"],
        output_names=["preprocessed_image"],
        dynamic_axes={
            "input_image": {0: "batch_size", 1: "height", 2: "width"},
            "preprocessed_image": {0: "batch_size"},
        },
    )

    # Initialize the PosePostprocess module
    postprocess_model = PosePostprocess(target_size=pose_target_size)
    det_dummy_input_d0 = torch.from_numpy(image).unsqueeze(0)
    det_dummy_input_d1 = torch.tensor([[10, 10, 90, 40]]).to(torch.int32)
    det_dummy_input_d2 = torch.rand(1, 17, 3)

    # Export to ONNX
    torch.onnx.export(
        postprocess_model,
        (det_dummy_input_d0, det_dummy_input_d1, det_dummy_input_d2),
        base_path + "pose_postprocess.onnx",
        opset_version=11,
        input_names=["input_image", "bbox", "keypoints"],
        output_names=["output_keypoints"],
        dynamic_axes={
            "input_image": {0: "batch_size", 1: "height", 2: "width"},
            "output_keypoints": {0: "batch_size"},
        },
    )


# ==================================================================================================

if __name__ == "__main__":
    main()