import torch
import numpy as np


def gaussian_kernel(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    n_samples = source.shape[0] + target.shape[0]
    total = torch.cat([source, target], dim=0)
    total0 = total.unsqueeze(0).expand(total.shape[0], total.shape[0], total.shape[1])
    total1 = total.unsqueeze(1).expand(total.shape[0], total.shape[0], total.shape[1])
    l2_distance = ((total0-total1)**2).sum(2)
    
    if fix_sigma:
        bandwidth = fix_sigma
    else:
        bandwidth = torch.sum(l2_distance.detach()) / (n_samples**2-n_samples)
        
    bandwidth /= kernel_mul ** (kernel_num // 2)
    bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]
    kernel_val = [torch.exp(-l2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
    return sum(kernel_val)#/len(kernel_val)

# def guassian_kernel(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
#     n_samples = int(source.size()[0])+int(target.size()[0])
#     total = torch.cat([source, target], dim=0)
#     print(total)
#     # 接下来就是扩维了，因为是计算每一个点之间的距离。
#     # 这里十分巧妙，第一个是按原来的数据进行增维，就是说复制这么多个点，有多少个点为一组进行复制
#     # 第二部分在第一维度进行扩张，把每一个点复制一共有多少个点的次数，非常方便！ 每一个点复制一组
#     # 那么再计算的时候直接相减，就得到每一个点和其他点的距离了。
#     total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
#     total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
#     print(total0)
#     print('mean:', torch.mean(total0))
#     L2_distance = ((total0-total1)**2).sum(2)  # 按特征维进行求和
#     if fix_sigma:
#         bandwidth = fix_sigma
#     else:
#         # torch.sum 直接是所有元素相加到一个元素
#         bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)  # 这里就是计算K(1/(m(m-1)))
#     # 以 fix_sigma为中值，以kernel_mul为倍数取kernel_num个值
#     bandwidth /= kernel_mul ** (kernel_num // 2)
#     print(bandwidth)
#     bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]
#     print(bandwidth_list)
#     print(L2_distance)
#     kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
#     return sum(kernel_val)  #返回度量矩阵。

def prob_map(sp, tp):

    total = torch.cat([sp, tp], dim=0).unsqueeze(1)
    total1 = torch.cat([sp, tp], dim=0).unsqueeze(0)
    prob = torch.mm(total, total1)
    return prob

def origin_mmd_soft(source, target, sp, tp, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    n = source.shape[0]
    used = (torch.ones((n,n)) - torch.eye(n)).cuda()
    kernels = gaussian_kernel(source, target, kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)
    #print(kernels)
    prob = prob_map(sp, tp)
    kernels_p = kernels * prob
    XX = kernels_p[:n, :n] * used
    YY = kernels_p[n:, n:] * used
    XY = kernels_p[:n, n:] * used

    ssp = sum(sp) * sum(sp)
    ttp = sum(tp) * sum(tp)
    stp = sum(sp) * sum(tp)
    XX = torch.sum(XX) * n/(n - 1) / ssp
    YY = torch.sum(YY) * n/(n - 1) / ttp
    XY = torch.sum(XY) * n/(n - 1) / stp

    loss = torch.abs(XX + YY - 2*XY)
    return loss

def mmd_soft(source, target, sp, tp, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    n = len(source)
    m = len(target)
    #used = (torch.ones((n,n)) - torch.eye(n)).cuda()
    kernels = gaussian_kernel(source, target, kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)
    #print(kernels)
    prob = prob_map(sp, tp)
    kernels_p = kernels * prob
    XX = kernels_p[:n, :n]
    YY = kernels_p[-m:, -m:]
    XY = kernels_p[:n, -m:]

    ssp = sum(sp) * sum(sp)
    ttp = sum(tp) * sum(tp)
    stp = sum(sp) * sum(tp)
    XX = torch.sum(XX) * n/(n - 1) / ssp
    YY = torch.sum(YY) * n/(n - 1) / ttp
    XY = torch.sum(XY) * n/(n - 1) / stp

    loss = torch.abs(XX + YY - 2*XY)
    return loss


def mmd_cal(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):

    n = source.shape[0]
    m = target.shape[0]
    kernels = gaussian_kernel(source, target,
                              kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)
    XX = kernels[:n, :n]
    YY = kernels[-m:, -m:]
    XY = kernels[:n, -m:]
    YX = kernels[-m:, :n]

    XX = torch.div(XX, n * n).sum(dim=1).view(1, -1)  # K_ss矩阵，Source<->Source
    XY = torch.div(XY, -n * m).sum(dim=1).view(1, -1)  # K_st矩阵，Source<->Target

    YX = torch.div(YX, -m * n).sum(dim=1).view(1, -1)  # K_ts矩阵,Target<->Source
    YY = torch.div(YY, m * m).sum(dim=1).view(1, -1)  # K_tt矩阵,Target<->Target

    loss = (XX + XY).sum() + (YX + YY).sum()
    return loss

if __name__ == "__main__":
    # 样本数量可以不同，特征数目必须相同

    # 100和90是样本数量，50是特征数目
    data_1 = torch.tensor([[2], [1]])
    data_2 = torch.tensor([[3], [2]])
    sp = torch.tensor((0.2, 0.9))
    tp = torch.tensor((0.3, 0.6))
    print("MMD Loss:",mmd_soft(data_1,data_2, sp, tp))

    data_1 = torch.tensor(np.random.normal(loc=0,scale=10,size=(100,50)))
    data_2 = torch.tensor(np.random.normal(loc=0,scale=9,size=(80,50)))

    print("MMD Loss:",mmd_cal(data_1,data_2))

