from abc import ABC
from typing import List, Union, Dict

import torch
from torch import nn

from pyro.distributions import (
    ConditionalDistribution,
    ConditionalTransform,
    ConditionalTransformedDistribution,
)
from torch.distributions import (
    Distribution,
    Transform,
    TransformedDistribution,
)


class DistributionModule(Distribution, nn.Module, ABC):
    def __init__(self, validate_args: bool = False):
        Distribution.__init__(self, validate_args=validate_args)
        nn.Module.__init__(self)

    def __repr__(self):
        return nn.Module.__repr__(self)


class ConditionalDistributionModule(ConditionalDistribution, nn.Module, ABC):
    def __init__(self):
        ConditionalDistribution.__init__(self)
        nn.Module.__init__(self)


class ConditionalTransformedDistributionModule(
    ConditionalTransformedDistribution, nn.Module
):
    def __init__(
        self,
        base_dist: Union[ConditionalDistribution, Distribution],
        transforms: Union[
            ConditionalTransform,
            List[ConditionalTransform],
            Dict[str, Union[ConditionalTransform, Transform]],
            Transform,
            List[Transform],
            None,
        ],
    ):
        if isinstance(transforms, ConditionalTransform) or isinstance(
            transforms, Transform
        ):
            transforms = [transforms]
        elif isinstance(transforms, dict):
            transforms = [transforms[k] for k in sorted(transforms.keys())]

        if transforms is not None:
            transforms = nn.ModuleList(list(transforms))
        else:
            transforms = []

        self._base_dist_repr = base_dist.__repr__()

        nn.Module.__init__(self)
        self.base_dist = base_dist
        self.transforms = transforms

    def condition(self, context):
        base_dist = (
            self.base_dist.condition(context)
            if isinstance(self.base_dist, ConditionalDistribution)
            else self.base_dist
        )
        transforms = [
            (
                t.condition(context)
                if isinstance(t, ConditionalTransform)
                else t
            )
            for t in self.transforms
        ]

        return TransformedDistribution(base_dist, transforms)

    def clear_cache(self):
        pass

    def __repr__(self):
        s = self.__class__.__name__ + "("
        s += "\n  (base_dist): " + str(self._base_dist_repr).replace(
            "\n", "\n  "
        )
        s += "\n  (transforms): " + str(self.transforms).replace("\n", "\n  ")
        s += "\n"
        return s


class TransformedDistributionModule(DistributionModule):
    def __init__(
        self,
        base_dist: Distribution,
        transforms: Union[
            Transform, List[Transform], Dict[str, Transform], None
        ],
    ):
        super().__init__()

        self.base_dist = base_dist

        if isinstance(transforms, ConditionalTransform) or isinstance(
            transforms, Transform
        ):
            transforms = [transforms]
        elif hasattr(transforms, "keys"):
            transforms = [transforms[k] for k in sorted(transforms.keys())]

        self.transforms = nn.ModuleList(transforms)

    def rsample(self, sample_shape=torch.Size()):
        return TransformedDistribution(
            base_distribution=self.base_dist,
            transforms=[t for t in self.transforms],
        ).rsample(sample_shape)

    def log_prob(self, value):
        return TransformedDistribution(
            base_distribution=self.base_dist,
            transforms=[t for t in self.transforms],
            validate_args=False,
        ).log_prob(value)

    def __repr__(self):
        s = self.__class__.__name__ + "("
        s += "\n  (base_dist): " + str(self.base_dist).replace("\n", "\n  ")
        s += "\n  (transforms): " + str(self.transforms).replace("\n", "\n  ")
        s += "\n)"
        return s
