"""
Implements the CLIP Image processor.
"""

from tvm import s_tir, tirx
from tvm.relax.frontend.nn import Module, Tensor, op
from tvm.script import tirx as T


def _var(dtype, size=1):
    return T.sblock_alloc_buffer((size,), dtype, scope="local")


class ImageProcessor(Module):
    def __init__(self):
        super().__init__()

    def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]):
        loop_x, loop_y = sch.get_loops(block)[-2:]
        xo, xi = sch.split(loop_x, factors=[tile[0], None])
        yo, yi = sch.split(loop_y, factors=[tile[1], None])
        sch.reorder(xo, yo, xi, yi)
        t = sch.fuse(xo, yo)
        ty, tx = sch.split(t, factors=[None, bdx])
        sch.bind(ty, "threadIdx.y")
        sch.bind(tx, "threadIdx.x")

    def resize(self, image: Tensor, params):  # image layout:NCHW
        assert 4 == image.ndim, "image should be 4D data tensor"
        assert 3 == image.shape[1], "image layout should be NCHW"

        def get_output_image_size(image: Tensor):
            h = image.shape[2]
            w = image.shape[3]

            if "height" in params and "width" in params:
                return (params["height"], params["width"])
            elif "shortest_edge" in params:
                short = tirx.Select(w < h, w, h)
                long = tirx.Select(w > h, w, h)
                requested_new_short = params["shortest_edge"]
                new_short, new_long = (
                    tirx.Cast("int64", requested_new_short),
                    tirx.Cast(
                        "int64",
                        requested_new_short
                        * tirx.div(
                            tirx.Cast("float32", long),
                            tirx.Cast("float32", short),
                        ),
                    ),
                )
                ret_h = tirx.Select(w <= h, new_long, new_short)
                ret_w = tirx.Select(w <= h, new_short, new_long)
                return (ret_h, ret_w)
            elif "hd_transform" in params:
                hd_num = 4 if "hd_num" not in params else params["hd_num"]
                pad_num = 336 if "pad_num" not in params else params["pad_num"]
                ratio = tirx.Select(
                    w > h,
                    tirx.div(tirx.Cast("float32", w), tirx.Cast("float32", h)),
                    tirx.div(tirx.Cast("float32", h), tirx.Cast("float32", w)),
                )

                scale = tirx.ceil(tirx.sqrt(tirx.Cast("float32", hd_num) * ratio))

                scale = tirx.Select(
                    (scale * tirx.ceil(tirx.div(scale, ratio))) > hd_num,
                    scale - 1,
                    scale,
                )
                scale = tirx.Cast("int64", scale)

                new_w = tirx.Select(
                    w >= h,
                    scale * pad_num,
                    tirx.Cast("int64", tirx.div(scale * pad_num, ratio)),
                )
                new_h = tirx.Select(
                    w >= h,
                    tirx.Cast("int64", tirx.div(new_w, ratio)),
                    scale * pad_num,
                )
                return (new_h, new_w)
            else:
                assert False, "not supported resize parameter"

        new_h, new_w = get_output_image_size(image)
        out = op.interpolate(image, (new_h, new_w), data_layout="NCHW", mode="linear")
        return out

    def crop(self, image: Tensor, crop_size):
        assert 4 == image.ndim, "image should be 4D data tensor"
        assert 3 == image.shape[1], "image layout should be NCHW"

        def create_crop_func(dtype):  # , top, bottom, left, right):
            @T.prim_func(s_tir=True)
            def crop_func(
                image: T.handle,
                out: T.handle,
                top: T.int64(),
                bottom: T.int64(),
                left: T.int64(),
                right: T.int64(),
            ):
                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
                out_buf = T.match_buffer(out, (n, c, bottom - top, right - left), dtype=dtype)
                out_h = bottom - top
                out_w = right - left
                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
                        for h_idx, w_idx in T.grid(out_h, out_w):
                            with T.sblock("crop"):
                                if (h_idx + T.int64(top)) < h and (w_idx + T.int64(left)) < w:
                                    T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
                                    T.reads(image_buf[n_idx, c_idx, h_idx + top, w_idx + left])
                                    out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[
                                        n_idx, c_idx, h_idx + top, w_idx + left
                                    ]

            sch = s_tir.Schedule(crop_func)
            self.apply_schedule(sch, sch.get_sblock("crop"))
            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)

        n, c, orig_height, orig_width = image.shape
        crop_height = crop_size["height"]
        crop_width = crop_size["width"]

        top = (orig_height - crop_height) // 2
        bottom = orig_height - top

        left = (orig_width - crop_width) // 2
        right = orig_width - left

        out = op.tensor_ir_op(
            create_crop_func(image.dtype),
            "crop",
            [image, top, bottom, left, right],
            [Tensor.placeholder([n, c, crop_height, crop_width], image.dtype)],
        )
        return out

    def rescale(self, image: Tensor, rescale_factor=1 / 255.0, o_dtype="float32"):
        assert 4 == image.ndim, "image should be 4D data tensor"
        assert 3 == image.shape[1], "image layout should be NCHW"

        def create_rescale_func(rescale_factor, dtype, o_dtype):
            @T.prim_func(s_tir=True)
            def rescale_func(image: T.handle, out: T.handle):
                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
                out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype)

                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
                        for h_idx, w_idx in T.grid(h, w):
                            with T.sblock("rescale"):
                                T.reads(image_buf[n_idx, c_idx, h_idx, w_idx])
                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
                                if h_idx < h and w_idx < w:
                                    out_buf[n_idx, c_idx, h_idx, w_idx] = (
                                        T.cast(
                                            image_buf[n_idx, c_idx, h_idx, w_idx],
                                            o_dtype,
                                        )
                                        * rescale_factor
                                    )

            sch = s_tir.Schedule(rescale_func)
            self.apply_schedule(sch, sch.get_sblock("rescale"))
            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)

        out = op.tensor_ir_op(
            create_rescale_func(rescale_factor, image.dtype, o_dtype),
            "rescale",
            [image],
            [Tensor.placeholder(image.shape, o_dtype)],
        )
        return out

    def normalize(self, image: Tensor, o_dtype="float32"):
        assert 4 == image.ndim, "image should be 4D data tensor"
        assert 3 == image.shape[1], "image layout should be NCHW"

        def create_normalize_func(dtype, o_dtype):
            @T.prim_func(s_tir=True)
            def normalize_func(image: T.handle, out: T.handle):
                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
                out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype)
                mean = _var(o_dtype, 3)
                stddev = _var(o_dtype, 3)

                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
                        for h_idx, w_idx in T.grid(h, w):
                            with T.sblock("normalize"):
                                T.reads(
                                    image_buf[n_idx, c_idx, h_idx, w_idx],
                                    mean[c_idx],
                                    stddev[c_idx],
                                )
                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
                                with T.init():
                                    mean[0] = 0.48145466
                                    stddev[0] = 0.26862954
                                    mean[1] = 0.4578275
                                    stddev[1] = 0.26130258
                                    mean[2] = 0.40821073
                                    stddev[2] = 0.27577711
                                if h_idx < h and w_idx < w:
                                    out_buf[n_idx, c_idx, h_idx, w_idx] = (
                                        T.cast(
                                            image_buf[n_idx, c_idx, h_idx, w_idx],
                                            o_dtype,
                                        )
                                        - mean[c_idx]
                                    ) / stddev[c_idx]

            sch = s_tir.Schedule(normalize_func)
            self.apply_schedule(sch, sch.get_sblock("normalize"))
            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)

        out = op.tensor_ir_op(
            create_normalize_func(image.dtype, o_dtype),
            "normalize",
            [image],
            [Tensor.placeholder(image.shape, o_dtype)],
        )
        return out

    def pad(self, image: Tensor, dtype="uint8"):
        assert 4 == image.ndim, "image should be 4D data tensor"
        assert 3 == image.shape[1], "image layout should be NCHW"

        def create_pad_func(left, right, fill=255):
            @T.prim_func(s_tir=True)
            def pad_func(image: T.handle, out: T.handle, t: T.int64(), b: T.int64()):
                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
                out_buf = T.match_buffer(out, (n, c, h + t + b, w + left + right), dtype=dtype)
                out_h = h + t + b
                out_w = w + left + right

                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
                        for h_idx, w_idx in T.grid(out_h, out_w):
                            with T.sblock("pad"):
                                T.reads(image_buf[n_idx, c_idx, h_idx, w_idx])
                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
                                if h_idx < t or h_idx > h + b or w_idx < left or w_idx > w + right:
                                    out_buf[n_idx, c_idx, h_idx, w_idx] = fill
                                else:
                                    out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[
                                        n_idx, c_idx, h_idx - t, w_idx - left
                                    ]

            sch = s_tir.Schedule(pad_func)
            self.apply_schedule(sch, sch.get_sblock("pad"))
            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)

        h = image.shape[2]
        tar = tirx.truncdiv(h + 335, 336) * 336
        t = tirx.div(tar - h, 2)
        b = tar - h - t
        left = 0
        right = 0

        n, c, h, w = image.shape
        out = op.tensor_ir_op(
            create_pad_func(left, right),
            "pad",
            [image, t, b],
            [Tensor.placeholder((n, c, tar, w), image.dtype)],
        )
        return out

    def preprocess(self, pixel_values):
        return pixel_values