import torch

import torch.nn as nn

import torch.nn.functional as F

import numpy as np

import math

class AttnNet(nn.Module):

    # Adapted from https://github.com/mahmoodlab/CLAM/blob/master/models/model_clam.py

    # Lu, M.Y., Williamson, D.F.K., Chen, T.Y. et al. Data-efficient and weakly supervised computational pathology on whole-slide images. Nat Biomed Eng 5, 555–570 (2021). https://doi.org/10.1038/s41551-020-00682-w

    def __init__(self, L=1024, D=256, dropout=False, p_dropout_atn=0.25, n_classes=1):

        super(AttnNet, self).__init__()

        self.attention_a = [nn.Linear(L, D), nn.Tanh()]

        self.attention_b = [nn.Linear(L, D), nn.Sigmoid()]

        if dropout:

            self.attention_a.append(nn.Dropout(p_dropout_atn))

            self.attention_b.append(nn.Dropout(p_dropout_atn))

        self.attention_a = nn.Sequential(*self.attention_a)

        self.attention_b = nn.Sequential(*self.attention_b)

        self.attention_c = nn.Linear(D, n_classes)

    def forward(self, x):

        a = self.attention_a(x)

        b = self.attention_b(x)

        A = a.mul(b)

        A = self.attention_c(A)  # N x n_classes

        return A

class Attn_Modality_Gated(nn.Module):

    # Adapted from https://github.com/mahmoodlab/PORPOISE

    def __init__(self, gate_h1, gate_h2, gate_h3, dim1_og, dim2_og, dim3_og, use_bilinear=[True,True,True], scale=[1,1,1], p_dropout_fc=0.25):

        super(Attn_Modality_Gated, self).__init__()

        self.gate_h1 = gate_h1 #[boolean]

        self.gate_h2 = gate_h2 #[boolean]

        self.gate_h3 = gate_h3 #[boolean]

        self.use_bilinear = use_bilinear #[boolean]

        # can perform attention on latent vectors of lower dimension

        dim1, dim2, dim3 = dim1_og//scale[0], dim2_og//scale[1], dim3_og//scale[2]

        # attention gate of each modality

        if self.gate_h1:

            self.linear_h1 = nn.Sequential(nn.Linear(dim1_og, dim1), nn.ReLU())

            self.linear_z1 = nn.Bilinear(dim1_og, dim2_og+dim3_og, dim1) if self.use_bilinear[0] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim1))

            self.linear_o1 = nn.Sequential(nn.Linear(dim1, dim1), nn.ReLU(), nn.Dropout(p=p_dropout_fc))

        else:

            self.linear_h1, self.linear_o1 = nn.Identity(), nn.Identity()  

        if self.gate_h2:

            self.linear_h2 = nn.Sequential(nn.Linear(dim2_og, dim2), nn.ReLU())

            self.linear_z2 = nn.Bilinear(dim2_og, dim1_og+dim3_og, dim2) if self.use_bilinear[1] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim2))

            self.linear_o2 = nn.Sequential(nn.Linear(dim2, dim2), nn.ReLU(), nn.Dropout(p=p_dropout_fc))

        else:

            self.linear_h2, self.linear_o2 = nn.Identity(), nn.Identity()  

        if self.gate_h3:

            self.linear_h3 = nn.Sequential(nn.Linear(dim3_og, dim3), nn.ReLU())

            self.linear_z3 = nn.Bilinear(dim3_og, dim1_og+dim2_og, dim3) if self.use_bilinear[2] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim3))

            self.linear_o3 = nn.Sequential(nn.Linear(dim3, dim3), nn.ReLU(), nn.Dropout(p=p_dropout_fc))

        else:

            self.linear_h3, self.linear_o3 = nn.Identity(), nn.Identity()  

    def forward(self, x1, x2, x3):

        if self.gate_h1:

            h1 = self.linear_h1(x1) #breaks colli of h1

            z1 = self.linear_z1(x1, torch.cat([x2,x3], dim=-1)) if self.use_bilinear[0] else self.linear_z1(torch.cat((x1, x2, x3), dim=-1)) #creates a vector combining both modalities

            o1 = self.linear_o1(nn.Sigmoid()(z1)*h1) #update modality input

        else:

            h1 = self.linear_h1(x1)

            o1 = self.linear_o1(h1)

        if self.gate_h2:

            h2 = self.linear_h2(x2)

            z2 = self.linear_z2(x2, torch.cat([x1,x3], dim=-1)) if self.use_bilinear[1] else self.linear_z2(torch.cat((x1, x2, x3), dim=-1))

            o2 = self.linear_o2(nn.Sigmoid()(z2)*h2)

        else:

            h2 = self.linear_h2(x2)

            o2 = self.linear_o2(h2)

        if self.gate_h3:

            h3 = self.linear_h3(x3)

            z3 = self.linear_z3(x3, torch.cat([x1,x2], dim=-1)) if self.use_bilinear[2] else self.linear_z3(torch.cat((x1, x2, x3), dim=-1))

            o3 = self.linear_o3(nn.Sigmoid()(z3)*h3)

        else:

            h3 = self.linear_h3(x3)

            o3 = self.linear_o3(h3)

        return o1, o2, o3

class FC_block(nn.Module):

    def __init__(self, dim_in, dim_out, act_layer=nn.ReLU, dropout=True, p_dropout_fc=0.25):

        super(FC_block, self).__init__()

        self.fc = nn.Linear(dim_in, dim_out)

        self.act = act_layer()

        self.drop = nn.Dropout(p_dropout_fc) if dropout else nn.Identity()

    def forward(self, x):

        x = self.fc(x)

        x = self.act(x)

        x = self.drop(x)

        return x

class Categorical_encoding(nn.Module):

    def __init__(self, taxonomy_in=3, embedding_dim=128, depth=1, act_fct='relu', dropout=True, p_dropout=0.25):

        super(Categorical_encoding, self).__init__()

        act_fcts = {'relu': nn.ReLU(),

        'elu' : nn.ELU(),

        'tanh': nn.Tanh(),

        'selu': nn.SELU(),

        }

        dropout_module = nn.AlphaDropout(p_dropout) if act_fct=='selu' else nn.Dropout(p_dropout)

        self.embedding = nn.Embedding(taxonomy_in, embedding_dim)

        fc_layers = []

        for d in range(depth):

            fc_layers.append(nn.Linear(embedding_dim//(2**d), embedding_dim//(2**(d+1))))

            fc_layers.append(dropout_module if dropout else nn.Identity())

            fc_layers.append(act_fcts[act_fct])

        self.fc_layers = nn.Sequential(*fc_layers)

    def forward(self, x):

        x = self.embedding(x)

        x = self.fc_layers(x)

        return x

class HECTOR(nn.Module):

    def __init__(

        self,

        input_feature_size=1024,

        precompression_layer=True,

        feature_size_comp = 512,

        feature_size_attn = 256,

        postcompression_layer=True,

        feature_size_comp_post = 128,

        dropout=True,

        p_dropout_fc=0.25,

        p_dropout_atn=0.25,

        n_classes=2,

        input_stage_size=6,

        embedding_dim_stage=16,

        depth_dim_stage=1,

        act_fct_stage='elu',

        dropout_stage=True,

        p_dropout_stage=0.25,

        input_mol_size=4,

        embedding_dim_mol=16,

        depth_dim_mol=1,

        act_fct_mol='elu',

        dropout_mol=True,

        p_dropout_mol=0.25,

        fusion_type='kron',

        use_bilinear=[True,True,True],

        gate_hist=False,

        gate_stage=False,

        gate_mol=False,

        scale=[1,1,1],

    ):

        super(HECTOR, self).__init__()

        self.fusion_type =fusion_type

        self.input_stage_size=input_stage_size

        self.use_bilinear = use_bilinear

        self.gate_hist = gate_hist

        self.gate_stage = gate_stage

        self.gate_mol = gate_mol

        # Reduce dimension of H&E patch features.

        if precompression_layer:

            self.compression_layer = nn.Sequential(*[FC_block(input_feature_size, feature_size_comp*4, p_dropout_fc=p_dropout_fc),

                                                    FC_block(feature_size_comp*4, feature_size_comp*2, p_dropout_fc=p_dropout_fc),

                                                    FC_block(feature_size_comp*2, feature_size_comp, p_dropout_fc=p_dropout_fc),])

            dim_post_compression = feature_size_comp

        else:

            self.compression_layer = nn.Identity()

            dim_post_compression = input_feature_size

        # Get embeddings of categorical features.

        self.encoding_stage_net = Categorical_encoding(taxonomy_in=self.input_stage_size, 

                                                embedding_dim=embedding_dim_stage, 

                                                depth=depth_dim_stage, 

                                                act_fct=act_fct_stage, 

                                                dropout=dropout_stage, 

                                                p_dropout=p_dropout_stage)

        self.out_stage_size = embedding_dim_stage//(2**depth_dim_stage)

        self.encoding_mol_net = Categorical_encoding(taxonomy_in=input_mol_size, 

                                                embedding_dim=embedding_dim_mol, 

                                                depth=depth_dim_mol, 

                                                act_fct=act_fct_mol, 

                                                dropout=dropout_mol, 

                                                p_dropout=p_dropout_mol)

        h_mol_size_out = embedding_dim_mol//(2**depth_dim_mol)

        # For survival tasks the attention scores are binary (set to class=1).

        self.attention_survival_net = AttnNet(

            L=dim_post_compression,

            D=feature_size_attn,

            dropout=dropout,

            p_dropout_atn=p_dropout_atn,

            n_classes=1,)

        # Attention gate on each modality.

        self.attn_modalities = Attn_Modality_Gated(

            gate_h1=self.gate_hist, 

            gate_h2=self.gate_stage, 

            gate_h3=self.gate_mol,

            dim1_og=dim_post_compression, 

            dim2_og=self.out_stage_size, 

            dim3_og=h_mol_size_out,

            scale=scale, 

            use_bilinear=self.use_bilinear)

        # Post-compression layer for H&E slide-level embedding before fusion.

        dim_post_compression = dim_post_compression//scale[0] if self.gate_hist else dim_post_compression

        self.post_compression_layer_he = FC_block(dim_post_compression, dim_post_compression//2, p_dropout_fc=p_dropout_fc)

        dim_post_compression = dim_post_compression//2

        # Post-compression layer.

        dim1, dim2, dim3 = dim_post_compression, self.out_stage_size//scale[1] if self.gate_stage else self.out_stage_size, h_mol_size_out//scale[2] if self.gate_mol else h_mol_size_out

        if self.fusion_type=='bilinear':

            head_size_in = (dim1+1)*(dim2+1)*(dim3+1)

        elif self.fusion_type=='kron':

            head_size_in = (dim1)*(dim2)*(dim3)

        elif self.fusion_type=='concat':

            head_size_in = dim1+dim2+dim3

        self.post_compression_layer = nn.Sequential(*[FC_block(head_size_in, feature_size_comp_post*2, p_dropout_fc=p_dropout_fc),

                                                        FC_block(feature_size_comp_post*2, feature_size_comp_post, p_dropout_fc=p_dropout_fc),])

        # Survival head.

        self.n_classes = n_classes

        self.classifier = nn.Linear(feature_size_comp_post, self.n_classes)

        # Init weights.

        self.apply(self._init_weights)

    def _init_weights(self, module):

        if isinstance(module, nn.Linear):

            nn.init.xavier_normal_(module.weight)

            if module.bias is not None:

                module.bias.data.zero_()

    def forward_attention(self, h):

        A_ = self.attention_survival_net(h)  # h shape is N_tilesxdim

        A_raw = torch.transpose(A_, 1, 0)  # K_attention_classesxN_tiles

        A = F.softmax(A_raw, dim=-1)  # #normalize attentions scores over tiles

        return A_raw, A

    def forward_fusion(self, h1, h2, h3):

        if self.fusion_type=='bilinear':

            # Append 1 to retain unimodal embeddings in the fusion

            h1 = torch.cat((h1, torch.ones(1, 1, dtype=torch.float, device=h1.device)), -1)

            h2 = torch.cat((h2, torch.ones(1, 1, dtype=torch.float, device=h2.device)), -1)

            h3 = torch.cat((h3, torch.ones(1, 1, dtype=torch.float, device=h3.device)), -1)

            return torch.kron(torch.kron(h1, h2), h3)

        elif self.fusion_type=='kron':

            return torch.kron(torch.kron(h1, h2), h3)

        elif self.fusion_type=='concat':

            return torch.cat([h1, h2, h3], dim=-1)

        else:

            print('Not implemeted') 

            #raise Exception ... 

    def forward_survival(self, logits):

        Y_hat = torch.topk(logits, 1, dim=1)[1]

        # Model outputs the hazards with sigmoid activation function.

        hazards = torch.sigmoid(logits) #size [1, n_classes] h(t|X) := P(T=t|T>=t,X)

        #S(t|X) := P(T>=t|X) = TT (1-h(s|X)) for s=1,t. This is computed for each discrete time point t. So for s=1 there is no cum prod. 

        survival = torch.cumprod(1 - hazards, dim=1) #size [1, n_classes]

        return hazards, survival, Y_hat

    def forward(self, h, stage, h_mol):

        # H&E embedding.

        h = self.compression_layer(h)

        # Attention MIL and first-order pooling.

        A_raw, A = self.forward_attention(h) # 1xN tiles

        h_hist = A @ h #torch.Size([1, dim_embedding]) [Sum over N(aihi,1), ..., Sum over N(aihi,dim_embedding)]

        # Stage learnable embedding.

        stage = self.encoding_stage_net(stage)

        # Compression h_mol.

        h_mol = self.encoding_mol_net(h_mol)

        # Attention gates on each modality.

        h_hist, stage, h_mol = self.attn_modalities(h_hist, stage, h_mol)

        # Post-compressiong H&E slide embedding.

        h_hist = self.post_compression_layer_he(h_hist)

        # Fusion.

        m = self.forward_fusion(h_hist, stage, h_mol)

        # Post-compression of multimodal embedding.

        m = self.post_compression_layer(m)

        # Survival head.

        logits  = self.classifier(m)

        hazards, survival, Y_hat = self.forward_survival(logits)

        return hazards, survival, Y_hat, A_raw, m