from .data import Data
from .model import TrainingBatch
import torch


def _shuffle(x: torch.Tensor) -> torch.Tensor:
    order = torch.randperm(len(x))
    return x[order]


class Batcher(object):
    def __init__(self, data: Data, batch_size: int=512,
        shuffle: bool=True) -> None:

        if not isinstance(data, Data):
            raise TypeError('data must be an instance of Data')

        self.data = data
        self.batch_size = int(batch_size)
        self.shuffle = bool(shuffle)

    def __iter__(self) -> TrainingBatch:
        edge_types = list(self.data.edge_types.values())

        edge_lists = [ [ adj_mat.indices().transpose(0, 1) \
            for adj_mat in et.adjacency_matrices ] \
                for et in edge_types ]

        if self.shuffle:
            edge_lists = [ [ _shuffle(lst) for lst in edge_lst ] \
                for edge_lst in edge_lists ]

        offsets = [ [ 0 ] * len(et.adjacency_matrices) \
            for et in edge_types ]

        while True:
            candidates = [ edge_idx for edge_idx, edge_ofs in enumerate(offsets) \
                if len([ rel_idx for rel_idx, rel_ofs in enumerate(edge_ofs) \
                    if rel_ofs < len(edge_lists[edge_idx][rel_idx]) ]) > 0 ]
            if len(candidates) == 0:
                break

            edge_idx = torch.randint(0, len(candidates), (1,)).item()
            edge_idx = candidates[edge_idx]
            candidates = [ rel_idx \
                for rel_idx, rel_ofs in enumerate(offsets[edge_idx]) \
                    if rel_ofs < len(edge_lists[edge_idx][rel_idx]) ]

            rel_idx = torch.randint(0, len(candidates), (1,)).item()
            rel_idx = candidates[rel_idx]

            lst = edge_lists[edge_idx][rel_idx]
            et = edge_types[edge_idx]
            ofs = offsets[edge_idx][rel_idx]
            lst = lst[ofs:ofs+self.batch_size]
            offsets[edge_idx][rel_idx] += self.batch_size

            b = TrainingBatch(et.vertex_type_row, et.vertex_type_column,
                rel_idx, lst, torch.full((len(lst),), self.data.target_value,
                    dtype=torch.float32))

            yield b