from typing import List, \
    Union, \
    Callable
from .data import Data
from .trainprep import PreparedData
import torch
from .weights import init_glorot


def _cat(matrices: List[torch.Tensor]):
    if len(matrices) == 0:
        raise ValueError('Empty list passed to _cat()')

    n = sum(a.is_sparse for a in matrices)
    if n != 0 and n != len(matrices):
        raise ValueError('All matrices must have the same layout (dense or sparse)')

    if not all(a.shape[1:] == matrices[0].shape[1:]):
        raise ValueError('All matrices must have the same dimensions apart from dimension 0')

    if not matrices[0].is_sparse:
        return torch.cat(matrices)

    total_rows = sum(a.shape[0] for a in matrices)
    indices = []
    values = []
    row_offset = 0

    for a in matrices:
        ind = a._indices().clone()
        val = a._values()
        ind[0] += row_offset
        ind = ind.transpose(0, 1)
        indices.append(ind)
        values.append(val)
        row_offset += a.shape[0]

    indices = torch.cat(indices).transpose(0, 1)
    values = torch.cat(values)

    res = _sparse_coo_tensor(indices, values)
    return res


class FastGraphConv(torch.nn.Module):
    def __init__(self,
        in_channels: int,
        out_channels: int,
        adjacency_matrix: List[torch.Tensor],
        **kwargs):

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.weight = torch.cat([
            init_glorot(in_channels, out_channels) \
            for _ in adjacency_matrix
        ], dim=1)
        self.adjacency_matrix = _cat(adjacency_matrix)

    def forward(self, x):
        x = torch.sparse.mm(x, self.weight) \
            if x.is_sparse \
            else torch.mm(x, self.weight)
        x = torch.sparse.mm(self.adjacency_matrix, x) \
            if self.adjacency_matrix.is_sparse \
            else torch.mm(self.adjacency_matrix, x)
        return x



class FastConvLayer(torch.nn.Module):
    adjacency_matrix: List[torch.Tensor]
    adjacency_matrix_backward: List[torch.Tensor]
    weight: List[torch.Tensor]
    weight_backward: List[torch.Tensor]

    def __init__(self,
        input_dim: List[int],
        output_dim: List[int],
        data: Union[Data, PreparedData],
        keep_prob: float = 1.,
        rel_activation: Callable[[torch.Tensor], torch.Tensor] = lambda x: x
        layer_activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.relu,
        **kwargs):

        super().__init__(**kwargs)

        self._check_params(input_dim, output_dim, data, keep_prob,
            rel_activation, layer_activation)

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.data = data
        self.keep_prob = keep_prob
        self.rel_activation = rel_activation
        self.layer_activation = layer_activation

        self.adjacency_matrix = None
        self.adjacency_matrix_backward = None
        self.weight = None
        self.weight_backward = None
        self.build()

    def build(self):
        self.adjacency_matrix = []
        self.adjacency_matrix_backward = []
        self.weight = []
        self.weight_backward = []
        for fam in self.data.relation_families:
            adj_mat = [ rel.adjacency_matrix \
                for rel in fam.relation_types \
                if rel.adjacency_matrix is not None     ]
            adj_mat_back = [ rel.adjacency_matrix_backward \
                for rel in fam.relation_types \
                if rel.adjacency_matrix_backward is not None ]
            weight = [ init_glorot(self.input_dim[fam.node_type_column],
                self.output_dim[fam.node_type_row]) \
                for _ in range(len(adj_mat)) ]
            weight_back = [ init_glorot(self.input_dim[fam.node_type_column],
                self.output_dim[fam.node_type_row]) \
                for _ in range(len(adj_mat_back)) ]
            adj_mat = torch.cat(adj_mat) \
                if len(adj_mat) > 0 \
                else None
            adj_mat_back = torch.cat(adj_mat_back) \
                if len(adj_mat_back) > 0 \
                else None
            self.adjacency_matrix.append(adj_mat)
            self.adjacency_matrix_backward.append(adj_mat_back)
            self.weight.append(weight)
            self.weight_back.append(weight_back)

    def forward(self, prev_layer_repr):
        for i, fam in enumerate(self.data.relation_families):
            repr_row = prev_layer_repr[fam.node_type_row]
            repr_column = prev_layer_repr[fam.node_type_column]

            adj_mat = self.adjacency_matrix[i]
            adj_mat_back = self.adjacency_matrix_backward[i]

            if adj_mat is not None:
                x = dropout(repr_column, keep_prob=self.keep_prob)
                x = torch.sparse.mm(x, self.weight[i]) \
                    if x.is_sparse \
                    else torch.mm(x, self.weight[i])
                x = torch.sparse.mm(adj_mat, repr_row) \
                    if adj_mat.is_sparse \
                    else torch.mm(adj_mat, repr_row)
                x = self.rel_activation(x)
                x = x.view(len(fam.relation_types), len(repr_row), -1)

            if adj_mat_back is not None:
                x = torch.sparse.mm(adj_mat_back, repr_row) \
                    if adj_mat_back.is_sparse \
                    else torch.mm(adj_mat_back, repr_row)

    @staticmethod
    def _check_params(input_dim, output_dim, data, keep_prob,
        rel_activation, layer_activation):

        if not isinstance(input_dim, list):
            raise ValueError('input_dim must be a list')

        if not output_dim:
            raise ValueError('output_dim must be specified')

        if not isinstance(output_dim, list):
            output_dim = [output_dim] * len(data.node_types)

        if not isinstance(data, Data) and not isinstance(data, PreparedData):
            raise ValueError('data must be of type Data or PreparedData')