# Code adapted from: https://github.com/gpastal24/ViTPose-Pytorch
from .topdown_heatmap_simple_head import TopdownHeatmapSimpleHead
import torch
from .backbone import ViT
import torchvision.transforms.functional as F
import torch.nn as nn
import tops

model_large = dict(
    type="TopDown",
    pretrained=None,
    backbone=dict(
        type="ViT",
        img_size=(256, 192),
        patch_size=16,
        embed_dim=1024,
        depth=24,
        num_heads=16,
        ratio=1,
        use_checkpoint=False,
        mlp_ratio=4,
        qkv_bias=True,
        drop_path_rate=0.5,
    ),
    keypoint_head=dict(
        type="TopdownHeatmapSimpleHead",
        in_channels=1024,
        num_deconv_layers=2,
        num_deconv_filters=(256, 256),
        num_deconv_kernels=(4, 4),
        extra=dict(
            final_conv_kernel=1,
        ),
        out_channels=17,
        loss_keypoint=dict(type="JointsMSELoss", use_target_weight=True),
    ),
    train_cfg=dict(),
    test_cfg=dict(
        flip_test=True,
        post_process="default",
        shift_heatmap=False,
        target_type="GaussianHeatmap",
        modulate_kernel=11,
        use_udp=True,
    ),
)


model_base = dict(
    type="TopDown",
    pretrained=None,
    backbone=dict(
        type="ViT",
        img_size=(256, 192),
        patch_size=16,
        embed_dim=768,
        depth=12,
        num_heads=12,
        ratio=1,
        use_checkpoint=False,
        mlp_ratio=4,
        qkv_bias=True,
        drop_path_rate=0.3,
    ),
    keypoint_head=dict(
        type="TopdownHeatmapSimpleHead",
        in_channels=768,
        num_deconv_layers=2,
        num_deconv_filters=(256, 256),
        num_deconv_kernels=(4, 4),
        extra=dict(
            final_conv_kernel=1,
        ),
        out_channels=17,
        loss_keypoint=dict(type="JointsMSELoss", use_target_weight=True),
    ),
    train_cfg=dict(),
    test_cfg=dict(),
)
model_huge = dict(
    type='TopDown',
    pretrained=None,
    backbone=dict(
        type='ViT',
        img_size=(256, 192),
        patch_size=16,
        embed_dim=1280,
        depth=32,
        num_heads=16,
        ratio=1,
        use_checkpoint=False,
        mlp_ratio=4,
        qkv_bias=True,
        drop_path_rate=0.55,
    ),
    keypoint_head=dict(
        type='TopdownHeatmapSimpleHead',
        in_channels=1280,
        num_deconv_layers=2,
        num_deconv_filters=(256, 256),
        num_deconv_kernels=(4, 4),
        extra=dict(final_conv_kernel=1, ),
        out_channels=17,
        loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),
    train_cfg=dict(),
    test_cfg=dict(
        flip_test=True,
        post_process='default',
        shift_heatmap=False,
        target_type="GaussianHeatmap",
        modulate_kernel=11,
        use_udp=True))


class VitPoseModel(nn.Module):
    def __init__(self, model_name):
        super().__init__()
        assert model_name  in ["vit_base", "vit_large", "vit_huge"]
        model = {
            "vit_base": model_base,
            "vit_large": model_large,
            "vit_huge": model_huge
        }[model_name]
        weight_url = {
            "vit_base": "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/90235a26-3b8c-427d-a264-c68155abecdcfcfcd8a9-0388-4575-b85b-607d3c0a9b149bef8f0f-a0f9-4662-a561-1b47ba5f1636",
            "vit_large": "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/a580a44c-0afd-43ac-a2cb-9956c32b1d1a78c51ecb-81bb-4345-8710-13904cb9dbbe0703db2d-8534-42e0-ac4d-518ab51fe7db",
            "vit_huge": "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/a33b6ada-4d2f-4ef7-8f83-b33f58b69f5b2a62e181-2131-467d-a900-027157a08571d761fad4-785b-4b84-8596-8932c7857e44"
        }[model_name]
        file_name = {
            "vit_base":  "vit-b-multi-coco-595b5e128b.pth",
            "vit_large": "vit-l-multi-coco-9475d27cec.pth",
            "vit_huge":  "vit-h-multi-coco-dbc06d4337.pth",
        }[model_name]
        # Set check_hash to true if you suspect a download error.
        weight_path = tops.download_file(
            weight_url, file_name=file_name, check_hash=True)

        self.keypoint_head = tops.to_cuda(TopdownHeatmapSimpleHead(
            in_channels=model["keypoint_head"]["in_channels"],
            out_channels=model["keypoint_head"]["out_channels"],
            num_deconv_filters=model["keypoint_head"]["num_deconv_filters"],
            num_deconv_kernels=model["keypoint_head"]["num_deconv_kernels"],
            num_deconv_layers=model["keypoint_head"]["num_deconv_layers"],
            extra=model["keypoint_head"]["extra"],
        ))
        # print(head)
        self.backbone = tops.to_cuda(ViT(
            img_size=model["backbone"]["img_size"],
            patch_size=model["backbone"]["patch_size"],
            embed_dim=model["backbone"]["embed_dim"],
            depth=model["backbone"]["depth"],
            num_heads=model["backbone"]["num_heads"],
            ratio=model["backbone"]["ratio"],
            mlp_ratio=model["backbone"]["mlp_ratio"],
            qkv_bias=model["backbone"]["qkv_bias"],
            drop_path_rate=model["backbone"]["drop_path_rate"],
        ))
        ckpt = torch.load(weight_path, map_location=tops.get_device())
        self.load_state_dict(ckpt["state_dict"])
        self.backbone.eval()
        self.keypoint_head.eval()

    def forward(self, img: torch.Tensor, boxes_ltrb: torch.Tensor):
        assert img.ndim == 3
        assert img.dtype == torch.uint8
        assert boxes_ltrb.ndim == 2 and boxes_ltrb.shape[1] == 4
        assert boxes_ltrb.dtype == torch.long
        boxes_ltrb = boxes_ltrb.clamp(0)
        padded_boxes = torch.zeros_like(boxes_ltrb)
        images = torch.zeros((len(boxes_ltrb), 3, 256, 192), device=img.device, dtype=torch.float32)

        for i, (x0, y0, x1, y1) in enumerate(boxes_ltrb):
            x1 = min(img.shape[-1], x1)
            y1 = min(img.shape[-2], y1)
            correction_factor = 256 / 192 * (x1 - x0) / (y1 - y0)
            if correction_factor > 1:
                # increase y side
                center = y0 + (y1 - y0) // 2
                length = (y1-y0).mul(correction_factor).round().long()
                y0_new = center - length.div(2).long()
                y1_new = center + length.div(2).long()
                image_crop = img[:, y0:y1, x0:x1]
                # print(y1,y2,x1,x2)
                pad = ((y0_new-y0).abs(), (y1_new-y1).abs())
#                pad = (int(abs(y0_new-y0))), int(abs(y1_new-y1))
                image_crop = torch.nn.functional.pad(image_crop, [*(0, 0), *pad])
                padded_boxes[i] = torch.tensor([x0, y0_new, x1, y1_new])
            else:
                center = x0 + (x1 - x0) // 2
                length = (x1-x0).div(correction_factor).round().long()
                x0_new = center - length.div(2).long()
                x1_new = center + length.div(2).long()
                image_crop = img[:, y0:y1, x0:x1]
                pad = ((x0_new-x0).abs(), (x1_new-x1).abs())
                image_crop = torch.nn.functional.pad(image_crop, [*pad, ])
                padded_boxes[i] = torch.tensor([x0_new, y0, x1_new, y1])
            image_crop = F.resize(image_crop.float(), (256, 192), antialias=True)
            image_crop = F.normalize(image_crop, mean=[0.485*255, 0.456*255,
                                     0.406*255], std=[0.229*255, 0.224*255, 0.225*255])
            images[i] = image_crop

        x = self.backbone(images)
        out = self.keypoint_head(x)
        pts = torch.empty((out.shape[0], out.shape[1], 3), dtype=torch.float32, device=img.device)
        # For each human, for each joint: y, x, confidence
        b, indices = torch.max(out, dim=2)
        b, indices = torch.max(b, dim=2)

        c, indicesc = torch.max(out, dim=3)
        c, indicesc = torch.max(c, dim=2)
        dim1 = torch.tensor(1./64, device=img.device)
        dim2 = torch.tensor(1./48, device=img.device)
        for i in range(0, out.shape[0]):
            pts[i, :, 0] = indicesc[i, :] * dim1 * (padded_boxes[i][3] - padded_boxes[i][1]) + padded_boxes[i][1]
            pts[i, :, 1] = indices[i, :] * dim2 * (padded_boxes[i][2] - padded_boxes[i][0]) + padded_boxes[i][0]
            pts[i, :, 2] = c[i, :]
        pts = pts[:, :, [1, 0, 2]]
        return pts