"""RTMW-x whole-body 2D keypoints (ONNX Runtime, FP16 model, no flip test) -> COCO-63.

The crop (top_down_affine and helpers) and SimCC decoding (get_simcc_maximum) are adapted from
RTMLib (https://github.com/Tau-J/rtmlib), licensed under the Apache License 2.0, see
licenses/RTMLib_LICENSE.txt.
"""

import hashlib
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
import torch
from tqdm import tqdm

from exmr.preprocessing.keypoints.mappings import COCO133_TO_COCO63

INPUT_SIZE = (288, 384)  # width, height
MEAN = (123.675, 116.28, 103.53)
STD = (58.395, 57.12, 57.375)
BOX_PADDING = 1.25


class RTMWKeypoints:
    """2D keypoints of one person per frame, given its bounding boxes.

    Args:
        checkpoint: RTMW-x 384x288 ONNX model with FP32 inputs/outputs
        sha256: expected checksum of the checkpoint (None to skip the check)
    """

    def __init__(self, checkpoint, sha256=None, batch_size=32, num_workers=8, device_id=0):
        if sha256 is not None and _sha256(checkpoint) != sha256:
            raise ValueError(f"Unexpected RTMW checkpoint (SHA256 mismatch): {checkpoint}")
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.session = ort.InferenceSession(
            str(checkpoint),
            providers=[("CUDAExecutionProvider", {"device_id": device_id}), "CPUExecutionProvider"],
        )
        if self.session.get_providers()[0] != "CUDAExecutionProvider":
            raise RuntimeError("ONNX Runtime CUDA provider unavailable for RTMW")

    def __call__(self, frame_dir, bbx_xyxy):
        """frame_dir with %06d.jpg frames, bbx_xyxy (L, 4) -> (L, 63, 3) [x, y, raw score]"""
        boxes = bbx_xyxy.cpu().numpy()
        loader = torch.utils.data.DataLoader(
            _FrameDataset(frame_dir, boxes),
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=False,
            **({"prefetch_factor": 4} if self.num_workers else {}),
        )
        inp = self.session.get_inputs()[0]
        out_names = [v.name for v in self.session.get_outputs()]
        result = np.empty((len(boxes), 63, 3), dtype=np.float32)
        cursor = 0
        with tqdm(total=len(boxes), desc="RTMW-x keypoints", unit="frame") as pbar:
            for batch in loader:
                images = batch["input"].numpy()
                n = len(images)
                # pad the last batch to the full batch size: ONNX Runtime + cuDNN 9 set up the
                # convolutions anew (~7 s) for every new input shape
                images = np.pad(images, [(0, self.batch_size - n)] + [(0, 0)] * 3, mode="edge")
                simcc_x, simcc_y = self.session.run(out_names, {inp.name: images})
                simcc_x, simcc_y = simcc_x[:n], simcc_y[:n]
                center, scale = batch["center"].numpy(), batch["scale"].numpy()
                xy, scores = _decode(simcc_x, simcc_y, center, scale)
                kp = np.concatenate((xy, scores[..., None]), axis=-1)
                result[cursor : cursor + len(kp)] = kp[:, COCO133_TO_COCO63]
                cursor += len(kp)
                pbar.update(len(kp))
        if cursor != len(boxes) or not np.isfinite(result).all():
            raise RuntimeError("Incomplete or non-finite RTMW output")
        return torch.from_numpy(result)


class _FrameDataset(torch.utils.data.Dataset):
    def __init__(self, frame_dir, boxes):
        self.paths = [str(Path(frame_dir) / f"{i:06d}.jpg") for i in range(len(boxes))]
        self.boxes = boxes
        self.mean = np.array(MEAN)  # float64 normalization, as in the evaluated setup
        self.std = np.array(STD)

    def __len__(self):
        return len(self.boxes)

    def __getitem__(self, idx):
        image = cv2.imread(self.paths[idx])
        if image is None:
            raise ValueError(f"Cannot decode {self.paths[idx]}")
        center, scale = _bbox_xyxy2cs(self.boxes[idx], padding=BOX_PADDING)
        crop, scale = _top_down_affine(INPUT_SIZE, scale, center, image)
        crop = (crop - self.mean) / self.std
        chw = np.ascontiguousarray(crop.transpose(2, 0, 1), dtype=np.float32)
        return {
            "input": torch.from_numpy(chw),
            "center": torch.from_numpy(np.asarray(center, dtype=np.float32)),
            "scale": torch.from_numpy(np.asarray(scale, dtype=np.float32)),
        }


def _decode(simcc_x, simcc_y, centers, scales):
    """SimCC outputs -> keypoints in image coordinates and raw scores."""
    locs, scores = _get_simcc_maximum(simcc_x, simcc_y)
    keypoints = locs / 2.0  # SimCC split ratio
    input_size = np.asarray(INPUT_SIZE, dtype=np.float32)
    centers = np.asarray(centers, dtype=np.float32)
    scales = np.asarray(scales, dtype=np.float32)
    keypoints = keypoints / input_size * scales[:, None, :]
    keypoints = keypoints + centers[:, None, :] - scales[:, None, :] / 2
    return keypoints, scores.astype(np.float32)


def _sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


# ----- adapted from RTMLib (rtmlib/tools/pose_estimation/{pre,post}_processings.py) ----- #


def _bbox_xyxy2cs(bbox, padding=1.0):
    """(4,) box [x1, y1, x2, y2] -> center (2,) and padded scale (2,)."""
    x1, y1, x2, y2 = np.hsplit(bbox[None, :], [1, 2, 3])
    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
    scale = np.hstack([x2 - x1, y2 - y1]) * padding
    return center[0], scale[0]


def _get_3rd_point(a, b):
    """b + (a - b) rotated by 90 degrees anticlockwise."""
    direction = a - b
    return b + np.r_[-direction[1], direction[0]]


def _get_warp_matrix(center, scale, output_size):
    """2x3 affine matrix that maps the box (center, scale) onto an output_size image."""
    src_w = scale[0]
    dst_w, dst_h = output_size

    src_dir = np.array([0.0, src_w * -0.5])  # no rotation
    dst_dir = np.array([0.0, dst_w * -0.5])

    shift = np.array((0.0, 0.0))
    src = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale * shift
    src[1, :] = center + src_dir + scale * shift
    src[2, :] = _get_3rd_point(src[0, :], src[1, :])

    dst = np.zeros((3, 2), dtype=np.float32)
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
    return cv2.getAffineTransform(np.float32(src), np.float32(dst))


def _top_down_affine(input_size, bbox_scale, bbox_center, img):
    """Crop the box from img to input_size (box grown to the input aspect ratio)."""
    w, h = input_size
    aspect_ratio = w / h
    b_w, b_h = np.hsplit(bbox_scale, [1])
    bbox_scale = np.where(
        b_w > b_h * aspect_ratio,
        np.hstack([b_w, b_w / aspect_ratio]),
        np.hstack([b_h * aspect_ratio, b_h]),
    )
    warp_mat = _get_warp_matrix(bbox_center, bbox_scale, output_size=(w, h))
    img = cv2.warpAffine(img, warp_mat, (int(w), int(h)), flags=cv2.INTER_LINEAR)
    return img, bbox_scale


def _get_simcc_maximum(simcc_x, simcc_y):
    """SimCC (N, K, W) x/y -> locations (N, K, 2) and scores (N, K) of the maxima."""
    N, K, _ = simcc_x.shape
    simcc_x = simcc_x.reshape(N * K, -1)
    simcc_y = simcc_y.reshape(N * K, -1)
    locs = np.stack((np.argmax(simcc_x, axis=1), np.argmax(simcc_y, axis=1)), axis=-1)
    locs = locs.astype(np.float32)
    vals = 0.5 * (np.amax(simcc_x, axis=1) + np.amax(simcc_y, axis=1))
    locs[vals <= 0.0] = -1
    return locs.reshape(N, K, 2), vals.reshape(N, K)