import cv2
import mediapipe as mp
import torch
import numpy as np
import torch.nn.functional as F
from rembg import remove
from rembg.session_factory import new_session
from PIL import Image
from torchvision.models import detection

from lib.pymafx.core import constants
from lib.common.cloth_extraction import load_segmentation
from torchvision import transforms


def transform_to_tensor(res, mean=None, std=None, is_tensor=False):
    all_ops = []
    if res is not None:
        all_ops.append(transforms.Resize(size=res))
    if not is_tensor:
        all_ops.append(transforms.ToTensor())
    if mean is not None and std is not None:
        all_ops.append(transforms.Normalize(mean=mean, std=std))
    return transforms.Compose(all_ops)


def aug_matrix(w1, h1, w2, h2):
    dx = (w2 - w1) / 2.0
    dy = (h2 - h1) / 2.0

    matrix_trans = np.array([[1.0, 0, dx], [0, 1.0, dy], [0, 0, 1.0]])

    scale = np.min([float(w2) / w1, float(h2) / h1])

    M = get_affine_matrix(center=(w2 / 2.0, h2 / 2.0), translate=(0, 0), scale=scale)

    M = np.array(M + [0.0, 0.0, 1.0]).reshape(3, 3)
    M = M.dot(matrix_trans)

    return M


def get_affine_matrix(center, translate, scale):
    cx, cy = center
    tx, ty = translate

    M = [1, 0, 0, 0, 1, 0]
    M = [x * scale for x in M]

    # Apply translation and of center translation: RSS * C^-1
    M[2] += M[0] * (-cx) + M[1] * (-cy)
    M[5] += M[3] * (-cx) + M[4] * (-cy)

    # Apply center translation: T * C * RSS * C^-1
    M[2] += cx + tx
    M[5] += cy + ty
    return M


def load_img(img_file):

    img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    if not img_file.endswith("png"):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)

    return img


def get_keypoints(image):

    def collect_xyv(x, body=True):
        lmk = x.landmark
        all_lmks = []
        for i in range(len(lmk)):
            visibility = lmk[i].visibility if body else 1.0
            all_lmks.append(torch.Tensor([lmk[i].x, lmk[i].y, lmk[i].z, visibility]))
        return torch.stack(all_lmks).view(-1, 4)

    mp_holistic = mp.solutions.holistic

    with mp_holistic.Holistic(
            static_image_mode=True,
            model_complexity=2,
    ) as holistic:
        results = holistic.process(image)

    fake_kps = torch.zeros(33, 4)

    result = {}
    result["body"] = collect_xyv(results.pose_landmarks) if results.pose_landmarks else fake_kps
    result["lhand"] = collect_xyv(results.left_hand_landmarks, False) if results.left_hand_landmarks else fake_kps
    result["rhand"] = collect_xyv(results.right_hand_landmarks, False) if results.right_hand_landmarks else fake_kps
    result["face"] = collect_xyv(results.face_landmarks, False) if results.face_landmarks else fake_kps

    return result


def get_pymafx(image, landmarks):

    # image [3,512,512]

    item = {'img_body': F.interpolate(image.unsqueeze(0), size=224, mode='bicubic', align_corners=True)[0]}

    for part in ['lhand', 'rhand', 'face']:
        kp2d = landmarks[part]
        kp2d_valid = kp2d[kp2d[:, 3] > 0.]
        if len(kp2d_valid) > 0:
            bbox = [min(kp2d_valid[:, 0]), min(kp2d_valid[:, 1]), max(kp2d_valid[:, 0]), max(kp2d_valid[:, 1])]
            center_part = [(bbox[2] + bbox[0]) / 2., (bbox[3] + bbox[1]) / 2.]
            scale_part = 2. * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2

        # handle invalid part keypoints
        if len(kp2d_valid) < 1 or scale_part < 0.01:
            center_part = [0, 0]
            scale_part = 0.5
            kp2d[:, 3] = 0

        center_part = torch.tensor(center_part).float()

        theta_part = torch.zeros(1, 2, 3)
        theta_part[:, 0, 0] = scale_part
        theta_part[:, 1, 1] = scale_part
        theta_part[:, :, -1] = center_part

        grid = F.affine_grid(theta_part, torch.Size([1, 3, 224, 224]), align_corners=False)
        img_part = F.grid_sample(image.unsqueeze(0), grid, align_corners=False).squeeze(0).float()

        item[f'img_{part}'] = img_part

        theta_i_inv = torch.zeros_like(theta_part)
        theta_i_inv[:, 0, 0] = 1. / theta_part[:, 0, 0]
        theta_i_inv[:, 1, 1] = 1. / theta_part[:, 1, 1]
        theta_i_inv[:, :, -1] = -theta_part[:, :, -1] / theta_part[:, 0, 0].unsqueeze(-1)
        item[f'{part}_theta_inv'] = theta_i_inv[0]

    return item


def expand_bbox(bbox, width, height, ratio=0.1):

    bbox = np.around(bbox).astype(np.int16)
    bbox_width = bbox[2] - bbox[0]
    bbox_height = bbox[3] - bbox[1]

    bbox[1] = max(bbox[1] - bbox_height * ratio, 0)
    bbox[3] = min(bbox[3] + bbox_height * ratio, height)
    bbox[0] = max(bbox[0] - bbox_width * ratio, 0)
    bbox[2] = min(bbox[2] + bbox_width * ratio, width)

    return bbox


def remove_floats(mask):

    # 1. find all the contours
    # 2. fillPoly "True" for the largest one
    # 3. fillPoly "False" for its childrens

    new_mask = np.zeros(mask.shape)
    cnts, hier = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cnt_index = sorted(range(len(cnts)), key=lambda k: cv2.contourArea(cnts[k]), reverse=True)
    body_cnt = cnts[cnt_index[0]]
    childs_cnt_idx = np.where(np.array(hier)[0, :, -1] == cnt_index[0])[0]
    childs_cnt = [cnts[idx] for idx in childs_cnt_idx]
    cv2.fillPoly(new_mask, [body_cnt], 1)
    cv2.fillPoly(new_mask, childs_cnt, 0)

    return new_mask


def process_image(img_file, hps_type, single, input_res=512):

    img_raw = load_img(img_file)

    in_height, in_width = img_raw.shape[:2]
    M = aug_matrix(in_width, in_height, input_res * 2, input_res * 2)

    # from rectangle to square by padding (input_res*2, input_res*2)
    img_square = cv2.warpAffine(img_raw, M[0:2, :], (input_res * 2, input_res * 2), flags=cv2.INTER_CUBIC)

    # detection for bbox
    detector = detection.maskrcnn_resnet50_fpn(weights=detection.MaskRCNN_ResNet50_FPN_V2_Weights)
    detector.eval()
    predictions = detector([torch.from_numpy(img_square).permute(2, 0, 1) / 255.])[0]

    if single:
        top_score = predictions["scores"][predictions["labels"] == 1].max()
        human_ids = torch.where(predictions["scores"] == top_score)[0]
    else:
        human_ids = torch.logical_and(predictions["labels"] == 1, predictions["scores"] > 0.9).nonzero().squeeze(1)

    boxes = predictions["boxes"][human_ids, :].detach().cpu().numpy()
    masks = predictions["masks"][human_ids, :, :].permute(0, 2, 3, 1).detach().cpu().numpy()

    width = boxes[:, 2] - boxes[:, 0]  #(N,)
    height = boxes[:, 3] - boxes[:, 1]  #(N,)
    center = np.array([(boxes[:, 0] + boxes[:, 2]) / 2.0, (boxes[:, 1] + boxes[:, 3]) / 2.0]).T  #(N,2)
    scale = np.array([width, height]).max(axis=0) / 90.

    img_icon_lst = []
    img_crop_lst = []
    img_hps_lst = []
    img_mask_lst = []
    uncrop_param_lst = []
    landmark_lst = []
    hands_visibility_lst = []
    img_pymafx_lst = []

    uncrop_param = {
        "center": center,
        "scale": scale,
        "ori_shape": [in_height, in_width],
        "box_shape": [input_res, input_res],
        "crop_shape": [input_res * 2, input_res * 2, 3],
        "M": M,
    }

    for idx in range(len(boxes)):

        # mask out the pixels of others
        if len(masks) > 1:
            mask_detection = (masks[np.arange(len(masks)) != idx]).max(axis=0)
        else:
            mask_detection = masks[0] * 0.

        img_crop, _ = crop(
            np.concatenate([img_square, (mask_detection < 0.4) * 255], axis=2), center[idx], scale[idx], [input_res, input_res])

        # get accurate segmentation mask of focus person
        img_rembg = remove(img_crop, post_process_mask=True, session=new_session("u2net"))
        img_mask = remove_floats(img_rembg[:, :, [3]])

        # required image tensors / arrays

        # img_icon  (tensor): (-1, 1),          [3,512,512]
        # img_hps   (tensor): (-2.11, 2.44),    [3,224,224]

        # img_np    (array): (0, 255),          [512,512,3]
        # img_rembg (array): (0, 255),          [512,512,4]
        # img_mask  (array): (0, 1),            [512,512,1]
        # img_crop  (array): (0, 255),          [512,512,4]

        mean_icon = std_icon = (0.5, 0.5, 0.5)
        img_np = (img_rembg[..., :3] * img_mask).astype(np.uint8)
        img_icon = transform_to_tensor(512, mean_icon, std_icon)(Image.fromarray(img_np)) * torch.tensor(img_mask).permute(
            2, 0, 1)
        img_hps = transform_to_tensor(224, constants.IMG_NORM_MEAN, constants.IMG_NORM_STD)(Image.fromarray(img_np))

        landmarks = get_keypoints(img_np)

        if hps_type == 'pymafx':
            img_pymafx_lst.append(
                get_pymafx(
                    transform_to_tensor(512, constants.IMG_NORM_MEAN, constants.IMG_NORM_STD)(Image.fromarray(img_np)),
                    landmarks))

        img_crop_lst.append(torch.tensor(img_crop).permute(2, 0, 1) / 255.0)
        img_icon_lst.append(img_icon)
        img_hps_lst.append(img_hps)
        img_mask_lst.append(torch.tensor(img_mask[..., 0]))
        uncrop_param_lst.append(uncrop_param)
        landmark_lst.append(landmarks['body'])

        hands_visibility = [True, True]
        if landmarks['lhand'][:, -1].mean() == 0.:
            hands_visibility[0] = False
        if landmarks['rhand'][:, -1].mean() == 0.:
            hands_visibility[1] = False
        hands_visibility_lst.append(hands_visibility)

    return_dict = {
        "img_icon": torch.stack(img_icon_lst).float(),  #[N, 3, res, res]
        "img_crop": torch.stack(img_crop_lst).float(),  #[N, 4, res, res]               
        "img_hps": torch.stack(img_hps_lst).float(),  #[N, 3, res, res]
        "img_raw": img_raw,  #[H, W, 3]
        "img_mask": torch.stack(img_mask_lst).float(),  #[N, res, res]
        "uncrop_param": uncrop_param,
        "landmark": torch.stack(landmark_lst),  #[N, 33, 4]
        "hands_visibility": hands_visibility_lst,
    }

    img_pymafx = {}

    if len(img_pymafx_lst) > 0:
        for idx in range(len(img_pymafx_lst)):
            for key in img_pymafx_lst[idx].keys():
                if key not in img_pymafx.keys():
                    img_pymafx[key] = [img_pymafx_lst[idx][key]]
                else:
                    img_pymafx[key] += [img_pymafx_lst[idx][key]]

        for key in img_pymafx.keys():
            img_pymafx[key] = torch.stack(img_pymafx[key]).float()

        return_dict.update({"img_pymafx": img_pymafx})

    return return_dict


def get_transform(center, scale, res):
    """Generate transformation matrix."""
    h = 100 * scale
    t = np.zeros((3, 3))
    t[0, 0] = float(res[1]) / h
    t[1, 1] = float(res[0]) / h
    t[0, 2] = res[1] * (-float(center[0]) / h + 0.5)
    t[1, 2] = res[0] * (-float(center[1]) / h + 0.5)
    t[2, 2] = 1

    return t


def transform(pt, center, scale, res, invert=0):
    """Transform pixel location to different reference."""
    t = get_transform(center, scale, res)
    if invert:
        t = np.linalg.inv(t)
    new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.0]).T
    new_pt = np.dot(t, new_pt)
    return np.around(new_pt[:2]).astype(np.int16)


def crop(img, center, scale, res):
    """Crop image according to the supplied bounding box."""

    img_height, img_width = img.shape[:2]

    # Upper left point
    ul = np.array(transform([0, 0], center, scale, res, invert=1))

    # Bottom right point
    br = np.array(transform(res, center, scale, res, invert=1))

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    new_img = np.zeros(new_shape)

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], img_width) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], img_height) - ul[1]

    # Range to sample from original image
    old_x = max(0, ul[0]), min(img_width, br[0])
    old_y = max(0, ul[1]), min(img_height, br[1])

    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1], old_x[0]:old_x[1]]
    new_img = F.interpolate(
        torch.tensor(new_img).permute(2, 0, 1).unsqueeze(0), res, mode='bilinear').permute(0, 2, 3,
                                                                                           1)[0].numpy().astype(np.uint8)

    return new_img, (old_x, new_x, old_y, new_y, new_shape)


def crop_segmentation(org_coord, res, cropping_parameters):
    old_x, new_x, old_y, new_y, new_shape = cropping_parameters

    new_coord = np.zeros((org_coord.shape))
    new_coord[:, 0] = new_x[0] + (org_coord[:, 0] - old_x[0])
    new_coord[:, 1] = new_y[0] + (org_coord[:, 1] - old_y[0])

    new_coord[:, 0] = res[0] * (new_coord[:, 0] / new_shape[1])
    new_coord[:, 1] = res[1] * (new_coord[:, 1] / new_shape[0])

    return new_coord


def corner_align(ul, br):

    if ul[1] - ul[0] != br[1] - br[0]:
        ul[1] = ul[0] + br[1] - br[0]

    return ul, br


def uncrop(img, center, scale, orig_shape):
    """'Undo' the image cropping/resizing.
    This function is used when evaluating mask/part segmentation.
    """

    res = img.shape[:2]

    # Upper left point
    ul = np.array(transform([0, 0], center, scale, res, invert=1))
    # Bottom right point
    br = np.array(transform(res, center, scale, res, invert=1))

    # quick fix
    ul, br = corner_align(ul, br)

    # size of cropped image
    crop_shape = [br[1] - ul[1], br[0] - ul[0]]
    new_img = np.zeros(orig_shape, dtype=np.uint8)

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1]

    # Range to sample from original image
    old_x = max(0, ul[0]), min(orig_shape[1], br[0])
    old_y = max(0, ul[1]), min(orig_shape[0], br[1])

    img = np.array(Image.fromarray(img.astype(np.uint8)).resize(crop_shape))

    new_img[old_y[0]:old_y[1], old_x[0]:old_x[1]] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]]

    return new_img


def rot_aa(aa, rot):
    """Rotate axis angle parameters."""
    # pose parameters
    R = np.array([
        [np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
        [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
        [0, 0, 1],
    ])
    # find the rotation of the body in camera frame
    per_rdg, _ = cv2.Rodrigues(aa)
    # apply the global rotation to the global orientation
    resrot, _ = cv2.Rodrigues(np.dot(R, per_rdg))
    aa = (resrot.T)[0]
    return aa


def flip_img(img):
    """Flip rgb images or masks.
    channels come last, e.g. (256,256,3).
    """
    img = np.fliplr(img)
    return img


def flip_kp(kp, is_smpl=False):
    """Flip keypoints."""
    if len(kp) == 24:
        if is_smpl:
            flipped_parts = constants.SMPL_JOINTS_FLIP_PERM
        else:
            flipped_parts = constants.J24_FLIP_PERM
    elif len(kp) == 49:
        if is_smpl:
            flipped_parts = constants.SMPL_J49_FLIP_PERM
        else:
            flipped_parts = constants.J49_FLIP_PERM
    kp = kp[flipped_parts]
    kp[:, 0] = -kp[:, 0]
    return kp


def flip_pose(pose):
    """Flip pose.
    The flipping is based on SMPL parameters.
    """
    flipped_parts = constants.SMPL_POSE_FLIP_PERM
    pose = pose[flipped_parts]
    # we also negate the second and the third dimension of the axis-angle
    pose[1::3] = -pose[1::3]
    pose[2::3] = -pose[2::3]
    return pose


def normalize_2d_kp(kp_2d, crop_size=224, inv=False):
    # Normalize keypoints between -1, 1
    if not inv:
        ratio = 1.0 / crop_size
        kp_2d = 2.0 * kp_2d * ratio - 1.0
    else:
        ratio = 1.0 / crop_size
        kp_2d = (kp_2d + 1.0) / (2 * ratio)

    return kp_2d


def visualize_landmarks(image, joints, color):

    img_w, img_h = image.shape[:2]

    for joint in joints:
        image = cv2.circle(image, (int(joint[0] * img_w), int(joint[1] * img_h)), 5, color)

    return image


def generate_heatmap(joints, heatmap_size, sigma=1, joints_vis=None):
    """
    param joints:  [num_joints, 3]
    param joints_vis: [num_joints, 3]
    return: target, target_weight(1: visible, 0: invisible)
    """
    num_joints = joints.shape[0]
    device = joints.device
    cur_device = torch.device(device.type, device.index)
    if not hasattr(heatmap_size, "__len__"):
        # width  height
        heatmap_size = [heatmap_size, heatmap_size]
    assert len(heatmap_size) == 2
    target_weight = np.ones((num_joints, 1), dtype=np.float32)
    if joints_vis is not None:
        target_weight[:, 0] = joints_vis[:, 0]
    target = torch.zeros(
        (num_joints, heatmap_size[1], heatmap_size[0]),
        dtype=torch.float32,
        device=cur_device,
    )

    tmp_size = sigma * 3

    for joint_id in range(num_joints):
        mu_x = int(joints[joint_id][0] * heatmap_size[0] + 0.5)
        mu_y = int(joints[joint_id][1] * heatmap_size[1] + 0.5)
        # Check that any part of the gaussian is in-bounds
        ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
        br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
        if (ul[0] >= heatmap_size[0] or ul[1] >= heatmap_size[1] or br[0] < 0 or br[1] < 0):
            # If not, just return the image as is
            target_weight[joint_id] = 0
            continue

        # # Generate gaussian
        size = 2 * tmp_size + 1
        # x = np.arange(0, size, 1, np.float32)
        # y = x[:, np.newaxis]
        # x0 = y0 = size // 2
        # # The gaussian is not normalized, we want the center value to equal 1
        # g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
        # g = torch.from_numpy(g.astype(np.float32))

        x = torch.arange(0, size, dtype=torch.float32, device=cur_device)
        y = x.unsqueeze(-1)
        x0 = y0 = size // 2
        # The gaussian is not normalized, we want the center value to equal 1
        g = torch.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))

        # Usable gaussian range
        g_x = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0]
        g_y = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1]
        # Image range
        img_x = max(0, ul[0]), min(br[0], heatmap_size[0])
        img_y = max(0, ul[1]), min(br[1], heatmap_size[1])

        v = target_weight[joint_id]
        if v > 0.5:
            target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]

    return target, target_weight