import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger

import wandb

from torch import nn
import torch
import torch.nn.functional as F
from torch.nn import TransformerEncoderLayer

import numpy as np
from scipy.stats import rankdata

from allennlp.modules.attention import CosineAttention, BilinearAttention, AdditiveAttention, DotProductAttention


from utils.loss_utils import NCALoss
from utils.train_utils import mean_reciprocal_rank, top_k_acc, masked_mean, masked_softmax, weighted_sum


class PatientNCA(pl.LightningModule):

    def __init__(self, hparams, embed_dim):
        super().__init__()
        self.hyperparameters = hparams

        # attention for collapsing set of phenotype embeddings
        self.attn_vector = nn.Parameter(torch.zeros((1, embed_dim), dtype=torch.float), requires_grad=True)   
        nn.init.xavier_uniform_(self.attn_vector)
        
        if self.hyperparameters['attention_type'] == 'bilinear':
            self.attention = BilinearAttention(embed_dim, embed_dim)
        elif self.hyperparameters['attention_type'] == 'additive':
            self.attention = AdditiveAttention(embed_dim, embed_dim)
        elif self.hyperparameters['attention_type'] == 'dotpdt':
            self.attention = DotProductAttention()

        # projection layers
        self.phen_project = nn.Linear(embed_dim, embed_dim)
        self.phen_project2 = nn.Linear(embed_dim, embed_dim)
        if self.hyperparameters['loss'] == 'patient_disease_NCA':
            self.disease_project = nn.Linear(embed_dim, embed_dim)
            self.disease_project2 = nn.Linear(embed_dim, embed_dim)

        self.leaky_relu = nn.LeakyReLU(hparams['leaky_relu'])

        self.loss = NCALoss(softmax_scale=self.hyperparameters['softmax_scale'], only_hard_distractors=self.hyperparameters['only_hard_distractors']) 
        
        if 'n_transformer_layers' in hparams and hparams['n_transformer_layers'] > 0:
            encoder_layer = TransformerEncoderLayer(d_model=embed_dim, nhead=hparams['n_transformer_heads'])
            self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=hparams['n_transformer_layers'])


    def forward(self, phenotype_embeddings, disease_embeddings, phenotype_mask=None, disease_mask=None): 
        assert phenotype_mask != None  
        if self.hyperparameters['loss'] == 'patient_disease_NCA':  assert disease_mask != None
        
        if 'n_transformer_layers' in self.hyperparameters and self.hyperparameters['n_transformer_layers'] > 0:
            phenotype_embeddings = self.transformer_encoder(phenotype_embeddings.transpose(0, 1), src_key_padding_mask=~phenotype_mask).transpose(0, 1)


        # attention weighted average of phenotype embeddings
        batched_attn = self.attn_vector.repeat(phenotype_embeddings.shape[0],1)
        attn_weights = self.attention(batched_attn, phenotype_embeddings, phenotype_mask)
        phenotype_embedding = weighted_sum(phenotype_embeddings, attn_weights)
        
        # project embeddings
        phenotype_embedding = self.phen_project2(self.leaky_relu(self.phen_project(phenotype_embedding)))
        if self.hyperparameters['loss'] == 'patient_disease_NCA': disease_embeddings = self.disease_project2(self.leaky_relu(self.disease_project(disease_embeddings)))
        else: disease_embeddings = None

        return phenotype_embedding, disease_embeddings, phenotype_mask, disease_mask, attn_weights

    def calc_loss(self, batch, phenotype_embedding, disease_embeddings, disease_mask, labels, use_candidate_list):
        if self.hyperparameters['loss'] == 'patient_disease_NCA':
            loss, softmax, labels, candidate_disease_idx, candidate_disease_embeddings = self.loss(phenotype_embedding, disease_embeddings, batch.batch_disease_nid, batch.batch_cand_disease_nid, disease_mask, labels, use_candidate_list=use_candidate_list)
        elif self.hyperparameters['loss'] == 'patient_patient_NCA':
            loss, softmax, labels, candidate_disease_idx, candidate_disease_embeddings = self.loss(phenotype_embedding, None, None, None, None, labels, use_candidate_list=False)
        else:
            raise NotImplementedError
        return loss, softmax, labels, candidate_disease_idx, candidate_disease_embeddings


    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hyperparameters['lr'])
        return optimizer