# coding=utf-8

import torch
from utils import *

def get_cmt(x, adj, s, mask=None):
  x = x.unsqueeze(0) if x.dim() == 2 else x
  adj = adj.unsqueeze(0) if adj.dim() == 2 else adj
  s = s.unsqueeze(0) if s.dim() == 2 else s

  batch_size, num_nodes, _ = x.size()

  if mask is not None:
    mask = mask.view(batch_size, num_nodes, 1).to(x.dtype)
    x, s = x * mask, s * mask

  out = torch.matmul(s.transpose(1, 2), x)
  out_adj = torch.matmul(torch.matmul(s.transpose(1, 2), adj), s)

  link_loss = adj - torch.matmul(s, s.transpose(1, 2)) + EPS
  link_loss = torch.norm(link_loss, p=2)
  link_loss = link_loss / adj.numel()

  ent_loss = (-s * torch.log(s + EPS)).sum(dim=-1).mean()

  return out, out_adj, link_loss, ent_loss