"""
Based on an older version of https://github.com/adrianjav/heterogeneous_vaes
"""

from __future__ import annotations

from typing import List
from functools import reduce

import torch
from utils.distributions import get_distribution_by_name, Base


def _get_distributions(dists_names) -> List[Base]:
    dists = []

    for i, name in enumerate(dists_names):
        is_gammatrick = name[-1] == '*'
        is_bernoullitrick = name[-1] == '+'

        if is_gammatrick or is_bernoullitrick:
            name = name[:-1]

        if 'categorical' in name or 'ordinal' in name:
            pos = name.find('(')
            num_probs = int(name[pos + 1 : name.find(')')])
            name = name[:pos]
        else:
            num_probs = 1

        if is_gammatrick:
            if num_probs == 1:
                dist_i = get_distribution_by_name('gammatrick')[name]()
            else:
                dist_i = get_distribution_by_name('gammatrick')[name](num_probs)
        elif is_bernoullitrick:
            if num_probs == 1:
                dist_i = get_distribution_by_name('bernoullitrick')[name]()
            else:
                dist_i = get_distribution_by_name('bernoullitrick')[name](num_probs)
        else:
            if num_probs == 1:
                dist_i = get_distribution_by_name(name)()
            else:
                dist_i = get_distribution_by_name(name)(num_probs)

        dists += [dist_i]

    return dists


class ProbabilisticModel(object):
    def __init__(self, dists_names):
        self.dists = _get_distributions(dists_names)
        self.indexes = reduce(list.__add__, [[[i, j] for j in range(d.num_dists)] for i, d in enumerate(self.dists)])

    @property
    def weights(self):
        return [d.weight for d in self]

    @weights.setter
    def weights(self, values):
        if isinstance(values, torch.Tensor):
            values = values.detach().tolist()

        for w, d in zip(values, self):
            d.weight = w

    def scale_data(self, x):
        new_x = []
        for i, d in enumerate(self):
            new_x.append(d >> x[:, i])
        return torch.stack(new_x, dim=-1)

    def __rshift__(self, data):
        return self.scale_data(data)

    def params_from_data(self, x, mask):
        params = []
        for i, d in enumerate(self):
            pos = self.gathered_index(i)
            data = x[..., i] if mask is None or mask[..., pos].all() else torch.masked_select(x[..., i], mask[..., pos])
            params += d.params_from_data(data)
        return params

    def preprocess_data(self, x, mask=None):
        new_x = []
        for i, dist_i in enumerate(self.dists):
            new_x += dist_i.preprocess_data(x[:, i], mask)

        for i in range(len(self.dists), x.size(1)):
            new_x += [x[:, i]]

        return torch.stack(new_x, 1)

    def gathered_index(self, index):
        return self.indexes[index][0]

    def __len__(self):
        return len(self.indexes)

    def __getitem__(self, item) -> Base:
        if isinstance(item, int):
            return self.__getitem__(self.indexes[item])

        return self.dists[item[0]][item[1]]

    @property
    def gathered(self):
        class GatherProbabilisticModel(object):
            def __init__(self, model):
                self.model = model

            def __len__(self):
                return len(self.model.dists)

            def __getitem__(self, item):
                offset = sum([d.num_dists for d in self.model.dists[: item]])
                idxs = range(offset, offset + self.model.dists[item].num_dists)

                return idxs, self.model.dists[item]

            @property
            def weights(self):
                return [d.weight for [_, d] in self]

            @weights.setter
            def weights(self, values):
                if isinstance(values, torch.Tensor):
                    values = values.detach().tolist()

                for w, [_, d] in zip(values, self):
                    d.weight = w

            def __iter__(self):
                offset = 0
                for i, d in enumerate(self.model.dists):
                    yield list(range(offset, offset + d.num_dists)), d
                    offset += d.num_dists

            def get_param_names(self):
                names = []
                for i, dist_i in enumerate(self.model.dists):
                    if dist_i.num_dists > 1 or dist_i.size_params[0] > 1:
                        param_name = dist_i.real_parameters[0]
                        num_classes = dist_i.size_params[0] if dist_i.num_dists == 1 else dist_i.num_dists
                        names += [f'{dist_i}_{param_name}{j}_dim{i}' for j in range(num_classes)]
                    else:
                        names += [f'{dist_i}_{v}_dim{i}' for v in dist_i.real_parameters]

                return names

            def scale_data(self, x):
                new_x = []
                for i, [_, d] in enumerate(self):
                    new_x.append(d >> x[:, i])
                return torch.stack(new_x, dim=-1)

            def __rshift__(self, data):
                return self.scale_data(data)

        return GatherProbabilisticModel(self)



