import cv2

import numpy as np
import tensorflow as tf
import tf2onnx

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

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

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


class BayerToRGB(tf.keras.layers.Layer):
    """Convert Bayer image to RGB
    See: https://stanford.edu/class/ee367/reading/Demosaicing_ICASSP04.pdf
    See: https://github.com/cheind/pytorch-debayer/blob/master/debayer/modules.py#L231
    """

    def __init__(self):
        super().__init__()
        self.layout = "RGGB"
        self.max_val = 255.0

        self.kernels = tf.constant(
            np.array(
                [
                    # G at R/B locations
                    [
                        [0, 0, -1, 0, 0],
                        [0, 0, 2, 0, 0],
                        [-1, 2, 4, 2, -1],
                        [0, 0, 2, 0, 0],
                        [0, 0, -1, 0, 0],
                    ],
                    # R/B at G in R/B rows and B/R columns
                    [
                        [0, 0, 0.5, 0, 0],
                        [0, -1, 0, -1, 0],
                        [-1, 4, 5, 4, -1],
                        [0, -1, 0, -1, 0],
                        [0, 0, 0.5, 0, 0],
                    ],
                    # R/B at G in B/R rows and R/B columns
                    [
                        [0, 0, 0.5, 0, 0],
                        [0, -1, 4, -1, 0],
                        [-1, 0, 5, 0, -1],
                        [0, -1, 4, -1, 0],
                        [0, 0, 0.5, 0, 0],
                    ],
                    # R/B at B/R in B/R rows and B/R columns
                    [
                        [0, 0, -1.5, 0, 0],
                        [0, 2, 0, 2, 0],
                        [-1.5, 0, 6, 0, -1.5],
                        [0, 2, 0, 2, 0],
                        [0, 0, -1.5, 0, 0],
                    ],
                ],
                dtype=np.float32,
            )
            .reshape(1, 4, 5, 5)
            .transpose(2, 3, 0, 1)
            / 8.0
        )
        self.index = tf.constant(
            np.array(
                # Describes the kernel indices that calculate the corresponding RGB values for
                # the 2x2 layout (RGGB)  sub-structure
                [
                    # Destination R
                    [
                        [4, 1],  # identity, R at G in R row  B column
                        [2, 3],  # R at G in B row R column, R at B in B row R column
                    ],
                    # Destination G
                    [
                        [0, 4],
                        [4, 0],
                    ],
                    # Destination B
                    [
                        [3, 2],
                        [1, 4],
                    ],
                ]
            ).reshape(1, 3, 2, 2)
        )

    def call(self, img):
        H, W = tf.shape(img)[1], tf.shape(img)[2]

        # Pad the image
        tpad = img[:, 0:2, :, :]
        bpad = img[:, H - 2 : H, :, :]
        ipad = tf.concat([tpad, img, bpad], axis=1)
        lpad = ipad[:, :, 0:2, :]
        rpad = ipad[:, :, W - 2 : W, :]
        ipad = tf.concat([lpad, ipad, rpad], axis=2)

        # Convolve with kernels
        planes = tf.nn.conv2d(ipad, self.kernels, strides=[1, 1, 1, 1], padding="VALID")

        # Concatenate identity kernel
        planes = tf.concat([planes, img], axis=-1)

        # Gather values
        index_repeated = tf.tile(self.index, multiples=[1, 1, H // 2, W // 2])
        index_repeated = tf.transpose(index_repeated, perm=[0, 2, 3, 1])
        row_indices, col_indices = tf.meshgrid(tf.range(H), tf.range(W), indexing="ij")
        index_tensor = tf.stack([row_indices, col_indices], axis=-1)
        index_tensor = tf.expand_dims(index_tensor, axis=0)
        index_tensor = tf.expand_dims(index_tensor, axis=-2)
        index_tensor = tf.repeat(index_tensor, repeats=3, axis=-2)
        index_repeated = tf.expand_dims(index_repeated, axis=-1)
        indices = tf.concat([tf.cast(index_tensor, tf.int64), index_repeated], axis=-1)
        rgb = tf.gather_nd(planes, indices, batch_dims=1)

        if self.max_val == 255.0:
            # Make value range valid again
            rgb = tf.round(rgb)

        return rgb


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


def bayer_resize(img, size):
    """Resize a Bayer image by splitting color channels"""

    # Split the image into 4 channels
    r = img[:, 0::2, 0::2, 0]
    g1 = img[:, 0::2, 1::2, 0]
    g2 = img[:, 1::2, 0::2, 0]
    b = img[:, 1::2, 1::2, 0]
    bsplit = tf.stack([r, g1, g2, b], axis=-1)

    # Resize the image
    # Make sure the target size is divisible by 2
    size = (size[0] // 2, size[1] // 2)
    bsized = tf.image.resize(bsplit, size=size, method="bilinear")

    # Create a bayer image again
    img = tf.nn.depth_to_space(bsized, block_size=2)

    return img


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


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

        self.b2rgb = BayerToRGB()
        self.target_size = target_size
        self.fill_value = fill_value

    def calc_params(self, ishape):
        img_h, img_w = ishape[1], ishape[2]
        target_h, target_w = self.target_size

        scale = tf.minimum(target_w / img_w, target_h / img_h)
        new_w = tf.round(tf.cast(img_w, scale.dtype) * scale)
        new_h = tf.round(tf.cast(img_h, scale.dtype) * scale)
        new_w = tf.cast(new_w, tf.int32)
        new_h = tf.cast(new_h, tf.int32)
        new_w = new_w - (new_w % 2)
        new_h = new_h - (new_h % 2)

        pad_w = target_w - new_w
        pad_h = target_h - new_h
        pad_left = tf.cast(tf.floor(tf.cast(pad_w, tf.float32) / 2.0), tf.int32)
        pad_top = tf.cast(tf.floor(tf.cast(pad_h, tf.float32) / 2.0), tf.int32)
        pad_right = pad_w - pad_left
        pad_bottom = pad_h - pad_top
        paddings = [pad_top, pad_bottom, pad_left, pad_right]

        return paddings, scale, (new_w, new_h)

    def call(self, img):
        paddings, _, (nw, nh) = self.calc_params(tf.shape(img))

        # Resize the image and convert to RGB
        img = bayer_resize(img, (nh, nw))
        img = self.b2rgb(img)

        # Pad the image
        pad_top, pad_bottom, pad_left, pad_right = paddings
        img = tf.pad(
            img,
            paddings=[[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]],
            mode="CONSTANT",
            constant_values=self.fill_value,
        )

        return img


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


class DetPreprocess(tf.keras.layers.Layer):
    def __init__(self, target_size, fill_value=114):
        super(DetPreprocess, self).__init__()
        self.letterbox = Letterbox(target_size, fill_value)

    def call(self, img):
        """img: tf.Tensor of shape [batch, H, W, C], dtype=tf.uint8"""

        # Cast to float32 since TensorRT does not support uint8 layers
        img = tf.cast(img, tf.float32)

        img = self.letterbox(img)
        return img


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


def rgb2bayer(img):
    bayer = np.zeros((img.shape[0], img.shape[1]), dtype=img.dtype)
    bayer[0::2, 0::2] = img[0::2, 0::2, 0]
    bayer[0::2, 1::2] = img[0::2, 1::2, 1]
    bayer[1::2, 0::2] = img[1::2, 0::2, 1]
    bayer[1::2, 1::2] = img[1::2, 1::2, 2]
    return bayer


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


def main():

    img_path = "/RapidPoseTriangulation/scripts/../data/h1/54138969-img_003201.jpg"
    image = cv2.imread(img_path, 3)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = rgb2bayer(image)
    image = np.expand_dims(image, axis=-1)
    image = np.asarray(image, dtype=np.uint8)

    # Initialize the DetPreprocess module
    preprocess_model = tf.keras.Sequential()
    preprocess_model.add(DetPreprocess(target_size=det_target_size))
    det_dummy_input_a0 = tf.convert_to_tensor(
        np.expand_dims(image, axis=0), dtype=tf.uint8
    )
    det_dummy_output_a0 = preprocess_model(det_dummy_input_a0)
    print("\n", det_dummy_output_a0.shape, "\n")

    output_a0 = det_dummy_output_a0.numpy()
    output_a0 = np.squeeze(output_a0, axis=0)
    output_a0 = np.asarray(output_a0, dtype=np.uint8)
    output_a0 = cv2.cvtColor(output_a0, cv2.COLOR_RGB2BGR)
    cv2.imwrite(base_path + "det_preprocess.jpg", output_a0)

    # Export to ONNX
    input_signature = [tf.TensorSpec([None, None, None, 1], tf.uint8, name="x")]
    _, _ = tf2onnx.convert.from_keras(
        preprocess_model,
        input_signature,
        opset=11,
        output_path=base_path + "det_preprocess.onnx",
        target=["tensorrt"],
    )


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

if __name__ == "__main__":
    main()