text: from __future__ import annotations
  
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  import pyro
  import pyro.distributions as distributions
  import pyro.distributions.constraints as constraints
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from pyro.nn import PyroModule, PyroParam
  from torch.distributions import transform_to
  
  
  def unravel_index(indices: torch.Tensor, shape_tensor: torch.Tensor) -> torch.Tensor:
      # https://github.com/pytorch/pytorch/blob/ed07f3b948a7c47be25c8a1f21037adfc58e7507/torch/functional.py#L1781
      # while waiting for unravel_index to get merged into PyTorch
      coefs = shape_tensor[1:].flipud().cumprod(dim=0).flipud()
      coefs = torch.cat((coefs, coefs.new_tensor((1,))), dim=0)
      coords = torch.div(indices[..., None], coefs, rounding_mode="trunc") % shape_tensor
  
      return coords
  
  
  class Tucker(PyroModule):
      def __init__(
          self,
          num_components: Union[int, List[int]],
          dists: Union[nn.ModuleList, List[distributions.Distribution]],
          component_logits: List[torch.Tensor] = None,
          dist_params: List[Dict[str, torch.Tensor]] = None,
      ) -> None:
          """Probabilistic model based on the Tucker decomposition. Instead of a core-tensor \
          that is `M`-dimensional it utilizes `M` different vectors of weights, where \
          the core tensor can be constructed as the outer product of the `M` vectors.
  
          Args:
              num_components (Union[int, List[int]]):
                  Amount of mixture components in each mixture. Either a single shared \
                  amount or list of specific amounts.
              dists (Union[nn.ModuleList, List[distributions.Distribution]]):
                  List of distributions to describe each dimension.
              component_logits (List[torch.Tensor], optional):
                  Mixture component weights as logits. \
                  Defaults to equally weighted mixtures.
              dist_params (List[Dict[str, torch.Tensor]], optional):
                  List of distribution parameters, formatted as if passed as **kwargs to \
                  the underlying Distribution. \
                  Defaults uniformly distributed random parameters, \
                  transformed to fit parameter constraints.
          """
          super().__init__()
          # Easy access to .device
          self.__empty__ = nn.Parameter(torch.empty(0))
  
          self.M = len(dists)
          self.Ks = (
              num_components
              if isinstance(num_components, list)
              else [num_components] * self.M
          )
  
          self._mask: Optional[torch.Tensor] = None
          self.components = nn.ModuleList()
  
          for i, dist in enumerate(dists):
              comp = PyroModule(f"dim_{i}")
              comp.dist = dist
              comp._logits = PyroParam(
                  torch.ones(self.Ks[i])
                  if component_logits is None
                  else component_logits[i],
                  constraint=constraints.real,
              )
  
              for param, constraint in dist.arg_constraints.items():
                  setattr(
                      comp,
                      param,
                      PyroParam(
                          transform_to(constraint)(torch.empty(self.Ks[i]).uniform_())
                          if dist_params is None
                          else dist_params[i][param],
                          constraint=constraint,
                      ),
                  )
              self.components.append(comp)
  
      def mask(self, mask: Union[torch.BoolTensor, torch.Tensor] = None) -> Tucker:
          """Sets mask for model in place. Similar to Distribution.mask
  
          Args:
              mask (Union[torch.BoolTensor, torch.Tensor], optional): \
                  None or masking tensor with shape (M,) or (N, M). Defaults to None.
  
          Returns:
              Tucker: self
          """
          assert (
              mask is None or 1 <= mask.dim() <= 2
          ), "Mask should be 1- or 2-dimensional or None."
          if mask is not None and mask.dim() == 1:
              mask = mask.unsqueeze(0)
          self._mask = mask
          return self
  
      def forward(
          self, data: torch.Tensor = None, n_samples: int = 1, impute: bool = False
      ) -> torch.Tensor:
          """Generate random samples if data is None. If data is defined, pass data through \
          model. Used for Pyro inference algorithms.
  
          Args:
              data (torch.Tensor, optional):
                  Input data. Defaults to None.
              n_samples (int, optional):
                  Amount of samples to generate if data is None. \
                  Defaults to 1.
              impute (bool, optional):
                  Should nan-values be imputed automatically. \
                  Defaults to False
  
          Returns:
              torch.Tensor: Either observed samples or generated samples.
          """
          device = self.__empty__.device
          if data is not None and data.dim() == 1:
              data = data.unsqueeze(0)
          N, M = (n_samples, self.M) if data is None else data.shape
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
  
          x = torch.empty((N, M), device=device)
          with pyro.plate("data", N):
              for m in pyro.markov(range(M)):
                  obs = data[..., m] if data is not None else None
                  comp = self.components[m]
                  k = pyro.sample(
                      f"k_{m+1}",
                      distributions.Categorical(logits=comp._logits),
                      infer={"enumerate": "parallel"},
                  )
                  params = {
                      p: v[k] for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs" if p != "_logits"
                  }
                  dist = comp.dist(**params)
                  if obs is not None and impute:
                      mask = ~torch.isnan(obs)
                      x_imputed = pyro.sample(f"imputed_x_{m}", dist.mask(False))
                      obs = obs.clone()
                      obs[~mask] = x_imputed[~mask].to(dtype=obs.dtype)
                  elif self._mask is not None:
                      dist = dist.mask(self._mask[:, m])
                  x[:, m] = pyro.sample(f"x_{m+1}", dist, obs=obs)
          return x
  
      def log_prob(self, x: torch.Tensor) -> torch.Tensor:
          """Evaluates log probability densities for each sample in a batch of samples
  
          Args:
              x (torch.Tensor): A single value (M,) or batch of values (N, M)
  
          Returns:
              torch.Tensor: Log probability of x
          """
          device = self.__empty__.device
          x = x.to(device)
  
          if x.dim() == 1:
              x = x.unsqueeze(0)
          N, M = x.shape
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
  
          log_p = torch.zeros(N, 1, device=x.device)
          for m in pyro.plate("features", M):
              comp = self.components[m]
              params = {p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs" if p != "_logits"}
              dist = comp.dist(**params)
              if self._mask is not None:
                  dist = dist.mask(self._mask[..., None, m])
              log_p = (
                  F.log_softmax(comp._logits, -1) + dist.log_prob(x[..., None, m]) + log_p
              )
              log_p = torch.logsumexp(log_p, -1, keepdim=True)
          return torch.logsumexp(log_p.view(N, -1), dim=-1)
  
      def impute(
          self, x: torch.Tensor, n_samples=10, log_likelihood=False, return_all=False
      ) -> torch.Tensor:
          """Impute missing data in a single sample by selecting highest likelihood \
              random realization out of `n` samples.
  
          Args:
              x (torch.Tensor): Input sample where `nan`s should be imputed away.
              n_samples (int): Amount of random samples to to use for imputation. \
                  Higher values lead to better estimation of the distribution of missing \
                  observables.
              log_likelihood (bool): Should the function return log likelihood as well. \
                  Defaults to False.
              return_all (bool): Should the function return all calculated samples. \
                  Defaults to False.
  
          Returns:
              torch.Tensor/Tuple(torch.Tensor, torch.Tensor): x with `nan`s replaced. \
                  (`x`, `log_likelihood`) if log likelihood is requested.
          """
          device = self.__empty__.device
  
          if x.dim() == 1:
              x = x.unsqueeze(0)
          N, M = x.shape
          assert N == 1, f"Can currently only impute one sample at a time."
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
          x = x.clone().to(device)
          x = x.repeat(n_samples, 1)
  
          with torch.no_grad():
              with pyro.plate("data", n_samples):
                  log_p = torch.zeros(n_samples, 1, device=x.device)
                  for m in pyro.plate("features", M):
                      comp = self.components[m]
                      if torch.isnan(x[0, m]):
                          log_pa = F.log_softmax(comp._logits, -1)
                          k = pyro.sample(
                              f"k{m+1}", distributions.Categorical(logits=log_pa)
                          )
                          params = {
                              p: v[k]
                              for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"
                              if p != "_logits"
                          }
                          dist = comp.dist(**params)
                          x_hat = pyro.sample(f"x{m+1}", dist)
                          x[:, m] = x_hat
                          log_p = log_pa + dist.log_prob(x_hat)[..., None] + log_p
                      else:
                          params = {
                              p: v
                              for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"
                              if p != "_logits"
                          }
                          dist = comp.dist(**params)
                          log_p = (
                              F.log_softmax(comp._logits, -1)
                              + dist.log_prob(x[..., None, m])
                              + log_p
                          )
                      log_p = torch.logsumexp(log_p, -1, keepdim=True)
  
                  log_p = torch.logsumexp(log_p, -1).squeeze()
                  if not return_all:
                      max_val, max_arg = torch.max(log_p, 0)
                      return x[max_arg] if not log_likelihood else (x[max_arg], max_val)
                  else:
                      return x if not log_likelihood else (x, log_p)
  
  
  class FullTucker(PyroModule):
      def __init__(
          self,
          num_components: Union[int, List[int]],
          dists: Union[nn.ModuleList, List[distributions.Distribution]],
          component_logits: torch.Tensor = None,
          dist_params: List[Dict[str, torch.Tensor]] = None,
      ) -> None:
          super().__init__()
          # Easy access to .device
          self.__empty__ = nn.Parameter(torch.empty(0))
  
          self.M = len(dists)
          self.Ks = (
              num_components
              if isinstance(num_components, list)
              else [num_components] * self.M
          )
  
          self._mask: Optional[torch.Tensor] = None
          self.components = nn.ModuleList()
  
          self.logits = PyroParam(
              torch.ones((1, np.prod(self.Ks)))
              if component_logits is None
              else component_logits,
              constraint=constraints.real,
          )
  
          for i, dist in enumerate(dists):
              comp = PyroModule(f"dim_{i}")
              comp.dist = dist
  
              for param, constraint in dist.arg_constraints.items():
                  setattr(
                      comp,
                      param,
                      PyroParam(
                          transform_to(constraint)(torch.empty(self.Ks[i]).uniform_())
                          if dist_params is None
                          else dist_params[i][param],
                          constraint=constraint,
                      ),
                  )
              self.components.append(comp)
  
      def mask(self, mask: Union[torch.BoolTensor, torch.Tensor] = None) -> Tucker:
          """Sets mask for model in place. Similar to Distribution.mask
  
          Args:
              mask (Union[torch.BoolTensor, torch.Tensor], optional): \
                  None or masking tensor with shape (M,) or (N, M). Defaults to None.
  
          Returns:
              CanonicalPolyadic: self
          """
          assert (
              mask is None or 1 <= mask.dim() <= 2
          ), "Mask should be 1- or 2-dimensional or None."
          if mask is not None and mask.dim() == 1:
              mask = mask.unsqueeze(0)
          self._mask = mask
          return self
  
      def forward(
          self, data: torch.Tensor = None, n_samples: int = 1, impute: bool = False
      ) -> torch.Tensor:
          """Generate random samples if data is None. If data is defined, pass data through \
          model. Used for Pyro inference algorithms.
  
          Args:
              data (torch.Tensor, optional):
                  Input data. Defaults to None.
              n_samples (int, optional):
                  Amount of samples to generate if data is None. \
                  Defaults to 1.
              impute (bool, optional):
                  Should nan-values be imputed automatically. \
                  Defaults to False
  
          Returns:
              torch.Tensor: Either observed samples or generated samples.
          """
          device = self.__empty__.device
          if data is not None and data.dim() == 1:
              data = data.unsqueeze(0)
          N, M = (n_samples, self.M) if data is None else data.shape
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
  
          x = torch.empty((N, M), device=device)
          shape_tensor = torch.LongTensor(self.Ks, device=device)
          with pyro.plate("data", N):
              k = pyro.sample(
                  "k",
                  distributions.Categorical(logits=self.logits),
                  infer={"enumerate": "parallel"},
              )
              # linear idx -> M-dimensional index
              idx = unravel_index(k, shape_tensor)
              for m in pyro.plate("features", M):
                  obs = data[..., m] if data is not None else None
                  comp = self.components[m]
                  params = {p: v[idx[..., m]] for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"}
                  dist = comp.dist(**params)
                  if obs is not None and impute:
                      mask = ~torch.isnan(obs)
                      x_imputed = pyro.sample(f"imputed_x_{m}", dist.mask(False))
                      obs = obs.clone()
                      obs[~mask] = x_imputed[~mask].to(dtype=obs.dtype)
                  elif self._mask is not None:
                      dist = dist.mask(self._mask[:, m])
                  x[:, m] = pyro.sample(f"x_{m}", dist, obs=obs)
          return x
  
      def log_prob(self, x: torch.Tensor) -> torch.Tensor:
          device = self.__empty__.device
          x = x.to(device)
  
          if x.dim() == 1:
              x = x.unsqueeze(0)
          N, M = x.shape
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
  
          log_p = torch.zeros(N, *self.Ks, device=x.device)
          for m in pyro.plate("features", M):
              comp = self.components[m]
              params = {p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"}
              dist = comp.dist(**params)
              if self._mask is not None:
                  dist = dist.mask(self._mask[..., None, m])
  
              # Change view to allow broadcasting
              # B x M1 x .. x Mm x .. x MM
              # +
              # B x  1 x .. x Mm x  1 x  1
              # =
              # B x M1 x .. x Mm x .. x MM
  
              shape = [1] * M
              shape[m] = self.Ks[m]
              log_p += dist.log_prob(x[..., None, m]).view(N, *shape)
          log_p += F.log_softmax(self.logits, -1).view(1, *self.Ks)
          return torch.logsumexp(log_p.view(N, -1), dim=-1)
  
      def impute(
          self, x: torch.Tensor, n_samples=10, log_likelihood=False, return_all=False
      ) -> torch.Tensor:
          """Impute missing data in a single sample by selecting highest likelihood \
              random realization out of `n` samples.
  
          Args:
              x (torch.Tensor): Input sample where `nan`s should be imputed away.
              n_samples (int): Amount of random samples to to use for imputation. \
                  Higher values lead to better estimation of the distribution of missing \
                  observables.
              log_likelihood (bool): Should the function return log likelihood as well. \
                  Defaults to False.
              return_all (bool): Should the function return all calculated samples. \
                  Defaults to False.
  
          Returns:
              torch.Tensor/Tuple(torch.Tensor, torch.Tensor): x with `nan`s replaced. \
                  (`x`, `log_likelihood`) if log likelihood is requested.
          """
          device = self.__empty__.device
  
          if x.dim() == 1:
              x = x.unsqueeze(0)
          N, M = x.shape
          assert N == 1, f"Can currently only impute one sample at a time."
          assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."
          x = x.clone().to(device)
          x = x.repeat(n_samples, 1)
  
          with torch.no_grad():
              with pyro.plate("data", n_samples):
                  idx = torch.arange(n_samples)
                  missing_dims = []
  
                  log_p = F.log_softmax(self.logits, -1).view(1, *self.Ks)
                  for m in range(M):
                      if torch.isnan(x[0, m]):
                          missing_dims.append(m)
                          continue
                      comp = self.components[m]
                      params = {p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"}
                      dist = comp.dist(**params)
                      shape = [1] * M
                      shape[m] = self.Ks[m]
                      log_p = log_p + dist.log_prob(x[..., None, m]).view(
                          n_samples, *shape
                      )
  
                  if len(missing_dims):
                      log_pxs = torch.logsumexp(
                          log_p, tuple(-x for x in range(1, self.M + 1)), keepdim=True
                      )
                      log_pa_xs = log_p - log_pxs
                      assert all(
                          torch.isclose(
                              torch.exp(
                                  torch.logsumexp(
                                      log_pa_xs, tuple(-x for x in range(1, self.M + 1))
                                  ).squeeze()
                              ),
                              torch.tensor(1, dtype=log_pa_xs.dtype),
                          )
                      )
                      k = pyro.sample(
                          f"k", distributions.Categorical(logits=log_pa_xs.flatten(1))
                      )
                      log_p = log_p.flatten(1)[idx, k]
  
                      k_unravel = unravel_index(k, torch.tensor(self.Ks, device=device))
                      for m in missing_dims:
                          comp = self.components[m]
                          params = {
                              p: v[k_unravel[..., m]]
                              for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"
                          }
                          dist = comp.dist(**params)
                          x_hat = pyro.sample(f"x{m+1}_hat", dist)
                          x[:, m] = x_hat
                          log_p = log_p + dist.log_prob(x_hat).view(n_samples, 1)
  
                  log_p = torch.logsumexp(log_p.view(n_samples, -1), -1).squeeze()
                  if not return_all:
                      max_val, max_arg = torch.max(log_p, 0)
                      return x[max_arg] if not log_likelihood else (x[max_arg], max_val)
                  else:
                      return x if not log_likelihood else (x, log_p)
  
  
  if __name__ == "__main__":
      import numpy as np
      from pyro.distributions import Normal
  
      from src.utils import fit
  
      rng = np.random.default_rng(42)
  
      sample = (
          rng.normal(0, 1, (1000, 8))
          + rng.normal(0.3, 0.7, (1000, 8))
          + rng.normal(0.6, 1.3, (1000, 8))
          + rng.normal(1.1, 0.2, (1000, 8))
      )
      # model = Tucker(list(np.arange(2, 10)), [Normal]*8)
      model = FullTucker(3, [Normal] * 8)
      lp = model.log_prob(torch.tensor(sample))
      p = np.exp(lp.detach())
      print(p.mean())
      fit(
          model,
          train_data=torch.tensor(sample),
          n_epochs=1,
          batch_size=128,
          use_svi=False,
          optimizer=torch.optim.Adam,
      )
      lp = model.log_prob(torch.tensor(sample))
      p = np.exp(lp.detach())
      print(p.mean())
  
      s = sample[0]
      print(s)
      print(model.log_prob(torch.tensor(s)))
      s[[1, 5, 7]] = np.nan
      x, l = model.impute(
          torch.tensor(s), n_samples=50, log_likelihood=True, return_all=False
      )
      print(x)
      print(l)
      print(model.log_prob(x))
      pass
  