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
