import torch
import torch.nn as nn
import torch.nn.functional as F


def iuvmap_clean(U_uv, V_uv, Index_UV, AnnIndex=None):

    Index_UV_max = torch.argmax(Index_UV, dim=1).float()
    recon_Index_UV = []
    for i in range(Index_UV.size(1)):
        if i == 0:
            recon_Index_UV_i = torch.min(
                F.threshold(Index_UV_max + 1, 0.5, 0), -F.threshold(-Index_UV_max - 1, -1.5, 0)
            )
        else:
            recon_Index_UV_i = torch.min(
                F.threshold(Index_UV_max, i - 0.5, 0), -F.threshold(-Index_UV_max, -i - 0.5, 0)
            ) / float(i)
        recon_Index_UV.append(recon_Index_UV_i)
    recon_Index_UV = torch.stack(recon_Index_UV, dim=1)

    if AnnIndex is None:
        recon_Ann_Index = None
    else:
        AnnIndex_max = torch.argmax(AnnIndex, dim=1).float()
        recon_Ann_Index = []
        for i in range(AnnIndex.size(1)):
            if i == 0:
                recon_Ann_Index_i = torch.min(
                    F.threshold(AnnIndex_max + 1, 0.5, 0), -F.threshold(-AnnIndex_max - 1, -1.5, 0)
                )
            else:
                recon_Ann_Index_i = torch.min(
                    F.threshold(AnnIndex_max, i - 0.5, 0), -F.threshold(-AnnIndex_max, -i - 0.5, 0)
                ) / float(i)
            recon_Ann_Index.append(recon_Ann_Index_i)
        recon_Ann_Index = torch.stack(recon_Ann_Index, dim=1)

    recon_U = recon_Index_UV * U_uv
    recon_V = recon_Index_UV * V_uv

    return recon_U, recon_V, recon_Index_UV, recon_Ann_Index


def iuv_map2img(U_uv, V_uv, Index_UV, AnnIndex=None, uv_rois=None, ind_mapping=None, n_part=24):
    device_id = U_uv.get_device()
    batch_size = U_uv.size(0)
    K = U_uv.size(1)
    heatmap_size = U_uv.size(2)

    Index_UV_max = torch.argmax(Index_UV, dim=1)
    if AnnIndex is None:
        Index_UV_max = Index_UV_max.to(torch.int64)
    else:
        AnnIndex_max = torch.argmax(AnnIndex, dim=1)
        Index_UV_max = Index_UV_max * (AnnIndex_max > 0).to(torch.int64)

    outputs = []

    for batch_id in range(batch_size):

        output = torch.zeros([3, U_uv.size(2), U_uv.size(3)], dtype=torch.float32).cuda(device_id)
        output[0] = Index_UV_max[batch_id].to(torch.float32)
        if ind_mapping is None:
            output[0] /= float(K - 1)
        else:
            for ind in range(len(ind_mapping)):
                output[0][output[0] == ind] = ind_mapping[ind] * (1. / n_part)

        for part_id in range(0, K):
            CurrentU = U_uv[batch_id, part_id]
            CurrentV = V_uv[batch_id, part_id]
            output[1,
                   Index_UV_max[batch_id] == part_id] = CurrentU[Index_UV_max[batch_id] == part_id]
            output[2,
                   Index_UV_max[batch_id] == part_id] = CurrentV[Index_UV_max[batch_id] == part_id]

        if uv_rois is None:
            outputs.append(output.unsqueeze(0))
        else:
            roi_fg = uv_rois[batch_id][1:]

            # x1 = roi_fg[0]
            # x2 = roi_fg[2]
            # y1 = roi_fg[1]
            # y2 = roi_fg[3]

            w = roi_fg[2] - roi_fg[0]
            h = roi_fg[3] - roi_fg[1]

            aspect_ratio = float(w) / h

            if aspect_ratio < 1:
                new_size = [heatmap_size, max(int(heatmap_size * aspect_ratio), 1)]
                output = F.interpolate(output.unsqueeze(0), size=new_size, mode='nearest')
                paddingleft = int(0.5 * (heatmap_size - new_size[1]))
                output = F.pad(
                    output, pad=(paddingleft, heatmap_size - new_size[1] - paddingleft, 0, 0)
                )
            else:
                new_size = [max(int(heatmap_size / aspect_ratio), 1), heatmap_size]
                output = F.interpolate(output.unsqueeze(0), size=new_size, mode='nearest')
                paddingtop = int(0.5 * (heatmap_size - new_size[0]))
                output = F.pad(
                    output, pad=(0, 0, paddingtop, heatmap_size - new_size[0] - paddingtop)
                )

            outputs.append(output)

    return torch.cat(outputs, dim=0)


def iuv_img2map(uvimages, uv_rois=None, new_size=None, n_part=24):
    device_id = uvimages.get_device()
    batch_size = uvimages.size(0)
    uvimg_size = uvimages.size(-1)

    Index2mask = [[0], [1, 2], [3], [4], [5], [6], [7, 9], [8, 10], [11, 13], [12, 14], [15, 17],
                  [16, 18], [19, 21], [20, 22], [23, 24]]

    part_ind = torch.round(uvimages[:, 0, :, :] * n_part)
    part_u = uvimages[:, 1, :, :]
    part_v = uvimages[:, 2, :, :]

    recon_U = []
    recon_V = []
    recon_Index_UV = []
    recon_Ann_Index = []

    for i in range(n_part + 1):
        if i == 0:
            recon_Index_UV_i = torch.min(
                F.threshold(part_ind + 1, 0.5, 0), -F.threshold(-part_ind - 1, -1.5, 0)
            )
        else:
            recon_Index_UV_i = torch.min(
                F.threshold(part_ind, i - 0.5, 0), -F.threshold(-part_ind, -i - 0.5, 0)
            ) / float(i)
        recon_U_i = recon_Index_UV_i * part_u
        recon_V_i = recon_Index_UV_i * part_v

        recon_Index_UV.append(recon_Index_UV_i)
        recon_U.append(recon_U_i)
        recon_V.append(recon_V_i)

    for i in range(len(Index2mask)):
        if len(Index2mask[i]) == 1:
            recon_Ann_Index_i = recon_Index_UV[Index2mask[i][0]]
        elif len(Index2mask[i]) == 2:
            p_ind0 = Index2mask[i][0]
            p_ind1 = Index2mask[i][1]
            # recon_Ann_Index[:, i, :, :] = torch.where(recon_Index_UV[:, p_ind0, :, :] > 0.5, recon_Index_UV[:, p_ind0, :, :], recon_Index_UV[:, p_ind1, :, :])
            # recon_Ann_Index[:, i, :, :] = torch.eq(part_ind, p_ind0) | torch.eq(part_ind, p_ind1)
            recon_Ann_Index_i = recon_Index_UV[p_ind0] + recon_Index_UV[p_ind1]

        recon_Ann_Index.append(recon_Ann_Index_i)

    recon_U = torch.stack(recon_U, dim=1)
    recon_V = torch.stack(recon_V, dim=1)
    recon_Index_UV = torch.stack(recon_Index_UV, dim=1)
    recon_Ann_Index = torch.stack(recon_Ann_Index, dim=1)

    if uv_rois is None:
        return recon_U, recon_V, recon_Index_UV, recon_Ann_Index

    recon_U_roi = []
    recon_V_roi = []
    recon_Index_UV_roi = []
    recon_Ann_Index_roi = []

    if new_size is None:
        M = uvimg_size
    else:
        M = new_size

    for i in range(batch_size):
        roi_fg = uv_rois[i][1:]

        # x1 = roi_fg[0]
        # x2 = roi_fg[2]
        # y1 = roi_fg[1]
        # y2 = roi_fg[3]

        w = roi_fg[2] - roi_fg[0]
        h = roi_fg[3] - roi_fg[1]

        aspect_ratio = float(w) / h

        if aspect_ratio < 1:
            w_size = max(int(uvimg_size * aspect_ratio), 1)
            w_margin = int((uvimg_size - w_size) / 2)

            recon_U_roi_i = recon_U[i, :, :, w_margin:w_margin + w_size]
            recon_V_roi_i = recon_V[i, :, :, w_margin:w_margin + w_size]
            recon_Index_UV_roi_i = recon_Index_UV[i, :, :, w_margin:w_margin + w_size]
            recon_Ann_Index_roi_i = recon_Ann_Index[i, :, :, w_margin:w_margin + w_size]
        else:
            h_size = max(int(uvimg_size / aspect_ratio), 1)
            h_margin = int((uvimg_size - h_size) / 2)

            recon_U_roi_i = recon_U[i, :, h_margin:h_margin + h_size, :]
            recon_V_roi_i = recon_V[i, :, h_margin:h_margin + h_size, :]
            recon_Index_UV_roi_i = recon_Index_UV[i, :, h_margin:h_margin + h_size, :]
            recon_Ann_Index_roi_i = recon_Ann_Index[i, :, h_margin:h_margin + h_size, :]

        recon_U_roi_i = F.interpolate(recon_U_roi_i.unsqueeze(0), size=(M, M), mode='nearest')
        recon_V_roi_i = F.interpolate(recon_V_roi_i.unsqueeze(0), size=(M, M), mode='nearest')
        recon_Index_UV_roi_i = F.interpolate(
            recon_Index_UV_roi_i.unsqueeze(0), size=(M, M), mode='nearest'
        )
        recon_Ann_Index_roi_i = F.interpolate(
            recon_Ann_Index_roi_i.unsqueeze(0), size=(M, M), mode='nearest'
        )

        recon_U_roi.append(recon_U_roi_i)
        recon_V_roi.append(recon_V_roi_i)
        recon_Index_UV_roi.append(recon_Index_UV_roi_i)
        recon_Ann_Index_roi.append(recon_Ann_Index_roi_i)

    recon_U_roi = torch.cat(recon_U_roi, dim=0)
    recon_V_roi = torch.cat(recon_V_roi, dim=0)
    recon_Index_UV_roi = torch.cat(recon_Index_UV_roi, dim=0)
    recon_Ann_Index_roi = torch.cat(recon_Ann_Index_roi, dim=0)

    return recon_U_roi, recon_V_roi, recon_Index_UV_roi, recon_Ann_Index_roi


def seg_img2map(segimages, uv_rois=None, new_size=None, n_part=24):
    device_id = segimages.get_device()
    batch_size = segimages.size(0)
    uvimg_size = segimages.size(-1)

    part_ind = torch.round(segimages[:, 0, :, :] * n_part)

    recon_Index_UV = []

    for i in range(n_part + 1):
        if i == 0:
            recon_Index_UV_i = torch.min(
                F.threshold(part_ind + 1, 0.5, 0), -F.threshold(-part_ind - 1, -1.5, 0)
            )
        else:
            recon_Index_UV_i = torch.min(
                F.threshold(part_ind, i - 0.5, 0), -F.threshold(-part_ind, -i - 0.5, 0)
            ) / float(i)

        recon_Index_UV.append(recon_Index_UV_i)

    recon_Index_UV = torch.stack(recon_Index_UV, dim=1)

    if uv_rois is None:
        return None, None, recon_Index_UV, None

    recon_Index_UV_roi = []

    if new_size is None:
        M = uvimg_size
    else:
        M = new_size

    for i in range(batch_size):
        roi_fg = uv_rois[i][1:]

        # x1 = roi_fg[0]
        # x2 = roi_fg[2]
        # y1 = roi_fg[1]
        # y2 = roi_fg[3]

        w = roi_fg[2] - roi_fg[0]
        h = roi_fg[3] - roi_fg[1]

        aspect_ratio = float(w) / h

        if aspect_ratio < 1:
            w_size = max(int(uvimg_size * aspect_ratio), 1)
            w_margin = int((uvimg_size - w_size) / 2)

            recon_Index_UV_roi_i = recon_Index_UV[i, :, :, w_margin:w_margin + w_size]
        else:
            h_size = max(int(uvimg_size / aspect_ratio), 1)
            h_margin = int((uvimg_size - h_size) / 2)

            recon_Index_UV_roi_i = recon_Index_UV[i, :, h_margin:h_margin + h_size, :]

        recon_Index_UV_roi_i = F.interpolate(
            recon_Index_UV_roi_i.unsqueeze(0), size=(M, M), mode='nearest'
        )

        recon_Index_UV_roi.append(recon_Index_UV_roi_i)

    recon_Index_UV_roi = torch.cat(recon_Index_UV_roi, dim=0)

    return None, None, recon_Index_UV_roi, None