import torch.nn as nn
import torch
from torch.nn.functional import mse_loss
# for NEW: losses when calculated on keypoint locations
# see https://kornia.readthedocs.io/en/latest/_modules/kornia/geometry/subpix/dsnt.html
# from kornia.geometry import dsnt            # old kornia version  
from kornia.geometry.subpix import dsnt     # kornia 0.4.0

def joints_mse_loss_orig(output, target, target_weight=None):
    batch_size = output.size(0)
    num_joints = output.size(1)
    heatmaps_pred = output.view((batch_size, num_joints, -1)).split(1, 1)
    heatmaps_gt = target.view((batch_size, num_joints, -1)).split(1, 1)

    loss = 0
    for idx in range(num_joints):
        heatmap_pred = heatmaps_pred[idx]
        heatmap_gt = heatmaps_gt[idx]
        if target_weight is None:
            loss += 0.5 * mse_loss(heatmap_pred, heatmap_gt, reduction='mean')
        else:
            loss += 0.5 * mse_loss(
                heatmap_pred.mul(target_weight[:, idx]),
                heatmap_gt.mul(target_weight[:, idx]),
                reduction='mean'
            )

    return loss / num_joints


class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight=True):
        super().__init__()
        self.use_target_weight = use_target_weight
        raise NotImplementedError

    def forward(self, output, target, target_weight):
        if not self.use_target_weight:
            target_weight = None
        return joints_mse_loss_orig(output, target, target_weight)




# ----- NEW: losses when calculated on keypoint locations instead of keypoint heatmaps -----


def joints_mse_loss_onKPloc(output, target, meta, target_weight=None):
    # debugging:
    # for old kornia version
    # output_softmax_2d = dsnt.spatial_softmax_2d(target, temperature=torch.tensor(100))
    # output_kp = dsnt.spatial_softargmax_2d(output_softmax_2d, normalized_coordinates=False) + 1 
    # print(output_kp[0])
    # print(meta['tpts'][0])
    # render gaussian
    # dsnt.render_gaussian_2d(meta['tpts'][0][0, :2].to('cpu'), torch.tensor(([5., 5.])).to('cpu'), [256, 256], False)
    # output_softmax_2d = dsnt.spatial_softmax_2d(output, temperature=torch.tensor(100))
    # target_norm = target / target.sum(axis=3).sum(axis=2)[:, :, None, None]
    # output_softmax_2d = dsnt.spatial_softmax_2d(output*10)       # (target, temperature=torch.tensor(10))
    # output_kp = dsnt.spatial_softargmax_2d(target_norm, normalized_coordinates=False) + 1 

    # normalize target heatmap
    '''target_sum = target.sum(axis=3).sum(axis=2)[:, :, None, None]
    target_sum[target_sum==0] = 1e-2
    target_norm = target / target_sum'''
    target_norm = target        # now we have normalized heatmaps

    # normalize predictions -> from logits to probability distribution
    output_norm = dsnt.spatial_softmax2d(output, temperature=torch.tensor(1))

    # heatmap loss (for normalization)
    heatmap_loss = joints_mse_loss_orig(output_norm, target_norm, target_weight)

    # keypoint distance loss (average distance in pixels)
    output_kp = dsnt.spatial_expectation2d(output_norm, normalized_coordinates=False) + 1   # (bs, 20, 2)
    target_kp = meta['tpts'].to(output_kp.device)        # (bs, 20, 3)
    output_kp_resh = output_kp.reshape((-1, 2))
    target_kp_resh = target_kp[:, :, :2].reshape((-1, 2))
    weights_resh = target_kp[:, :, 2].reshape((-1))
    # dist_loss = (((output_kp_resh - target_kp_resh)**2).sum(axis=1).sqrt()*weights_resh)[weights_resh>0].sum() / min(weights_resh[weights_resh>0].sum(), 1e-5)
    dist_loss = (((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt()*weights_resh[weights_resh>0]).sum() / max(weights_resh[weights_resh>0].sum(), 1e-5)


    # return heatmap_loss*100 # + 0.0001*dist_loss 

    # import pdb; pdb.set_trace()


    '''import matplotlib as mpl
    mpl.use('Agg')
    import matplotlib.pyplot as plt

    img_np = output_norm[0, :, :, :].detach().cpu().numpy().transpose(1, 2, 0)[:, :, :3] 
    img_np = img_np * 255./ img_np.max()
    # plot image
    plt.imshow(img_np)     
    plt.savefig('./debugging_output/test_output.png')
    plt.close()

    img_np = target_norm[0, :, :, :].detach().cpu().numpy().transpose(1, 2, 0)[:, :, :3] 
    img_np = img_np * 255./ img_np.max()
    # plot image
    plt.imshow(img_np)     
    plt.savefig('./debugging_output/test_gt.png')
    plt.close()'''

    # print(heatmap_loss*100)
    # print(dist_loss * 1e-4)

    # distlossonly: return dist_loss * 1e-4 
    # both: return dist_loss * 1e-4 + heatmap_loss*100 
    return dist_loss * 1e-4 + heatmap_loss*100 




class JointsMSELoss_onKPloc(nn.Module):
    def __init__(self, use_target_weight=True):
        super().__init__()
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        if not self.use_target_weight:
            target_weight = None
        return joints_mse_loss_onKPloc(output, target, meta, target_weight)





# ----- NEW: lsegmentation loss -----

import torch.nn.functional as F

'''def resize2d(img, size):
    return (F.adaptive_avg_pool2d(Variable(img,volatile=True), size)).data
    # F.adaptive_avg_pool2d(meta['silh'], (64,64))).data'''

def segmentation_loss(output, meta):
    # output: (6, 2, 64, 64)
    # meta.keys(): ['index', 'center', 'scale', 'pts', 'tpts', 'target_weight', 'breed_index', 'silh']
    # prepare target silhouettes
    target_silh = meta['silh']
    target_silh_l = target_silh.to(torch.long)
    criterion_ce = nn.CrossEntropyLoss()
    if output.shape[2] == 64:
        target_silh_64 = F.adaptive_avg_pool2d(target_silh, (64,64))
        target_silh_64[target_silh_64>0.5] = 1
        target_silh_64[target_silh_64<=0.5] = 0
        target_silh_64_l = target_silh_64.to(torch.long)
        loss_silh_64 = criterion_ce(output, target_silh_64_l)       # 0.7
        return loss_silh_64
    else:
        loss_silh_l = criterion_ce(output, target_silh_l)       # 0.7
        return loss_silh_l