import collections
import io

import numpy

import chainer
from chainer import cuda


def make_array(tokens, vocab, add_eos=True):
    unk_id = vocab["<unk>"]
    eos_id = vocab["<eos>"]
    ids = [vocab.get(token, unk_id) for token in tokens]
    if add_eos:
        ids.append(eos_id)
    return numpy.array(ids, "i")


def transform_to_array(dataset, vocab, with_label=True):
    if with_label:
        return [
            (make_array(tokens, vocab), numpy.array([cls], "i"))
            for tokens, cls in dataset
        ]
    else:
        return [make_array(tokens, vocab) for tokens in dataset]


def convert_seq(batch, device=None, with_label=True):
    def to_device_batch(batch):
        if device is None:
            return batch
        elif device < 0:
            return [chainer.dataset.to_device(device, x) for x in batch]
        else:
            xp = cuda.cupy.get_array_module(*batch)
            concat = xp.concatenate(batch, axis=0)
            sections = numpy.cumsum([len(x) for x in batch[:-1]], dtype="i")
            concat_dev = chainer.dataset.to_device(device, concat)
            batch_dev = cuda.cupy.split(concat_dev, sections)
            return batch_dev

    if with_label:
        return {
            "xs": to_device_batch([x for x, _ in batch]),
            "ys": to_device_batch([y for _, y in batch]),
        }
    else:
        return to_device_batch([x for x in batch])
