import torch
import torch.nn as nn 
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.parameter import Parameter
import torch.optim as optim
import math
import time
import numpy as np


class GraphConvolution(nn.Module):
    """
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    """
    
    """GraphConvolution.

        Parameters
        ----------
        in_features : int 
            input dimension.
        out_features : int 
            output dimension.
        support : 
            support for graph convolution.    
        bias : boolean optional
            include bias term, default False.

        Returns
        -------
        None. 

        """
    
    def __init__(self, in_features, out_features, support, bias=False):
        super(GraphConvolution, self).__init__()
        self.support = support
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        #glorot
        init_range = np.sqrt(6.0 / (self.in_features + self.out_features))
        self.weight.data.uniform_(-init_range, init_range)
        # similar to kaiming normal/uniform
        stdv = 1. / math.sqrt(self.weight.size(1))
        #self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input, adj):
        """forward function.

        Parameters
        ----------
        input : 
            node features.
        adj : 
            adjacency matrix.

        Returns
        -------
        output : float
            output of graph convolution layer.

        """
        #convolution
        if input.is_sparse:
            #sparse input features
            support = torch.spmm(input, self.weight)
        else:
            support = torch.mm(input, self.weight)
        
        #adj should always be sparse
        #### add a ReLU or other activation!!
        output = torch.spmm(adj, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -> ' \
               + str(self.out_features) + ')'
    
    
class GCN(nn.Module):
    """
    dropout + GC + activation
    """
    def __init__(self, nfeat, nhid1, nout, bias=False, dropout=0., act=F.relu):
        super(GCN, self).__init__()
            
        self.gc1 = GraphConvolution(nfeat, nhid1, bias)
        self.gc2 = GraphConvolution(nhid1, nout, bias)
        self.dropout = dropout
        self.act = act
        

    def forward(self, x, adj):
        """forward function.

        Parameters
        ----------
        x : 
            node features.
        adj : 
            adjacency matrix.

        Returns
        -------
        output : float
            output of graph convolution network.

        """
        #self.nnz = x._nnz() if sparse_inputs else None
        
        #dropout + convolution
        x=dropout_layer(x, self.dropout)
        x = self.gc1(x, adj)
        x = self.act(x)
        x=dropout_layer(x, self.dropout)
        x=self.gc2(x, adj)
        
        return x
    

class DSTGLearner(object):
    """DSTGLearner.

        Parameters
        ----------
        nfeat : int 
            input dimension.
        nhid1 : int
            number of units in the hidden layer (graph convolution).
        nout : int 
            output dimension.
        bias : boolean optional
            include bias term, default False.
        dropout : float optional
            dropout rate, default 0.
        act : optional
            activation function, default torch functional relu.

        Returns
        -------
        None. 

        """
    def __init__(self, nfeat, nhid1, nout, bias=False, dropout=0., act=F.relu, device = torch.device("cuda" if torch.cuda.is_available() else "cpu")):
        super(DSTGLearner, self).__init__()
        self.device = device
        self.model = GCN(nfeat, nhid1, nout, bias, dropout, act).to(device)
            
    def fit(self, X, adj, labels, labels_mask=None, lr=0.005, max_epochs=100, weight_decay=0):
        """fit function for model training.

        Parameters
        ----------
        X : 
            node features.
        adj : 
            adjacency matrix.
        labels :
            labels of pseudo spots (and real spots if provided). 
        labels_mask : optional
            mask to indicate which samples to use for training.
        lr : float optional
            learning rate.
        max_epochs : int optional
            maximum number of epochs to train.
        weight_decay : float optional
            weight decay parameter for optimization (Adam). 
            
        Returns
        -------
        None.

        """
        X=X.to(self.device)
        adj=adj.to(self.device)
        labels=labels.to(self.device)
        labels_mask=labels_mask.to(self.device)
        
        #device = self.device
        model = self.model
        model.train()
        optimizer = optim.Adam(model.parameters(),lr=lr, weight_decay=weight_decay)
        # optimizer = RAdam(model.parameters(), lr=self.args.lr)
        
        #X=torch.FloatTensor(X)
        #adj=torch.FloatTensor(adj)
        
        for epoch in range(max_epochs):
            t = time.time()
            
            y_hat = model(X, adj)
            loss = masked_softmax_cross_entropy(y_hat, labels, labels_mask)
            
            if (epoch+1) % 5 == 0:
                print("Epoch:", '%04d' % (epoch + 1), "train_loss=",
                  "{:.5f}".format(loss),
                  "time=", "{:.5f}".format(time.time() - t))
            #    _ = model(X, adj)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
    def predict(self, X, adj):
        """prediction function.
        Parameters
        ----------
        X : 
            node features.
        adj : 
            adjacency matrix.
     
        Returns
        -------
        
        pred : torch tensor
            predictions of cell-type proportions.
        
        """
        X=X.to(self.device)
        adj=adj.to(self.device)
        fX=self.model(X,adj)
        pred = F.softmax(fX, dim=1)
        return pred
    
    
    def score(self, X, adj, labels, labels_mask=None):
        """score.

        Parameters
        ----------
        X : 
            node features.
        adj : 
            adjacency matrix.
        labels :
            pseudo spot labels (and real labels if provided).
        labels_mask : optional
            mask to indicate which samples to use in scoring.

        Returns
        -------
        loss : float
            cross entropy loss between predicted and true labels.

        """
        X=X.to(self.device)
        adj=adj.to(self.device)
        labels=labels.to(self.device)
        labels_mask=labels_mask.to(self.device)
        model = self.model
        model.eval()
        y_hat = model(X, adj)
        pred = F.softmax(y_hat, dim=1)
        loss = masked_softmax_cross_entropy(y_hat, labels, labels_mask)
        return loss

    
def dropout_layer(x, dropout):
    """dropout_layer.
    Parameters
    ----------
    x:  
        input to dropout layer.
    dropout : float
        dropout rate (between 0 and 1).

    Returns
    -------
    out : torch tensor
        dropout output.
    """
    if x.is_sparse:
        #sparse input features
        out = sparse_dropout(x, dropout)
        return out
    else:
        out = F.dropout(x, dropout)
        return out
        
        
def sparse_dropout(x, dropout):
    """sparse_dropout.
    Parameters
    ----------
    x:  
        input to dropout layer.
    dropout : float
        dropout rate (between 0 and 1).

    Returns
    -------
    out * (1. / (1-dropout)) : torch tensor
        dropout output.
    """
                        
    noise_shape = x._nnz()
    random_tensor = (1-dropout)+torch.rand(noise_shape).to(x.device)
    dropout_mask = torch.floor(random_tensor).type(torch.bool)
    i = x._indices()
    v = x._values()

    i = i[:, dropout_mask]
    v = v[dropout_mask]

    out = torch.sparse.FloatTensor(i, v, x.shape).to(x.device)
    return out * (1. / (1-dropout))


"""Softmax cross-entropy loss with masking."""
def masked_softmax_cross_entropy(preds, labels, mask):
    """masked_softmax_cross_entropy.
    Parameters
    ----------
    preds:  
        cell-type proportion predictions from dstg model.
    labels : 
        true cell-type proportion labels.
    mask : 
        mask to indicate which samples to use in computing loss.

    Returns
    -------
    loss : float
        cross entropy loss between true and predicted cell-type proportions (mean reduced).
    """
    if(mask is None):
        loss = F.cross_entropy(preds, labels, reduction='mean')
        return loss
    else: 
        loss = F.cross_entropy(preds, labels, reduction='none')
        mask = mask.type(torch.float32)
        mask /= torch.mean(mask)
        loss *= mask
        loss = torch.mean(loss)
        return loss