<gh_stars>0
# This File is for all encoder models
from torch import nn
from torch.nn import functional as F
import torch
import numpy as np


# Encoder for AutoEncoder
class Encoder_AE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_AE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5 = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                     batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.conv5(out)
        return out


# Encoder for Gaussian VAE
class Encoder_GVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_GVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                        batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_logvar = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        mu = self.conv5_mu(out)
        logvar = self.conv5_logvar(out)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def kl_div(self, mu, logvar):
        # KL  = 1/2(- logvar + var + m^2  - 1 )
        kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        kld = kld/(self.N * (self.P - 2) * (self.P - 2))
        return kld

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        kld = self.kl_div(mu, logvar)
        return z, kld, mu, logvar


# Encoder for Laplacian VAE
class Encoder_LVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_LVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                        batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_logb = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        mu = self.conv5_mu(out)
        logb = self.conv5_logb(out)
        return mu, logb

    def reparameterize(self, mu, logb):
        eps1 = torch.rand_like(mu)
        eps2 = torch.rand_like(mu)
        eps = -torch.log(eps1 + 1e-10) - (-torch.log(eps2 + 1e-10))
        sample = mu + eps * logb.exp()  # Laplace(mu, b)
        return sample

    def kl_div(self, mu, logb):
        b = logb.exp()
        # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1
        kld = torch.sum(2 ** 0.5 * b * torch.exp(-mu.abs() / b) + 2 ** 0.5 * mu.abs() - logb - 0.5 * np.log(2) - 1.0)
        kld = kld / (self.N * (self.P - 2) * (self.P - 2))
        return kld

    def forward(self, x):
        mu, logb = self.encode(x)
        z = self.reparameterize(mu, logb)
        kld = self.kl_div(mu, logb)
        return z, kld, mu, logb


# Encoder for Gaussian Mixture VAE
class Encoder_GMVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_GMVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        self.n_classes = 3  # Number of noise levels
        # input size (CxWxH) = (NxPxP)
        # encode input RSS
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        # Generate h, w and c
        self.conv5_h_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_h_logvar = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                              batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_w_mu = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.conv5_w_logvar = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                              batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.qc = nn.Linear(128*(self.P-2)*(self.P-2), self.n_classes)   # [batch_size, n_classes]
        # prior generator
        self.deconv1 = self.deconv_layer(32, 16, 3, padding=0)  # shape: [batch_size, 16, P-2, P-2]
        # prior h for each cluster
        self.conv_ph_mu = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]
        self.conv_ph_logvar = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def deconv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                     batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode1(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        return out

    def infer_h(self, x):
        h_mu = self.conv5_h_mu(x)
        h_logvar = self.conv5_h_logvar(x)
        return h_mu, h_logvar

    def infer_w(self, x):
        w_mu = self.conv5_w_mu(x)
        w_logvar = self.conv5_w_logvar(x)
        return w_mu, w_logvar

    def infer_c(self, x):
        x = torch.flatten(x, start_dim=1)
        qc = F.softmax(self.qc(x), dim=-1)
        return qc

    def q_net(self, x):
        out = self.encode1(x)
        h_mu, h_logvar = self.infer_h(out)
        w_mu, w_logvar = self.infer_w(out)
        qc = self.infer_c(out)
        return h_mu, h_logvar, w_mu, w_logvar, qc

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def priorGenerator(self, w_sample):
        # w_sample: [batch_size, 32, P-4, P-4]
        batchSize = w_sample.size(0)
        h = self.deconv1(w_sample)
        ph_mu = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        ph_logvar = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        for i in range(self.n_classes):
            ph_mu[:, :, :, :, i] = self.conv_ph_mu[i](h)
            ph_logvar[:, :, :, :, i] = self.conv_ph_logvar[i](h)
        return ph_mu, ph_logvar

    def kl_w_loss(self, w_mu, w_logvar):
        # w_mu, w_logvar: [batch_size, 32, P-4, P-4]
        # KL  = 1/2(- logvar + var + m^2  - 1 )
        kl = -0.5 * torch.sum(1 + w_logvar - w_mu.pow(2) - w_logvar.exp())/(32 * (self.P-4) * (self.P-4))
        return kl

    def kl_c_loss(self, prob_c):
        # prob_c [batch_size, n_classes]
        kl = torch.sum(prob_c * (torch.log(self.n_classes * prob_c + 1e-10)))
        kl = torch.max(kl, torch.as_tensor(0.5, device=prob_c.get_device()))
        return kl

    def kl_h_loss(self, h_mu, h_logvar, ph_mu, ph_logvar, prob_c):
        # h_mu, h_logvar: [batch_size, N, P-2, P-2]
        # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes]
        # prob_c: [batch_size, n_classes]
        h_mu = h_mu.unsqueeze(-1)
        h_mu = h_mu.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        h_logvar = h_logvar.unsqueeze(-1)
        h_logvar = h_logvar.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        # KL  = 1/2 * (logvar2 - logvar1 + (var1 + (m1-m2)^2)/var2  - 1)
        # KL(h||ph)
        # shape: [batch_size, N, P-2, P-2, n_classes]
        KLD_QX_PX = 0.5 * (ph_logvar - h_logvar + (h_logvar.exp() + (h_mu - ph_mu).pow(2)) / ph_logvar.exp() - 1)
        # shape: [batch_size, n_classes]
        KLD_QX_PX = torch.sum(KLD_QX_PX, dim=(1, 2, 3))
        kl = torch.sum(KLD_QX_PX * prob_c)/(self.N * (self.P-2) * (self.P-2))
        return kl

    def forward(self, x):
        h_mu, h_logvar, w_mu, w_logvar, prob_c = self.q_net(x)
        h_sample = self.reparameterize(h_mu, h_logvar)    # h_sample: [batch_size, N, P-2, P-2]
        w_sample = self.reparameterize(w_mu, w_logvar)    # w_sample: [batch_size, 32, P-4, P-4]
        ph_mu, ph_logvar = self.priorGenerator(w_sample)  # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes]
        kl_loss_w = self.kl_w_loss(w_mu, w_logvar)
        kl_loss_c = self.kl_c_loss(prob_c)
        kl_loss_h = self.kl_h_loss(h_mu, h_logvar, ph_mu, ph_logvar, prob_c)
        kld = kl_loss_w + kl_loss_c + kl_loss_h
        return h_sample, kld, kl_loss_w, kl_loss_h, kl_loss_c


# Encoder for Laplacian Mixture VAE
class Encoder_LMVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_LMVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        self.n_classes = 3  # Number of noise levels
        # input size (CxWxH) = (NxPxP)
        # encode input RSS
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        # Generate h, w and c
        self.conv5_h_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_h_logb = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_w_mu = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.conv5_w_logb = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.qc = nn.Linear(128*(self.P-2)*(self.P-2), self.n_classes)   # [batch_size, n_classes]
        # prior generator
        self.deconv1 = self.deconv_layer(32, 16, 3, padding=0)  # shape: [batch_size, 16, P-2, P-2]
        # prior h for each cluster
        self.conv_ph_mu = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]
        self.conv_ph_logb = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def deconv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                     batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode1(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        return out

    def infer_h(self, x):
        h_mu = self.conv5_h_mu(x)
        h_logb = self.conv5_h_logb(x)
        return h_mu, h_logb

    def infer_w(self, x):
        w_mu = self.conv5_w_mu(x)
        w_logb = self.conv5_w_logb(x)
        return w_mu, w_logb

    def infer_c(self, x):
        x = torch.flatten(x, start_dim=1)
        qc = F.softmax(self.qc(x), dim=-1)
        return qc

    def q_net(self, x):
        out = self.encode1(x)
        h_mu, h_logb = self.infer_h(out)
        w_mu, w_logb = self.infer_w(out)
        qc = self.infer_c(out)
        return h_mu, h_logb, w_mu, w_logb, qc

    def reparameterize(self, mu, logb):
        eps1 = torch.rand_like(mu)
        eps2 = torch.rand_like(mu)
        eps = -torch.log(eps1 + 1e-10) - (-torch.log(eps2 + 1e-10))
        sample = mu + eps * logb.exp()  # Laplace(mu, b)
        return sample

    def priorGenerator(self, w_sample):
        # w_sample: [batch_size, 32, P-4, P-4]
        batchSize = w_sample.size(0)
        h = self.deconv1(w_sample)
        ph_mu = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        ph_logb = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        for i in range(self.n_classes):
            ph_mu[:, :, :, :, i] = self.conv_ph_mu[i](h)
            ph_logb[:, :, :, :, i] = self.conv_ph_logb[i](h)
        return ph_mu, ph_logb

    def kl_w_loss(self, w_mu, w_logb):
        b = w_logb.exp()
        logb = w_logb
        mu = w_mu
        # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1
        kl = torch.sum(2 ** 0.5 * b * torch.exp(-mu.abs() / b) + 2 ** 0.5 * mu.abs() - logb - 0.5 * np.log(2) - 1.0)/(32 * (self.P-4) * (self.P-4))
        return kl

    def kl_c_loss(self, prob_c):
        # prob_c [batch_size, n_classes]
        kl = torch.sum(prob_c * (torch.log(self.n_classes * prob_c + 1e-10)))
        kl = torch.max(kl, torch.as_tensor(0.5, device=prob_c.get_device()))
        return kl

    def kl_h_loss(self, h_mu, h_logb, ph_mu, ph_logb, prob_c):
        # h_mu, h_logb: [batch_size, N, P-2, P-2]
        # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes]
        # prob_c: [batch_size, n_classes]
        h_mu = h_mu.unsqueeze(-1)
        h_mu = h_mu.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        h_logb = h_logb.unsqueeze(-1)
        h_logb = h_logb.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        mu1 = h_mu
        mu2 = ph_mu
        b1 = h_logb.exp()
        b2 = ph_logb.exp()
        logb1 = h_logb
        logb2 = ph_logb
        # KL: (b1*exp(-|mu1-mu2|/b1)+|mu1-mu2|)/b2 + log(b2/b1) - 1
        # KL(h||ph)
        # shape: [batch_size, N, P-2, P-2, n_classes]
        KLD_QX_PX = (b1 * torch.exp(-torch.abs(mu1 - mu2) / b1) + torch.abs(mu1 - mu2)) / b2 + logb2 - logb1 - 1.0
        KLD_QX_PX = torch.sum(KLD_QX_PX, dim=(1, 2, 3))
        kl = torch.sum(KLD_QX_PX * prob_c)/(self.N * (self.P-2) * (self.P-2))
        return kl

    def forward(self, x):
        h_mu, h_logb, w_mu, w_logb, prob_c = self.q_net(x)
        h_sample = self.reparameterize(h_mu, h_logb)    # h_sample: [batch_size, N, P-2, P-2]
        w_sample = self.reparameterize(w_mu, w_logb)    # w_sample: [batch_size, 32, P-4, P-4]
        ph_mu, ph_logb = self.priorGenerator(w_sample)  # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes]
        kl_loss_w = self.kl_w_loss(w_mu, w_logb)
        kl_loss_c = self.kl_c_loss(prob_c)
        kl_loss_h = self.kl_h_loss(h_mu, h_logb, ph_mu, ph_logb, prob_c)
        kld = kl_loss_w + kl_loss_c + kl_loss_h
        return h_sample, kld, kl_loss_w, kl_loss_h, kl_loss_c

# This File is for all encoder models
from torch import nn
from torch.nn import functional as F
import torch
import numpy as np


# Encoder for AutoEncoder
class Encoder_AE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_AE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5 = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                     batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.conv5(out)
        return out


# Encoder for Gaussian VAE
class Encoder_GVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_GVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                        batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_logvar = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        mu = self.conv5_mu(out)
        logvar = self.conv5_logvar(out)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def kl_div(self, mu, logvar):
        # KL  = 1/2(- logvar + var + m^2  - 1 )
        kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        kld = kld/(self.N * (self.P - 2) * (self.P - 2))
        return kld

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        kld = self.kl_div(mu, logvar)
        return z, kld, mu, logvar


# Encoder for Laplacian VAE
class Encoder_LVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_LVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        # input size (CxWxH) = (NxPxP)
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        self.conv5_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                        batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_logb = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        mu = self.conv5_mu(out)
        logb = self.conv5_logb(out)
        return mu, logb

    def reparameterize(self, mu, logb):
        eps1 = torch.rand_like(mu)
        eps2 = torch.rand_like(mu)
        eps = -torch.log(eps1 + 1e-10) - (-torch.log(eps2 + 1e-10))
        sample = mu + eps * logb.exp()  # Laplace(mu, b)
        return sample

    def kl_div(self, mu, logb):
        b = logb.exp()
        # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1
        kld = torch.sum(2 ** 0.5 * b * torch.exp(-mu.abs() / b) + 2 ** 0.5 * mu.abs() - logb - 0.5 * np.log(2) - 1.0)
        kld = kld / (self.N * (self.P - 2) * (self.P - 2))
        return kld

    def forward(self, x):
        mu, logb = self.encode(x)
        z = self.reparameterize(mu, logb)
        kld = self.kl_div(mu, logb)
        return z, kld, mu, logb


# Encoder for Gaussian Mixture VAE
class Encoder_GMVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_GMVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        self.n_classes = 3  # Number of noise levels
        # input size (CxWxH) = (NxPxP)
        # encode input RSS
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        # Generate h, w and c
        self.conv5_h_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_h_logvar = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                              batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_w_mu = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.conv5_w_logvar = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                              batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.qc = nn.Linear(128*(self.P-2)*(self.P-2), self.n_classes)   # [batch_size, n_classes]
        # prior generator
        self.deconv1 = self.deconv_layer(32, 16, 3, padding=0)  # shape: [batch_size, 16, P-2, P-2]
        # prior h for each cluster
        self.conv_ph_mu = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]
        self.conv_ph_logvar = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def deconv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                     batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode1(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        return out

    def infer_h(self, x):
        h_mu = self.conv5_h_mu(x)
        h_logvar = self.conv5_h_logvar(x)
        return h_mu, h_logvar

    def infer_w(self, x):
        w_mu = self.conv5_w_mu(x)
        w_logvar = self.conv5_w_logvar(x)
        return w_mu, w_logvar

    def infer_c(self, x):
        x = torch.flatten(x, start_dim=1)
        qc = F.softmax(self.qc(x), dim=-1)
        return qc

    def q_net(self, x):
        out = self.encode1(x)
        h_mu, h_logvar = self.infer_h(out)
        w_mu, w_logvar = self.infer_w(out)
        qc = self.infer_c(out)
        return h_mu, h_logvar, w_mu, w_logvar, qc

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def priorGenerator(self, w_sample):
        # w_sample: [batch_size, 32, P-4, P-4]
        batchSize = w_sample.size(0)
        h = self.deconv1(w_sample)
        ph_mu = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        ph_logvar = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        for i in range(self.n_classes):
            ph_mu[:, :, :, :, i] = self.conv_ph_mu[i](h)
            ph_logvar[:, :, :, :, i] = self.conv_ph_logvar[i](h)
        return ph_mu, ph_logvar

    def kl_w_loss(self, w_mu, w_logvar):
        # w_mu, w_logvar: [batch_size, 32, P-4, P-4]
        # KL  = 1/2(- logvar + var + m^2  - 1 )
        kl = -0.5 * torch.sum(1 + w_logvar - w_mu.pow(2) - w_logvar.exp())/(32 * (self.P-4) * (self.P-4))
        return kl

    def kl_c_loss(self, prob_c):
        # prob_c [batch_size, n_classes]
        kl = torch.sum(prob_c * (torch.log(self.n_classes * prob_c + 1e-10)))
        kl = torch.max(kl, torch.as_tensor(0.5, device=prob_c.get_device()))
        return kl

    def kl_h_loss(self, h_mu, h_logvar, ph_mu, ph_logvar, prob_c):
        # h_mu, h_logvar: [batch_size, N, P-2, P-2]
        # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes]
        # prob_c: [batch_size, n_classes]
        h_mu = h_mu.unsqueeze(-1)
        h_mu = h_mu.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        h_logvar = h_logvar.unsqueeze(-1)
        h_logvar = h_logvar.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        # KL  = 1/2 * (logvar2 - logvar1 + (var1 + (m1-m2)^2)/var2  - 1)
        # KL(h||ph)
        # shape: [batch_size, N, P-2, P-2, n_classes]
        KLD_QX_PX = 0.5 * (ph_logvar - h_logvar + (h_logvar.exp() + (h_mu - ph_mu).pow(2)) / ph_logvar.exp() - 1)
        # shape: [batch_size, n_classes]
        KLD_QX_PX = torch.sum(KLD_QX_PX, dim=(1, 2, 3))
        kl = torch.sum(KLD_QX_PX * prob_c)/(self.N * (self.P-2) * (self.P-2))
        return kl

    def forward(self, x):
        h_mu, h_logvar, w_mu, w_logvar, prob_c = self.q_net(x)
        h_sample = self.reparameterize(h_mu, h_logvar)    # h_sample: [batch_size, N, P-2, P-2]
        w_sample = self.reparameterize(w_mu, w_logvar)    # w_sample: [batch_size, 32, P-4, P-4]
        ph_mu, ph_logvar = self.priorGenerator(w_sample)  # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes]
        kl_loss_w = self.kl_w_loss(w_mu, w_logvar)
        kl_loss_c = self.kl_c_loss(prob_c)
        kl_loss_h = self.kl_h_loss(h_mu, h_logvar, ph_mu, ph_logvar, prob_c)
        kld = kl_loss_w + kl_loss_c + kl_loss_h
        return h_sample, kld, kl_loss_w, kl_loss_h, kl_loss_c


# Encoder for Laplacian Mixture VAE
class Encoder_LMVAE(nn.Module):
    def __init__(self, M, P):
        super(Encoder_LMVAE, self).__init__()
        self.M = M  # The number of RF nodes
        self.N = int(M*(M-1))  # The number of links
        self.P = P  # Each node has P measurement positions
        self.n_classes = 3  # Number of noise levels
        # input size (CxWxH) = (NxPxP)
        # encode input RSS
        self.conv1 = self.conv_layer(self.N, 32, 3, padding='same')
        self.conv2 = self.conv_layer(32, 64, 3, padding='same')
        self.conv3 = self.conv_layer(64, 128, 3, padding='same')  # shape: [batch_size, 128, P, P]
        self.conv4 = self.conv_layer(128, 128, 3, padding='valid')  # shape: [batch_size, 128, P-2, P-2]
        # Generate h, w and c
        self.conv5_h_mu = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_h_logb = self.conv_layer(128, self.N, 1, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, N, P-2, P-2]
        self.conv5_w_mu = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                          batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.conv5_w_logb = self.conv_layer(128, 32, 3, padding='valid', activation='Linear',
                                            batch_norm=False)  # shape: [batch_size, 32, P-4, P-4]
        self.qc = nn.Linear(128*(self.P-2)*(self.P-2), self.n_classes)   # [batch_size, n_classes]
        # prior generator
        self.deconv1 = self.deconv_layer(32, 16, 3, padding=0)  # shape: [batch_size, 16, P-2, P-2]
        # prior h for each cluster
        self.conv_ph_mu = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]
        self.conv_ph_logb = nn.ModuleList(
            [nn.Conv2d(16, self.N, 1, padding='valid', bias=True) for i in range(self.n_classes)])  # shape: [batch_size, N, P-2, P-2]

    def conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                   batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def deconv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='LeakyReLU',
                     batch_norm=True):
        layer_out = []
        bias = True
        if batch_norm:
            bias = False
        layer_out.append(nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias))
        if batch_norm:
            layer_out.append(nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01))
        if activation == 'LeakyReLU':
            layer_out.append(nn.LeakyReLU(0.3))
        elif activation == 'ReLU':
            layer_out.append(nn.ReLU())
        elif activation == 'Linear':
            pass
        elif activation == 'Sigmoid':
            layer_out.append(nn.Sigmoid())
        else:
            raise NotImplementedError('Activation Function {} not understood.'.format(activation))
        return nn.Sequential(*layer_out)

    def encode1(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        return out

    def infer_h(self, x):
        h_mu = self.conv5_h_mu(x)
        h_logb = self.conv5_h_logb(x)
        return h_mu, h_logb

    def infer_w(self, x):
        w_mu = self.conv5_w_mu(x)
        w_logb = self.conv5_w_logb(x)
        return w_mu, w_logb

    def infer_c(self, x):
        x = torch.flatten(x, start_dim=1)
        qc = F.softmax(self.qc(x), dim=-1)
        return qc

    def q_net(self, x):
        out = self.encode1(x)
        h_mu, h_logb = self.infer_h(out)
        w_mu, w_logb = self.infer_w(out)
        qc = self.infer_c(out)
        return h_mu, h_logb, w_mu, w_logb, qc

    def reparameterize(self, mu, logb):
        eps1 = torch.rand_like(mu)
        eps2 = torch.rand_like(mu)
        eps = -torch.log(eps1 + 1e-10) - (-torch.log(eps2 + 1e-10))
        sample = mu + eps * logb.exp()  # Laplace(mu, b)
        return sample

    def priorGenerator(self, w_sample):
        # w_sample: [batch_size, 32, P-4, P-4]
        batchSize = w_sample.size(0)
        h = self.deconv1(w_sample)
        ph_mu = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        ph_logb = torch.empty((batchSize, self.N, self.P-2, self.P-2, self.n_classes), device=w_sample.get_device(), requires_grad=False)
        for i in range(self.n_classes):
            ph_mu[:, :, :, :, i] = self.conv_ph_mu[i](h)
            ph_logb[:, :, :, :, i] = self.conv_ph_logb[i](h)
        return ph_mu, ph_logb

    def kl_w_loss(self, w_mu, w_logb):
        b = w_logb.exp()
        logb = w_logb
        mu = w_mu
        # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1
        kl = torch.sum(2 ** 0.5 * b * torch.exp(-mu.abs() / b) + 2 ** 0.5 * mu.abs() - logb - 0.5 * np.log(2) - 1.0)/(32 * (self.P-4) * (self.P-4))
        return kl

    def kl_c_loss(self, prob_c):
        # prob_c [batch_size, n_classes]
        kl = torch.sum(prob_c * (torch.log(self.n_classes * prob_c + 1e-10)))
        kl = torch.max(kl, torch.as_tensor(0.5, device=prob_c.get_device()))
        return kl

    def kl_h_loss(self, h_mu, h_logb, ph_mu, ph_logb, prob_c):
        # h_mu, h_logb: [batch_size, N, P-2, P-2]
        # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes]
        # prob_c: [batch_size, n_classes]
        h_mu = h_mu.unsqueeze(-1)
        h_mu = h_mu.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        h_logb = h_logb.unsqueeze(-1)
        h_logb = h_logb.expand(-1, self.N, self.P-2, self.P-2, self.n_classes)
        mu1 = h_mu
        mu2 = ph_mu
        b1 = h_logb.exp()
        b2 = ph_logb.exp()
        logb1 = h_logb
        logb2 = ph_logb
        # KL: (b1*exp(-|mu1-mu2|/b1)+|mu1-mu2|)/b2 + log(b2/b1) - 1
        # KL(h||ph)
        # shape: [batch_size, N, P-2, P-2, n_classes]
        KLD_QX_PX = (b1 * torch.exp(-torch.abs(mu1 - mu2) / b1) + torch.abs(mu1 - mu2)) / b2 + logb2 - logb1 - 1.0
        KLD_QX_PX = torch.sum(KLD_QX_PX, dim=(1, 2, 3))
        kl = torch.sum(KLD_QX_PX * prob_c)/(self.N * (self.P-2) * (self.P-2))
        return kl

    def forward(self, x):
        h_mu, h_logb, w_mu, w_logb, prob_c = self.q_net(x)
        h_sample = self.reparameterize(h_mu, h_logb)    # h_sample: [batch_size, N, P-2, P-2]
        w_sample = self.reparameterize(w_mu, w_logb)    # w_sample: [batch_size, 32, P-4, P-4]
        ph_mu, ph_logb = self.priorGenerator(w_sample)  # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes]
        kl_loss_w = self.kl_w_loss(w_mu, w_logb)
        kl_loss_c = self.kl_c_loss(prob_c)
        kl_loss_h = self.kl_h_loss(h_mu, h_logb, ph_mu, ph_logb, prob_c)
        kld = kl_loss_w + kl_loss_c + kl_loss_h
        return h_sample, kld, kl_loss_w, kl_loss_h, kl_loss_c
# This File is for all encoder models # Encoder for AutoEncoder # The number of RF nodes # The number of links # Each node has P measurement positions # input size (CxWxH) = (NxPxP) # shape: [batch_size, 128, P, P] # shape: [batch_size, 128, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # Encoder for Gaussian VAE # The number of RF nodes # The number of links # Each node has P measurement positions # input size (CxWxH) = (NxPxP) # shape: [batch_size, 128, P, P] # shape: [batch_size, 128, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # KL  = 1/2(- logvar + var + m^2  - 1 ) # Encoder for Laplacian VAE # The number of RF nodes # The number of links # Each node has P measurement positions # input size (CxWxH) = (NxPxP) # shape: [batch_size, 128, P, P] # shape: [batch_size, 128, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # Laplace(mu, b) # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1 # Encoder for Gaussian Mixture VAE # The number of RF nodes # The number of links # Each node has P measurement positions # Number of noise levels # input size (CxWxH) = (NxPxP) # encode input RSS # shape: [batch_size, 128, P, P] # shape: [batch_size, 128, P-2, P-2] # Generate h, w and c # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, 32, P-4, P-4] # shape: [batch_size, 32, P-4, P-4] # [batch_size, n_classes] # prior generator # shape: [batch_size, 16, P-2, P-2] # prior h for each cluster # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # w_sample: [batch_size, 32, P-4, P-4] # w_mu, w_logvar: [batch_size, 32, P-4, P-4] # KL  = 1/2(- logvar + var + m^2  - 1 ) # prob_c [batch_size, n_classes] # h_mu, h_logvar: [batch_size, N, P-2, P-2] # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes] # prob_c: [batch_size, n_classes] # KL  = 1/2 * (logvar2 - logvar1 + (var1 + (m1-m2)^2)/var2  - 1) # KL(h||ph) # shape: [batch_size, N, P-2, P-2, n_classes] # shape: [batch_size, n_classes] # h_sample: [batch_size, N, P-2, P-2] # w_sample: [batch_size, 32, P-4, P-4] # ph_mu, ph_logvar: [batch_size, N, P-2, P-2, n_classes] # Encoder for Laplacian Mixture VAE # The number of RF nodes # The number of links # Each node has P measurement positions # Number of noise levels # input size (CxWxH) = (NxPxP) # encode input RSS # shape: [batch_size, 128, P, P] # shape: [batch_size, 128, P-2, P-2] # Generate h, w and c # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, 32, P-4, P-4] # shape: [batch_size, 32, P-4, P-4] # [batch_size, n_classes] # prior generator # shape: [batch_size, 16, P-2, P-2] # prior h for each cluster # shape: [batch_size, N, P-2, P-2] # shape: [batch_size, N, P-2, P-2] # Laplace(mu, b) # w_sample: [batch_size, 32, P-4, P-4] # KL: sqrt(2)*b*exp(-|mu|/b) + sqrt(2)*|mu| - log(b) - log(2)/2 - 1 # prob_c [batch_size, n_classes] # h_mu, h_logb: [batch_size, N, P-2, P-2] # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes] # prob_c: [batch_size, n_classes] # KL: (b1*exp(-|mu1-mu2|/b1)+|mu1-mu2|)/b2 + log(b2/b1) - 1 # KL(h||ph) # shape: [batch_size, N, P-2, P-2, n_classes] # h_sample: [batch_size, N, P-2, P-2] # w_sample: [batch_size, 32, P-4, P-4] # ph_mu, ph_logb: [batch_size, N, P-2, P-2, n_classes]