text: from random import gauss
  from pandas import Categorical
  import torch
  import torch.nn as nn
  
  from opacus import PrivacyEngine
  
  # from torch.distributions.bernoulli import Bernoulli
  from torch.distributions.normal import Normal
  
  from tqdm import tqdm
  
  
  class Encoder(nn.Module):
      """Encoder, takes in x
      and outputs mu_z, sigma_z
      (diagonal Gaussian variational posterior assumed)
      """
  
      def __init__(
          self, input_dim, latent_dim, hidden_dim=32, activation=nn.Tanh, device="gpu",
      ):
          super().__init__()
          if device == "gpu":
              self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
              print(f"Encoder: {device} specified, {self.device} used")
          else:
              self.device = torch.device("cpu")
              print(f"Encoder: {device} specified, {self.device} used")
          output_dim = 2 * latent_dim
          self.latent_dim = latent_dim
          self.net = nn.Sequential(
              nn.Linear(input_dim, hidden_dim),
              activation(),
              nn.Linear(hidden_dim, hidden_dim),
              activation(),
              nn.Linear(hidden_dim, output_dim),
          )
  
      def forward(self, x):
          outs = self.net(x)
          mu_z = outs[:, : self.latent_dim]
          logsigma_z = outs[:, self.latent_dim :]
          return mu_z, logsigma_z
  
  
  class Decoder(nn.Module):
      """Decoder, takes in z and outputs reconstruction"""
  
      def __init__(
          self,
          latent_dim,
          num_continuous,
          num_categories=[0],
          hidden_dim=32,
          activation=nn.Tanh,
          device="gpu",
      ):
          super().__init__()
  
          output_dim = num_continuous + sum(num_categories)
          self.num_continuous = num_continuous
          self.num_categories = num_categories
  
          if device == "gpu":
              self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
              print(f"Decoder: {device} specified, {self.device} used")
          else:
              self.device = torch.device("cpu")
              print(f"Decoder: {device} specified, {self.device} used")
  
          self.net = nn.Sequential(
              nn.Linear(latent_dim, hidden_dim),
              activation(),
              nn.Linear(hidden_dim, hidden_dim),
              activation(),
              nn.Linear(hidden_dim, output_dim),
          )
  
      def forward(self, z):
          return self.net(z)
  
  
  class Noiser(nn.Module):
      def __init__(self, num_continuous):
          super().__init__()
          self.output_logsigma_fn = nn.Linear(num_continuous, num_continuous, bias=True)
          torch.nn.init.zeros_(self.output_logsigma_fn.weight)
          torch.nn.init.zeros_(self.output_logsigma_fn.bias)
          self.output_logsigma_fn.weight.requires_grad = False
  
      def forward(self, X):
          return self.output_logsigma_fn(X)
  
  
  class VAE(nn.Module):
      """Combines encoder and decoder into full VAE model"""
  
      def __init__(self, encoder, decoder, lr=1e-3):
          super().__init__()
          self.encoder = encoder.to(encoder.device)
          self.decoder = decoder.to(decoder.device)
          self.device = encoder.device
          self.num_categories = self.decoder.num_categories
          self.num_continuous = self.decoder.num_continuous
          self.noiser = Noiser(self.num_continuous).to(decoder.device)
          self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
          self.lr = lr
  
      def reconstruct(self, X):
          mu_z, logsigma_z = self.encoder(X)
  
          x_recon = self.decoder(mu_z)
          return x_recon
  
      def generate(self, N):
          z_samples = torch.randn_like(
              torch.ones((N, self.encoder.latent_dim)), device=self.device
          )
          x_gen = self.decoder(z_samples)
          x_gen_ = torch.ones_like(x_gen, device=self.device)
          i = 0
  
          for v in range(len(self.num_categories)):
              x_gen_[
                  :, i : (i + self.num_categories[v])
              ] = torch.distributions.one_hot_categorical.OneHotCategorical(
                  logits=x_gen[:, i : (i + self.num_categories[v])]
              ).sample()
              i = i + self.num_categories[v]
  
          x_gen_[:, -self.num_continuous :] = x_gen[
              :, -self.num_continuous :
          ] + torch.exp(self.noiser(x_gen[:, -self.num_continuous :])) * torch.randn_like(
              x_gen[:, -self.num_continuous :]
          )
          return x_gen_
  
      def loss(self, X):
          mu_z, logsigma_z = self.encoder(X)
  
          p = Normal(torch.zeros_like(mu_z), torch.ones_like(mu_z))
          q = Normal(mu_z, torch.exp(logsigma_z))
  
          divergence_loss = torch.sum(torch.distributions.kl_divergence(q, p))
  
          s = torch.randn_like(mu_z)
          z_samples = mu_z + s * torch.exp(logsigma_z)
  
          x_recon = self.decoder(z_samples)
  
          categoric_loglik = 0
          if sum(self.num_categories) != 0:
              i = 0
  
              for v in range(len(self.num_categories)):
  
                  categoric_loglik += -torch.nn.functional.cross_entropy(
                      x_recon[:, i : (i + self.num_categories[v])],
                      torch.max(X[:, i : (i + self.num_categories[v])], 1)[1],
                  ).sum()
                  i = i + self.decoder.num_categories[v]
  
          gauss_loglik = 0
          if self.decoder.num_continuous != 0:
              gauss_loglik = (
                  Normal(
                      loc=x_recon[:, -self.num_continuous :],
                      scale=torch.exp(self.noiser(x_recon[:, -self.num_continuous :])),
                  )
                  .log_prob(X[:, -self.num_continuous :])
                  .sum()
              )
  
          reconstruct_loss = -(categoric_loglik + gauss_loglik)
  
          elbo = divergence_loss + reconstruct_loss
  
          return (elbo, reconstruct_loss, divergence_loss, categoric_loglik, gauss_loglik)
  
      def train(
          self,
          x_dataloader,
          n_epochs,
          logging_freq=1,
          patience=5,
          delta=10,
          filepath=None,
      ):
          # mean_norm = 0
          # counter = 0
          log_elbo = []
          log_reconstruct = []
          log_divergence = []
          log_cat_loss = []
          log_num_loss = []
  
          # EARLY STOPPING #
          min_elbo = 0.0  # For early stopping workflow
          patience = patience  # How many epochs patience we give for early stopping
          stop_counter = 0  # Counter for stops
          delta = delta  # Difference in elbo value
  
          for epoch in range(n_epochs):
  
              train_loss = 0.0
              divergence_epoch_loss = 0.0
              reconstruction_epoch_loss = 0.0
              categorical_epoch_reconstruct = 0.0
              numerical_epoch_reconstruct = 0.0
  
              for batch_idx, (Y_subset,) in enumerate(tqdm(x_dataloader)):
                  self.optimizer.zero_grad()
                  (
                      elbo,
                      reconstruct_loss,
                      divergence_loss,
                      categorical_reconstruc,
                      numerical_reconstruct,
                  ) = self.loss(Y_subset.to(self.encoder.device))
                  elbo.backward()
                  self.optimizer.step()
  
                  train_loss += elbo.item()
                  divergence_epoch_loss += divergence_loss.item()
                  reconstruction_epoch_loss += reconstruct_loss.item()
                  categorical_epoch_reconstruct += categorical_reconstruc.item()
                  numerical_epoch_reconstruct += numerical_reconstruct.item()
  
                  # counter += 1
                  # l2_norm = 0
                  # for p in self.parameters():
                  #     if p.requires_grad:
                  #         p_norm = p.grad.detach().data.norm(2)
                  #         l2_norm += p_norm.item() ** 2
                  # l2_norm = l2_norm ** 0.5  # / Y_subset.shape[0]
                  # mean_norm = (mean_norm * (counter - 1) + l2_norm) / counter
  
              log_elbo.append(train_loss)
              log_reconstruct.append(reconstruction_epoch_loss)
              log_divergence.append(divergence_epoch_loss)
              log_cat_loss.append(categorical_epoch_reconstruct)
              log_num_loss.append(numerical_epoch_reconstruct)
  
              if epoch == 0:
  
                  min_elbo = train_loss
  
              if train_loss < (min_elbo - delta):
  
                  min_elbo = train_loss
                  stop_counter = 0  # Set counter to zero
                  if filepath != None:
                      self.save(filepath)  # Save best model if we want to
  
              else:  # elbo has not improved
  
                  stop_counter += 1
  
              if epoch % logging_freq == 0:
                  print(
                      f"\tEpoch: {epoch:2}. Elbo: {train_loss:11.2f}. Reconstruction Loss: {reconstruction_epoch_loss:11.2f}. KL Divergence: {divergence_epoch_loss:11.2f}. Categorical Loss: {categorical_epoch_reconstruct:11.2f}. Numerical Loss: {numerical_epoch_reconstruct:11.2f}"
                  )
                  # print(f"\tMean norm: {mean_norm}")
              # self.mean_norm = mean_norm
  
              if stop_counter == patience:
  
                  n_epochs = epoch + 1
  
                  break
  
          return (
              n_epochs,
              log_elbo,
              log_reconstruct,
              log_divergence,
              log_cat_loss,
              log_num_loss,
          )
  
      def diff_priv_train(
          self,
          x_dataloader,
          n_epochs,
          C=1e16,
          noise_scale=None,
          target_eps=1,
          target_delta=1e-5,
          logging_freq=1,
          sample_rate=0.1,
          patience=5,
          delta=10,
          filepath=None,
      ):
          if noise_scale is not None:
              self.privacy_engine = PrivacyEngine(
                  self,
                  sample_rate=sample_rate,
                  alphas=[1 + x / 10.0 for x in range(1, 100)] + list(range(12, 64)),
                  noise_multiplier=noise_scale,
                  max_grad_norm=C,
              )
          else:
              self.privacy_engine = PrivacyEngine(
                  self,
                  sample_rate=sample_rate,
                  alphas=[1 + x / 10.0 for x in range(1, 100)] + list(range(12, 64)),
                  target_epsilon=target_eps,
                  target_delta=target_delta,
                  epochs=n_epochs,
                  max_grad_norm=C,
              )
          self.privacy_engine.attach(self.optimizer)
  
          log_elbo = []
          log_reconstruct = []
          log_divergence = []
          log_cat_loss = []
          log_num_loss = []
  
          # EARLY STOPPING #
          min_elbo = 0.0  # For early stopping workflow
          patience = patience  # How many epochs patience we give for early stopping
          stop_counter = 0  # Counter for stops
          delta = delta  # Difference in elbo value
  
          for epoch in range(n_epochs):
              train_loss = 0.0
              divergence_epoch_loss = 0.0
              reconstruction_epoch_loss = 0.0
              categorical_epoch_reconstruct = 0.0
              numerical_epoch_reconstruct = 0.0
              # print(self.get_privacy_spent(target_delta))
  
              for batch_idx, (Y_subset,) in enumerate(tqdm(x_dataloader)):
  
                  self.optimizer.zero_grad()
                  (
                      elbo,
                      reconstruct_loss,
                      divergence_loss,
                      categorical_reconstruct,
                      numerical_reconstruct,
                  ) = self.loss(Y_subset.to(self.encoder.device))
                  elbo.backward()
                  self.optimizer.step()
  
                  train_loss += elbo.item()
                  divergence_epoch_loss += divergence_loss.item()
                  reconstruction_epoch_loss += reconstruct_loss.item()
                  categorical_epoch_reconstruct += categorical_reconstruct.item()
                  numerical_epoch_reconstruct += numerical_reconstruct.item()
  
                  # print(self.get_privacy_spent(target_delta))
                  # print(loss.item())
  
              log_elbo.append(train_loss)
              log_reconstruct.append(reconstruction_epoch_loss)
              log_divergence.append(divergence_epoch_loss)
              log_cat_loss.append(categorical_epoch_reconstruct)
              log_num_loss.append(numerical_epoch_reconstruct)
  
              if epoch == 0:
  
                  min_elbo = train_loss
  
              if train_loss < (min_elbo - delta):
  
                  min_elbo = train_loss
                  stop_counter = 0  # Set counter to zero
                  if filepath != None:
                      self.save(filepath)  # Save best model if we want to
  
              else:  # elbo has not improved
  
                  stop_counter += 1
  
              if epoch % logging_freq == 0:
                  print(
                      f"\tEpoch: {epoch:2}. Elbo: {train_loss:11.2f}. Reconstruction Loss: {reconstruction_epoch_loss:11.2f}. KL Divergence: {divergence_epoch_loss:11.2f}. Categorical Loss: {categorical_epoch_reconstruct:11.2f}. Numerical Loss: {numerical_epoch_reconstruct:11.2f}"
                  )
                  # print(f"\tMean norm: {mean_norm}")
  
              if stop_counter == patience:
  
                  n_epochs = epoch + 1
                  break
  
          return (
              n_epochs,
              log_elbo,
              log_reconstruct,
              log_divergence,
              log_cat_loss,
              log_num_loss,
          )
  
      def get_privacy_spent(self, delta):
          if hasattr(self, "privacy_engine"):
              return self.privacy_engine.get_privacy_spent(delta)
          else:
              print(
                  """This VAE object does not a privacy_engine attribute.
                  Run diff_priv_train to create one."""
              )
  
      def save(self, filename):
          torch.save(self.state_dict(), filename)
  
      def load(self, filename):
          self.load_state_dict(torch.load(filename))
  