from tkinter import Y
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.functional import relu

from models.pointnet_utils import PointNetEncoder
from models.pointnet2_utils import PointNetSetAbstraction,PointNetFeaturePropagation

class ECGnet(nn.Module):
    def __init__(self, in_ch=3+4, out_ch=3, num_input=1024, z_dims=16):
        super(ECGnet, self).__init__()


        self.encoder_signal = CRNN()

        # decode for signal
        self.elu = nn.ELU(inplace=True)
        self.fc1 = nn.Linear(z_dims, 256*2)
        self.fc2 = nn.Linear(256*2, 512*2)
        self.up = nn.Upsample(size=(8, 512), mode='bilinear')
        self.deconv = DoubleDeConv(1, 1)

        self.decoder_MI = nn.Sequential(
            nn.Linear(z_dims, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, out_ch),
        )

        
    def reparameterize(self, mu, log_var):
        """
        :param mu: mean from the encoder's latent space
        :param log_var: log variance from the encoder's latent space
        """
        std = torch.exp(0.5*log_var) # standard deviation
        eps = torch.randn(log_var.shape).to(std.device) # `randn_like` as we need the same size
        sample = mu + (eps * std) # sampling as if coming from the input space
        return sample

    def decode_signal(self, latent_z): # P(x|z, c)
        '''
        z: (bs, latent_size)
        '''
        inputs = latent_z
        f = self.elu(self.fc1(inputs))
        f = self.elu(self.fc2(f))
        u = self.up(f.reshape(f.shape[0], 1, 8, -1))
        dc = self.deconv(u)

        return dc
    
    def forward(self, partial_input, signal_input):   

        mu_signal, std_signal = self.encoder_signal(signal_input)
        latent_z_signal = self.reparameterize(mu_signal, std_signal)
        y_ECG = self.decode_signal(latent_z_signal)
        y_MI = self.decoder_MI(latent_z_signal)
        y_MI = nn.Softmax(dim=1)(y_MI)

        return y_MI, y_ECG, mu_signal, std_signal

class InferenceNet(nn.Module):
    def __init__(self, in_ch=3+4, out_ch=3, num_input=1024, z_dims=16):
        super(InferenceNet, self).__init__()

        self.z_dims = z_dims

        # PointNet++ Encoder
        self.sa1 = PointNetSetAbstraction(npoint=num_input, radius=0.2, nsample=64, in_channel=in_ch, mlp=[64, 64, 128], group_all=False)
        self.sa2 = PointNetSetAbstraction(128, 0.4, 64, 128 + 3, [128, 128, 256], False)
        self.sa3 = PointNetSetAbstraction(16, 0.8, 32, 256 + 3, [256, 512, 1024], False)
        self.fc11 = nn.Linear(1024*16, z_dims*2)

        # PointNet++ Decoder
        self.fc12 = nn.Linear(z_dims*2, 1024) # feat_ECG = H*feat_MI + epsilon
        self.fp3 = PointNetFeaturePropagation(1280, [256, 256])
        self.fp2 = PointNetFeaturePropagation(384, [256, 128])
        self.fp1 = PointNetFeaturePropagation(128, [128, 128, 128])
        self.conv1 = nn.Conv1d(128, 128, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.drop1 = nn.Dropout(0.5)
        self.conv2 = nn.Conv1d(128, out_ch, 1) 

        self.decoder_geometry = BetaVAE_Decoder(num_input, num_input//4, in_ch, z_dims) # in_ch -> out_ch*3
        
        self.encoder_signal = CRNN()

        # decode for signal
        self.elu = nn.ELU(inplace=True)
        self.fc1 = nn.Linear(z_dims, 256*2)
        self.fc2 = nn.Linear(256*2, 512*2)
        self.up = nn.Upsample(size=(8, 512), mode='bilinear')
        self.deconv = DoubleDeConv(1, 1)

    def reparameterize(self, mu, log_var):
        """
        :param mu: mean from the encoder's latent space
        :param log_var: log variance from the encoder's latent space
        """
        std = torch.exp(0.5*log_var) # standard deviation
        eps = torch.randn(log_var.shape).to(std.device) # `randn_like` as we need the same size
        sample = mu + (eps * std) # sampling as if coming from the input space
        return sample

    def decode_signal(self, latent_z): # P(x|z, c)
        '''
        z: (bs, latent_size)
        '''
        inputs = latent_z
        f = self.elu(self.fc1(inputs))
        f = self.elu(self.fc2(f))
        u = self.up(f.reshape(f.shape[0], 1, 8, -1))
        dc = self.deconv(u)

        return dc
    
    def forward(self, partial_input, signal_input):  
        num_points = partial_input.shape[-1]
        # extract ecg features
        mu_signal, std_signal = self.encoder_signal(signal_input)
        latent_z_signal = mu_signal # self.reparameterize(mu_signal, std_signal)

        # extract point cloud features      
        l0_xyz = partial_input[:,:3,:] 
        l0_points = partial_input[:,3:,:] 
        l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
        features = self.fc11(l3_points.view(-1, 1024*16))
        mu_geometry = features[:, : self.z_dims]
        std_geometry = features[:, self.z_dims:] + 1e-6
        latent_geometry = self.reparameterize(mu_signal, std_signal)
        # latent_geometry = self.fc11(l3_points.view(-1, 1024*16))

        # fuse two features 
        # mu = torch.cat((mu_geometry, mu_signal), dim=1)
        # log_var = torch.cat((std_geometry, std_signal), dim=1)
        # latent_z = self.reparameterize(mu, log_var)
        latent_z = torch.cat((latent_z_signal, latent_geometry), dim=1)

        # segment point cloud
        anatomy_signal_feat = F.relu(self.fc12(latent_z))
        anatomy_signal_feat = anatomy_signal_feat.view(-1, 1024, 1).repeat(1, 1, num_points)      
        l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, anatomy_signal_feat)
        l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
        l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)
        y_seg = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))
        y_seg = self.conv2(y_seg)   
        y_seg = nn.Softmax(dim=1)(y_seg)     

        # reconstruct point cloud and ecg
        y_coarse, y_detail = self.decoder_geometry(latent_geometry)
        y_coarse, y_detail = nn.Sigmoid()(y_coarse), nn.Sigmoid()(y_detail)  
        y_ECG = self.decode_signal(latent_z_signal)

        return y_seg, y_coarse, y_detail, y_ECG, mu_signal, std_signal

class CRNN(nn.Module):
    '''
    nh: default=256, 'size of the LSTM hidden state'
    imgH: default=8, 'the height of the input image to network'
    imgW: default=256, 'the width of the input image to network'

    :param class_labels: list[n_class]
    :return: (n_batch, n_class)
    '''

    def __init__(self, n_lead=8, z_dims=16):
        super(CRNN, self).__init__()

        n_out = 128
        self.z_dims = z_dims

        self.cnn = nn.Sequential(
            nn.Conv1d(n_lead, n_out, kernel_size=16, stride=2, padding=2),
            nn.BatchNorm1d(n_out),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv1d(n_out, n_out*2, kernel_size=16, stride=2, padding=2),
            nn.BatchNorm1d(n_out*2),
            nn.LeakyReLU(0.2, inplace=True)
            )
            

        self.rnn = BidirectionalLSTM(256, z_dims*4, z_dims*2)
        # self.rnn = nn.Sequential(
        #     BidirectionalLSTM(512, nh, nh),
        #     BidirectionalLSTM(nh, nh, 1))


    def forward(self, input):
        # conv features
        conv = self.cnn(input)
        b, c, w = conv.size()
        conv = conv.permute(2, 0, 1)  # [w, b, c]

        # rnn features
        output = self.rnn(conv).permute(1, 0, 2)
        features = torch.max(output, 1)[0]
        mean = features[:, : self.z_dims]
        std = features[:, self.z_dims:] + 1e-6
    
        return mean, std


    def backward_hook(self, module, grad_input, grad_output):
        for g in grad_input:
            g[g != g] = 0   # replace all nan/inf in gradients to zero

class BidirectionalLSTM(nn.Module):

    def __init__(self, nIn, nHidden, nOut):
        super(BidirectionalLSTM, self).__init__()

        self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
        self.embedding = nn.Linear(nHidden * 2, nOut)

    def forward(self, input):
        recurrent, _ = self.rnn(input)
        T, b, h = recurrent.size()
        t_rec = recurrent.view(T * b, h)

        output = self.embedding(t_rec)  # [T * b, nOut]
        output = output.view(T, b, -1)

        return output

class PointNet(nn.Module):
    def __init__(self, num_classes=10, n_signal=10, n_param=4, n_ECG=128):
        super(PointNet, self).__init__()
        self.k = num_classes
        self.n_signal = n_signal
        self.feat = PointNetEncoder(global_feat=False, feature_transform=True, channel=4)
        self.conv1 = torch.nn.Conv1d(1024+64+n_ECG, 512, 1)
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

        self.ECG_model = CRNN()

        self.inference_model = nn.Sequential(
            nn.Linear(1024+n_ECG, 512),
            nn.Dropout(0.5),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.Dropout(0.5),
            nn.ReLU(),
            nn.Linear(256, self.n_signal*n_param),
            nn.Sigmoid()
            )


    def forward(self, x, signal):
        n_pts = x.size()[2]
        anatomy_signal_feature, global_feature, trans_feat = self.feat(x)
        ECG_feature = self.ECG_model(signal)
        ECG_feature_extend = ECG_feature.repeat(1, 1, n_pts)
        
        anatomy_signal_feat = torch.cat([anatomy_signal_feature, ECG_feature_extend], 1)
        y1 = F.relu(self.bn1(self.conv1(anatomy_signal_feat)))
        y1 = F.relu(self.bn2(self.conv2(y1)))
        y1 = F.relu(self.bn3(self.conv3(y1)))
        y1 = self.conv4(y1)
        y1 = y1.transpose(2,1).contiguous()
        out_ATM = y1 #nn.Sigmoid()(y1)

        return out_ATM

class PointNet_plusplus(nn.Module):
    def __init__(self, num_classes=10, n_signal=10, n_param=4, n_ECG=128):
        super(PointNet_plusplus, self).__init__()
        self.n_signal = n_signal
        self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=64, in_channel= 3 + 4, mlp=[64, 64, 128], group_all=False)
        self.sa2 = PointNetSetAbstraction(128, 0.4, 64, 128 + 3, [128, 128, 256], False)
        self.sa3 = PointNetSetAbstraction(16, 0.8, 32, 256 + 3, [256, 512, 1024], False)
        self.fp3 = PointNetFeaturePropagation(1280+n_ECG, [256, 256])
        self.fp2 = PointNetFeaturePropagation(384, [256, 128])
        self.fp1 = PointNetFeaturePropagation(128, [128, 128, 128])
        self.conv1 = nn.Conv1d(128, 128, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.drop1 = nn.Dropout(0.5)
        self.conv2 = nn.Conv1d(128, num_classes, 1)

        self.ECG_model = CRNN()
        self.inference_model = nn.Sequential(
            nn.Linear(1024+n_ECG, 512),
            nn.Dropout(0.5),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.Dropout(0.5),
            nn.ReLU(),
            nn.Linear(256, self.n_signal*n_param),
            nn.Sigmoid())

    def forward(self, x, signal):
        l0_points = x
        l0_xyz = x[:,:3,:] 

        ECG_feature = self.ECG_model(signal)
 
        # Set Abstraction layers
        l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)

        ECG_feature_extend = ECG_feature.repeat(1, 1, l3_points.size()[2])  
        anatomy_signal_feat = torch.cat([l3_points, ECG_feature_extend], 1)

        # Feature Propagation layers
        l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, anatomy_signal_feat)
        l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
        l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)

        y1 = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))
        y1 = self.conv2(y1)        
        out_ATM = y1 #nn.Sigmoid()(y1)
        out_ATM = out_ATM.permute(0, 2, 1)

        return out_ATM

class BetaVAE(nn.Module):
    def __init__(self, in_ch=4, num_input=1024, num_class=2, z_dims=16):
        super(BetaVAE, self).__init__()

        self.encoder = BetaVAE_Encoder(in_ch, z_dims)
        self.decoder = BetaVAE_Decoder_new(num_input, num_class)

    def forward(self, x):
        latent_z = self.encoder(x)
        y = self.decoder(latent_z)
        return y

class BetaVAE_Encoder(nn.Module):
    def __init__(self, in_ch, z_dims):
        super(BetaVAE_Encoder, self).__init__()
        self.z_dims = z_dims
        self.mlp_conv1 = mlp_conv(in_ch, layer_dims=[128, 256])
        self.mlp_conv2 = mlp_conv(512, layer_dims=[512, 1024])

        self.fc1 = nn.Linear(1024, 1024)
        self.fc2 = nn.Linear(1024, 256)
        self.fc3 = nn.Linear(256, z_dims*2) 

    def forward(self, inputs):
        num_points = [inputs.shape[2]]
        features = self.mlp_conv1(inputs)
        features_global = point_maxpool(features, num_points, keepdim=True)
        features_global = point_unpool(features_global, num_points)
        features = torch.cat([features, features_global], dim=1)
        features = self.mlp_conv2(features)
        features = point_maxpool(features, num_points)

        features = features.view(features.size()[0], -1)
        features = self.fc1(features)
        features = self.fc2(features)
        features = self.fc3(features)
        mean = features[:, : self.z_dims]
        std = features[:, self.z_dims:] + 1e-6

        return mean, std

class BetaVAE_Decoder_new(nn.Module):
        def __init__(self, num_input, num_class=2, z_dims=16*2):
            super(BetaVAE_Decoder_new, self).__init__()
            self.out_ch = num_class
            self.n_pts = num_input 
            self.mlp = mlp(in_channels=z_dims, layer_dims=[128, 256, 512, 1024,  self.n_pts * self.out_ch]) 

        def forward(self, features):
            y = self.mlp(features).reshape(-1, self.out_ch, self.n_pts)
            
            return nn.Softmax(dim=1)(y)

class BetaVAE_Decoder_plus(nn.Module):
        def __init__(self, num_dense, num_coarse, out_ch, z_dims):
            super(BetaVAE_Decoder_plus, self).__init__()
            self.out_ch = out_ch
            self.num_coarse = num_coarse
            self.grid_size = int(np.sqrt(num_dense//num_coarse))
            self.num_fine = num_dense

            # PointNet++ Decoder
            self.fc12 = nn.Linear(z_dims*2, 1024)
            self.fp3 = PointNetFeaturePropagation(1280, [256, 256])
            self.fp2 = PointNetFeaturePropagation(384, [256, 128])
            self.fp1 = PointNetFeaturePropagation(128, [128, 128, 128])
            self.conv1 = nn.Conv1d(128, 128, 1)
            self.bn1 = nn.BatchNorm1d(128)
            self.drop1 = nn.Dropout(0.5)
            self.conv2 = nn.Conv1d(128, out_ch, 1) 


        def forward(self, latent_z, l0_xyz, l1_xyz, l2_xyz, l3_xyz):
            anatomy_signal_feat = F.relu(self.fc12(latent_z))
            coarse = anatomy_signal_feat.view(-1, 1024, 1).repeat(1, 1, self.num_coarse)      
            l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, coarse)
            l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
            l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)
            fine = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))
            fine = self.conv2(fine)   

            return coarse, fine

class BetaVAE_Decoder(nn.Module):
        def __init__(self, num_dense, num_coarse, out_ch, z_dims):
            super(BetaVAE_Decoder, self).__init__()
            self.out_ch = out_ch
            self.num_coarse = num_coarse
            self.grid_size = int(np.sqrt(num_dense//num_coarse))
            self.num_fine = num_dense

            self.mlp = mlp(in_channels=z_dims, layer_dims=[256, 512, 1024, 2048,  self.num_coarse * self.out_ch])
            x = torch.linspace(-0.05, 0.05, self.grid_size)
            y = torch.linspace(-0.05, 0.05, self.grid_size)
            self.grid = torch.cat(torch.meshgrid(x, y), dim=0).view(1, 2, self.grid_size ** 2)
            # self.grid = torch.stack(torch.meshgrid(x, y), dim=2)
            # self.grid = torch.reshape(self.grid.transpose(1, 0), [-1, 2]).unsqueeze(0)

            self.mlp_conv3 = mlp_conv(z_dims+2+out_ch, layer_dims=[512, 512, out_ch]) # here "+2" refers  to the two axes of grid

        def forward(self, latent_z):
            features = latent_z
            coarse = self.mlp(features).reshape(-1, self.num_coarse, self.out_ch)
            point_feat = coarse.unsqueeze(2).repeat(1, 1, self.grid_size * 2, 1)
            point_feat = point_feat.reshape(-1, self.out_ch, self.num_fine)

            grid_feat = self.grid.unsqueeze(2).repeat(features.shape[0], 1, self.num_coarse, 1).to(features.device)
            grid_feat = grid_feat.reshape(features.shape[0], -1, self.num_fine) 
            global_feat = features.unsqueeze(2).repeat(1, 1, self.num_fine)
            feat = torch.cat([grid_feat, point_feat, global_feat], dim=1)
          
            center = point_feat.reshape(-1, self.num_fine, self.out_ch)
            fine = self.mlp_conv3(feat).transpose(1, 2) + center

            return coarse, fine

def point_maxpool(features, npts, keepdim=True):
    splitted = torch.split(features, npts[0], dim=1)
    outputs = [torch.max(f, dim=2, keepdim=keepdim)[0] for f in splitted] # modified by Lei in 2022/02/10
    return torch.cat(outputs, dim=0)
    # return torch.max(features, dim=2, keepdims=keepdims)[0]

def point_unpool(features, npts):
    features = torch.split(features, features.shape[0], dim=0)
    outputs = [f.repeat(1, 1, npts[i]) for i, f in enumerate(features)]
    # outputs = [torch.tile(f, [1, 1, npts[i]]) for i, f in enumerate(features)]
    return torch.cat(outputs, dim=0)
    # return features.repeat([1, 1, 256])

class mlp_conv(nn.Module):
    def __init__(self, in_channels, layer_dims):
        super(mlp_conv, self).__init__()
        self.layer_dims = layer_dims
        for i, out_channels in enumerate(self.layer_dims):
            layer = torch.nn.Conv1d(in_channels, out_channels, kernel_size=1)
            setattr(self, 'conv_' + str(i), layer)
            in_channels = out_channels

    def __call__(self, inputs):
        outputs = inputs
        dims = len(self.layer_dims)
        for i in range(dims):
            layer = getattr(self, 'conv_' + str(i))
            if i == dims - 1:
                outputs = layer(outputs)
            else:
                outputs = relu(layer(outputs))
        return outputs

class mlp(nn.Module):
    def __init__(self, in_channels, layer_dims):
        super(mlp, self).__init__()
        self.layer_dims = layer_dims
        for i, out_channels in enumerate(layer_dims):
            layer = torch.nn.Linear(in_channels, out_channels)
            setattr(self, 'fc_' + str(i), layer)
            in_channels = out_channels

    def __call__(self, inputs):
        outputs = inputs
        dims = len(self.layer_dims)
        for i in range(dims):
            layer = getattr(self, 'fc_' + str(i))
            if i == dims - 1:
                outputs = layer(outputs)
            else:
                outputs = relu(layer(outputs))
        return outputs

class DoubleDeConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleDeConv, self).__init__()
        self.conv = nn.Sequential(
            nn.ConvTranspose2d(in_ch, out_ch, kernel_size=(3, 3), padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ELU(inplace=True),
            nn.ConvTranspose2d(out_ch, out_ch, kernel_size=(3, 3), padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ELU(inplace=True)
        )

    def forward(self, input):
        return self.conv(input)

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=(3, 3), padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ELU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=(3, 3), padding=1),
            nn.BatchNorm2d(out_ch),
            # nn.ELU(inplace=True)
        )

    def forward(self, input):
        return self.conv(input)


if __name__ == "__main__":
    x = torch.rand(3, 4, 2048)
    conditions = torch.rand(3, 2, 1)

    network = BetaVAE()
    y_coarse, y_detail = network(x, conditions)
    print(y_coarse.size(), y_detail.size())