import sys
import os
import collections

import numpy as np

lib_path = 'I:/code'
if not os.path.exists(lib_path):
  lib_path = '/media/6T/.tianle/.lib'
if os.path.exists(lib_path) and lib_path not in sys.path:
  sys.path.append(lib_path)

import torch
import torch.nn as nn

from dl.models.basic_models import DenseLinear, get_list, get_attr
from dl.utils.train import cosine_similarity

class AutoEncoder(nn.Module):
  r"""Factorization autoencoder
  
  Args:
  
  Shape:
  
  Attributes:
  
  Examples::
  
  
  """
  def __init__(self, in_dim, hidden_dims, num_classes, dense=True, residual=False, residual_layers='all',
    decoder_norm=False, decoder_norm_dim=0, uniform_decoder_norm=False, nonlinearity=nn.ReLU(), 
    last_nonlinearity=True, bias=True):
    super(AutoEncoder, self).__init__()
    self.encoder = DenseLinear(in_dim, hidden_dims, nonlinearity=nonlinearity, last_nonlinearity=last_nonlinearity, 
      dense=dense, residual=residual, residual_layers=residual_layers, forward_input=False, return_all=False, 
      return_layers=None, bias=bias)
    self.decoder_norm = decoder_norm
    self.uniform_decoder_norm = uniform_decoder_norm
    if self.decoder_norm:
      self.decoder = nn.utils.weight_norm(nn.Linear(hidden_dims[-1], in_dim), 'weight', dim=decoder_norm_dim)
      if self.uniform_decoder_norm:
        self.decoder.weight_g.data = self.decoder.weight_g.new_ones(1) # This changed the tensor shape, but it's ok
        self.decoder.weight_g.requires_grad_(False)
    else:
      self.decoder = nn.Linear(hidden_dims[-1], in_dim)
    self.classifier = nn.Linear(hidden_dims[-1], num_classes)
    
  def forward(self, x):
    out = self.encoder(x)
    return self.classifier(out), self.decoder(out)


class MultiviewAE(nn.Module):
  r"""Multiview autoencoder. 

  Args:
    in_dims: a list (or iterable) of integers
    hidden_dims: a list of ints if every view has the same hidden_dims; otherwise a list of lists of ints
    out_dim: for classification, out_dim = num_cls
    fuse_type: default 'sum', add up the outputs of all encoders; require all ouputs has the same dimensions
      if 'cat', concatenate the outputs of all encoders
    dense, residual, residual_layers, nonlinearity, last_nonlinearity, bias are passed to DenseLinear
    decoder_norm: if True, add forward prehook torch.nn.utils.weight_norm  to decoder (a nn.Linear module)
    decoder_norm_dim: default 0; pass to torch.nn.utils.weight_norm
    uniform_decoder_norm: if True, ensure that decoder weight norm is 1 for dim=decoder_norm_dim

  Shape:
    Input: can be a list of tensors or a single tensor which will be splitted into a list
    Output: two heads: score matrix of shape (N, out_dim), concatenated decoder output: (N, sum(in_dims))

  Attributes:
    A list of DenseLinear modules as encoders and decoders
    An nn.Linear as output layer (e.g., class score matrix)

  Examples:
    >>> x = torch.randn(10, 5)
    >>> model = MultiviewAE([2,3], [5, 5], 7)
    >>> y = model(x)
    >>> y[0].shape, y[1].shape

  """
  def __init__(self, in_dims, hidden_dims, out_dim, fuse_type='sum', dense=False, residual=True, 
    residual_layers='all', decoder_norm=False, decoder_norm_dim=0, uniform_decoder_norm=False, 
    nonlinearity=nn.ReLU(), last_nonlinearity=True, bias=True):
    super(MultiviewAE, self).__init__()
    self.num_views = len(in_dims)
    self.in_dims = in_dims
    self.out_dim = out_dim
    self.fuse_type = fuse_type
    if not isinstance(hidden_dims[0], collections.Iterable):
      # hidden_dims is a list of ints, which means all views have the same hidden dims
      hidden_dims = [hidden_dims] * self.num_views
    self.hidden_dims = hidden_dims
    assert len(self.hidden_dims) == self.num_views and isinstance(self.hidden_dims[0], collections.Iterable)
    self.encoders = nn.ModuleList()
    self.decoders = nn.ModuleList()
    for in_dim, hidden_dim in zip(in_dims, hidden_dims):
      self.encoders.append(DenseLinear(in_dim, hidden_dim, nonlinearity=nonlinearity, 
        last_nonlinearity=last_nonlinearity, dense=dense, forward_input=False, return_all=False, 
        return_layers=None, bias=bias, residual=residual, residual_layers=residual_layers))
      decoder = nn.Linear(hidden_dim[-1], in_dim)
      if decoder_norm:
        torch.nn.utils.weight_norm(decoder, 'weight', dim=decoder_norm_dim)
        if uniform_decoder_norm:
          decoder.weight_g.data = decoder.weight_g.new_ones(decoder.weight_g.size())
          decoder.weight_g.requires_grad_(False)
      self.decoders.append(decoder)
    self.fuse_dims = [hidden_dim[-1] for hidden_dim in self.hidden_dims]
    if self.fuse_type == 'sum':
      fuse_dim = self.fuse_dims[0]
      for d in self.fuse_dims:
        assert d == fuse_dim
    elif self.fuse_type == 'cat':
      fuse_dim = sum(self.fuse_dims)
    else:
      raise ValueError(f"fuse_type should be 'sum' or 'cat', but is {fuse_type}")
    self.output = nn.Linear(fuse_dim, out_dim)

  def forward(self, xs):
    if isinstance(xs, torch.Tensor):
      xs = xs.split(self.in_dims, dim=1)
    # assert len(xs) == self.num_views
    encoder_out = []
    decoder_out = []
    for i, x in enumerate(xs):
      out = self.encoders[i](x)
      encoder_out.append(out)
      decoder_out.append(self.decoders[i](out))
    if self.fuse_type == 'sum':
      out = torch.stack(encoder_out, dim=-1).mean(dim=-1)
    else:
      out = torch.cat(encoder_out, dim=-1)
    out = self.output(out)
    return out, torch.cat(decoder_out, dim=-1), torch.cat(encoder_out, dim=-1)


def get_interaction_loss(interaction_mat, w, loss_type='graph_laplacian', normalize=True):
  """Calculate loss on the inconsistency between feature representations w (N*D) 
  and feature interaction network interaction_mat (N*N)
  A trivial solution is all features (row vectors of w) have cosine similarity = 1 or distance = 0
  
  Args:
    interaction_mat: non-negative symmetric torch.Tensor with shape (N, N)
    w: feature representation tensor with shape (N, D)
    normalize: if True, call w = w / w.norm(p=2, dim=1, keepdim=True) /np.sqrt(w.size(0)) 
      for loss_type = 'graph_laplacian' or 'dot_product',
        this makes sure w.norm() = 1 and the row vectors of w have the same norm: len(torch.unique(w.norm(dim=1)))==1
      call loss = loss / w.size(0) for loss_type = 'cosine_similarity'; 
      By doing this we ensure the number of features is factored out; 
      this is useful for combining losses from multi-views.

  See Loss_feature_interaction for more documentation

  """
  if loss_type == 'cosine_similarity':
    # -(|cos(w,w)| * interaction_mat).sum()
    cos = cosine_similarity(w).abs() # get the absolute value of cosine simiarity
    loss = -(cos * interaction_mat).sum()
    if normalize:
      loss = loss / w.size(0)
  elif loss_type == 'graph_laplacian':
    # trace(w' * L * w)
    if normalize:
      w = w / w.norm(p=2, dim=1, keepdim=True) / np.sqrt(w.size(0))
      interaction_mat = interaction_mat / interaction_mat.norm() # this will ensure interaction_mat is normalized
    diag = torch.diag(interaction_mat.sum(dim=1))
    L_interaction_mat = diag - interaction_mat
    loss = torch.diagonal(torch.mm(torch.mm(w.t(), L_interaction_mat), w)).sum()
  elif loss_type == 'dot_product':
    # pairwise distance mat * interaction mat
    if normalize:
      w = w / w.norm(p=2, dim=1, keepdim=True) / np.sqrt(w.size(0))
    d = torch.sum(w*w, dim=1) # if normalize is True, then d is a vector of the same element 1/w.size(0)
    dist = d.unsqueeze(1) + d - 2*torch.mm(w, w.t())
    loss = (dist * interaction_mat).sum()
    # loss = (dist / dist.norm() * interaction_mat).sum() # This is an alternative to 'normalize' loss
  else:
    raise ValueError(f"loss_type can only be 'cosine_similarity', "
                     f"graph_laplacian' or 'dot_product', but is {loss_type}")
  return loss


class Loss_feature_interaction(nn.Module):
  r"""A customized loss function for a graph Laplacian constraint on the feature interaction network
    For factorization autoencoder model, the decoder weights can be seen as feature representations;
    This loss measures the inconsistency between learned feature representations and their interaction network.
    A trivial solution is all features have cosine similarity = 1 or distance = 0

  Args:
    interaction_mat: torch.Tensor of shape (N, N), a non-negative (symmetric) matrix; 
      or a list of matrices; each is an interaction mat; 
      To control the magnitude of the loss, it is preferred to have argument interaction_mat.norm() = 1
    loss_type: if loss_type == 'cosine_similarity', calculate -(cos(m, m).abs() * interaction_mat).sum()
               if loss_type == 'graph_laplacian' (faster), calculate trace(m' * L * m)
               if loss_type == 'dot_product', calculate dist(m) * interaction_mat 
                 where dist(m) is the pairwise distance matrix of features; the name 'dot_product' is misleading
              If all features have norm 1, all three types are equivalent in a sense
              cosine_similarity is preferred because the magnitude of features are implicitly ignored, 
               while the other two will be affected by the magnitude of features.
    weight_path: default ['decoder', 'weight'], with the goal to get w = model.decoder.weight
    normalize: pass it to get_interaction_loss; 
      if True, call w = w / w.norm(p=2, dim=1, keepdim=True) / np.sqrt(w.size(0))
        for loss_type 'graph_laplacian' or 'dot_product',
          this makes sure each row vector of w has the same norm, and w.norm() = 1
        call loss = loss / w.size(0) for loss_type = 'cosine_similarity'; 
      By doing this we ensure the number of features is factored out; 
      this is useful for combining losses from multi-views.
  
  Inputs:
    model: the above defined AutoEnoder model or other model
    or given weight matrix w
    if interaction_mat has shape (N,N), then w has shape (N, D)

  Returns:
    loss: torch.Tensor that can call loss.backward()
  """

  def __init__(self, interaction_mat, loss_type='graph_laplacian', weight_path=['decoder', 'weight'], 
    normalize=True):
    super(Loss_feature_interaction, self).__init__()
    self.loss_type = loss_type
    self.weight_path = weight_path
    self.normalize = normalize
    # If interaction_mat is a list, self.sections will be the used for splitting the weight matrix
    self.sections = None # when interaction_mat is a single matrix, self.sections is None
    if isinstance(interaction_mat, (list, tuple)):
      if normalize: # ensure interaction_mat is normalized
        interaction_mat = [m/m.norm() for m in interaction_mat]
      self.sections = [m.shape[0] for m in interaction_mat]
    else:
      if normalize: # ensure interaction_mat is normalized
        interaction_mat = interaction_mat / interaction_mat.norm()
    if self.loss_type == 'graph_laplacian':
      # precalculate self.L_interaction_mat save some compute for each forward pass
      if self.sections is None:
        diag = torch.diag(interaction_mat.sum(dim=1))
        self.L_interaction_mat = diag - interaction_mat # Graph Laplacian; should I normalize it?
      else:
        self.L_interaction_mat = []
        for mat in interaction_mat:
          diag = torch.diag(mat.sum(dim=1))
          self.L_interaction_mat.append(diag - mat)
    else: # we don't need to store interaction_mat for loss_type=='graph_laplacian'
      self.interaction_mat = interaction_mat
  
  def forward(self, model=None, w=None):
    if w is None:
      w = get_attr(model, self.weight_path)
    if self.sections is None:
      # There is only one interaction matrix; self.interaction_mat is a torch.Tensor
      if self.loss_type == 'graph_laplacian':
        # Used precalculated L_interaction_mat to save some time
        if self.normalize:
          # interaction_mat had already been normalized during initialization
          w = w / w.norm(p=2, dim=1, keepdim=True) / np.sqrt(w.size(0))
        return torch.diagonal(torch.mm(torch.mm(w.t(), self.L_interaction_mat), w)).sum()
      else:
        return get_interaction_loss(self.interaction_mat, w, loss_type=self.loss_type, normalize=self.normalize)
    else:
      # self.interaction_mat is a list of torch.Tensors
      if isinstance(w, torch.Tensor):
        w = w.split(self.sections, dim=0)
      if self.loss_type == 'graph_laplacian': # handle 'graph_laplacian' differently to save time during training
        loss = 0
        for w_, L in zip(w, self.L_interaction_mat):
          if self.normalize: # make sure w_.norm() = 1 and each row vector of w_ has the same norm
            w_ = w_ / w_.norm(p=2, dim=1, keepdim=True) / np.sqrt(w_.size(0))
          loss += torch.diagonal(torch.mm(torch.mm(w_.t(), L), w_)).sum()  
        return loss
      # for the case 'cosine_similarity' and 'dot_product'
      return sum([get_interaction_loss(mat, w_, loss_type=self.loss_type, normalize=self.normalize) 
                  for mat, w_ in zip(self.interaction_mat, w)])


class Loss_view_similarity(nn.Module):
  r"""The input is a multi-view representation of the same set of patients, 
      i.e., a set of matrices with shape (num_samples, feature_dim). feature_dim can be different for each view
    This loss will penalize the inconsistency among different views.
    This is somewhat limited, because different views should have both shared and complementary information
      This loss only encourages the shared information across views, 
      which may or may not be good for certain applications.
    A trivial solution for this is multi-view representation are all the same; then loss -> -1
    The two loss_types 'circle' and 'hub' can be quite different and unstable.
      'circle' tries to make all feature representations across views have high cosine similarity,
      while 'hub' only tries to make feature representations within each view have high cosine similarity;
      by multiplying 'mean-feature' target with 'hub' loss_type, it might 'magically' capture both within-view and 
        cross-view similarity; set as default choice; but my limited experimental results do not validate this;
        instead, 'circle' and 'hub' are dominant, while explicit_target and cal_target do not make a big difference 
    Cosine similarity are used here; To do: other similarity metrics

  Args:
    sections: a list of integers (or an int); this is used to split the input matrix into chunks;
      each chunk corresponds to one view representation.
      If input xs is not a torch.Tensor, this will not be used; assume xs to be a list of torch.Tensors
      sections being an int implies all feature dim are the same, set sections = feature_dim, NOT num_sections!
    loss_type: supose there are three views x1, x2, x3; let s_ij = cos(x_i,x_j), s_i = cos(x_i,x_i)
      if loss_type=='cicle', similarity = s12*s23*target if fusion_type=='multiply'; s12+s23 if fusion_type=='sum'                   
        This is fastest but requires x1, x2, x3 have the same shape
      if loss_type=='hub', similarity=s1*s2*s3*target if fusion_type=='multiply'; 
        similarity=|s1|+|s2|+|s3|+|target| if fusion_type=='sum'
        Implicitly, target=1 (fusion_type=='multiply) or 0 (fusion_type=='sum') if explicit_target is False
        if graph_laplacian is False:
          loss = - similarity.abs().mean()
        else:
          s = similarity.abs(); L_s = torch.diag(sum(s, axis=1)) - s #graph laplacian
          loss = sum_i(x_i * L_s * x_i^T)
    explicit_target: if False, target=1 (fusion_type=='multiply) or 0 (fusion_type=='sum') implicitly
      if True, use given target or calculate it from xs
      # to do handle the case when we only use the explicitly given target
    cal_target: if 'mean-similarity', target = (cos(x1,x1) + cos(x2,x2) + cos(x3,x3))/3
                if 'mean-feature', x = (x1+x2+x3)/3; target = cos(x,x); this requires x1,x2,x3 have the same shape
    target: default None; only used when explicit_target is True
      This saves computation if target is provided in advance or passed as input
    fusion_type: if 'multiply', similarity=product(similarities); if 'sum', similarity=sum(|similarities|);
      work with loss_type
    graph_laplacian:  if graph_laplacian is False:
          loss = - similarity.abs().mean()
        else:
          s = similarity.abs(); L_s = torch.diag(sum(s, axis=1)) - s #graph laplacian
          loss = sum_i(x_i * L_s * x_i^T)

  Inputs:
    xs: a set of torch.Tensor matrices of (num_samples, feature_dim), 
      or a single matrix with self.sections being specified
    target: the target cosine similarity matrix; default None; 
      if not given, first check if self.targets is given; 
        if self.targets is None, then calulate it according to cal_target;
      only used when self.explicit_target is True

  Output:
    loss = -similarity.abs().mean() if graph_laplacian is False # Is this the right way to do it?
      = sum_i(x_i * L_s * x_i^T) if graph_laplacian is True # call get_interaction_loss()
    
  """
  def __init__(self, sections=None, loss_type='hub', explicit_target=False, 
    cal_target='mean-feature', target=None, fusion_type='multiply', graph_laplacian=False):
    super(Loss_view_similarity, self).__init__()
    self.sections = sections
    if self.sections is not None:
      if not isinstance(self.sections, int):
        assert len(self.sections) >= 2  
    self.loss_type = loss_type
    assert self.loss_type in ['circle', 'hub']
    self.explicit_target = explicit_target
    self.cal_target = cal_target
    self.target = target
    self.fusion_type = fusion_type
    self.graph_laplacian = graph_laplacian
    # I got nan losses easily for whenever graph_laplacian is True, especially the following case; did not know why
    # probably I need normalize similarity during every forward?
    assert not (fusion_type=='multiply' and graph_laplacian) and not (loss_type=='circle' and graph_laplacian)

  def forward(self, xs, target=None):
    if isinstance(xs, torch.Tensor):
      # make sure xs is a list of tensors corresponding to multiple views
      # this requires self.sections to valid
      xs = xs.split(self.sections, dim=1) 
    # assert len(xs) >= 2 # comment this to save time for many forward passes
    similarity = 1
    if self.loss_type == 'circle':
      # assert xs[i-1].shape == xs[i].shape
      # this saves computation
      similarity_mats = [cosine_similarity(xs[i-1], xs[i]) for i in range(1, len(xs))]
      similarity_mats = [(m+m.t())/2 for m in similarity_mats] # make it symmetric
    elif self.loss_type == 'hub':
      similarity_mats = [cosine_similarity(x) for x in xs]
    if self.fusion_type=='multiply':
      for m in similarity_mats:
        similarity = similarity * m # element multiplication ensures the larget value to be 1
    elif self.fusion_type=='sum':
      similarity = sum(similarity_mats) / len(similarity_mats) # calculate mean to ensure the largest value to be 1

    if self.explicit_target:
      if target is None:
        if self.target is None:
          if self.cal_target == 'mean-similarity':
            target = torch.stack(similarity_mats, dim=0).mean(0)
          elif self.cal_target == 'mean-feature':
            x = torch.stack(xs, -1).mean(-1) # the list of view matrices must have the same dimension
            target = cosine_similarity(x)
          else:
            raise ValueError(f'cal_target should be mean-similarity or mean-feature, but is {self.cal_target}')
        else:
          target = self.target
      if self.fusion_type=='multiply':
        similarity = similarity * target
      elif self.fusion_type=='sum':
        similarity = (len(similarity_mats)*similarity + target) / (len(similarity_mats) + 1) # Moving average
    similarity = similarity.abs() # ensure similarity to be non-negative
    if self.graph_laplacian:
      # Easily get nan loss when it is True; do not know why
      return sum([get_interaction_loss(similarity, w, loss_type='graph_laplacian', normalize=True) for w in xs]) / len(xs)
    else:
      return -similarity.mean() # to ensure the loss is within range [-1, 0]