import pytest

import selfies as sf

class Entry:

    def __init__(self, selfies, symbols, label, one_hot):

        self.selfies = selfies

        self.symbols = symbols

        self.label = label

        self.one_hot = one_hot

@pytest.fixture()

def dataset():

    stoi = {"[nop]": 0, "[O]": 1, ".": 2, "[C]": 3, "[F]": 4}

    itos = {i: c for c, i in stoi.items()}

    pad_to_len = 4

    entries = [

        Entry(selfies="",

              symbols=[],

              label=[0, 0, 0, 0],

              one_hot=[[1, 0, 0, 0, 0],

                       [1, 0, 0, 0, 0],

                       [1, 0, 0, 0, 0],

                       [1, 0, 0, 0, 0]]),

        Entry(selfies="[C][C][C]",

              symbols=["[C]", "[C]", "[C]"],

              label=[3, 3, 3, 0],

              one_hot=[[0, 0, 0, 1, 0],

                       [0, 0, 0, 1, 0],

                       [0, 0, 0, 1, 0],

                       [1, 0, 0, 0, 0]]),

        Entry(selfies="[C].[C]",

              symbols=["[C]", ".", "[C]"],

              label=[3, 2, 3, 0],

              one_hot=[[0, 0, 0, 1, 0],

                       [0, 0, 1, 0, 0],

                       [0, 0, 0, 1, 0],

                       [1, 0, 0, 0, 0]]),

        Entry(selfies="[C][O][C][F]",

              symbols=["[C]", "[O]", "[C]", "[F]"],

              label=[3, 1, 3, 4],

              one_hot=[[0, 0, 0, 1, 0],

                       [0, 1, 0, 0, 0],

                       [0, 0, 0, 1, 0],

                       [0, 0, 0, 0, 1]]),

        Entry(selfies="[C][O][C]",

              symbols=["[C]", "[O]", "[C]"],

              label=[3, 1, 3, 0],

              one_hot=[[0, 0, 0, 1, 0],

                       [0, 1, 0, 0, 0],

                       [0, 0, 0, 1, 0],

                       [1, 0, 0, 0, 0]])

    ]

    return entries, (stoi, itos, pad_to_len)

@pytest.fixture()

def dataset_flat_hots(dataset):

    flat_hots = []

    for entry in dataset[0]:

        hot = [elm for vec in entry.one_hot for elm in vec]

        flat_hots.append(hot)

    return flat_hots

def test_len_selfies(dataset):

    for entry in dataset[0]:

        assert sf.len_selfies(entry.selfies) == len(entry.symbols)

def test_split_selfies(dataset):

    for entry in dataset[0]:

        assert list(sf.split_selfies(entry.selfies)) == entry.symbols

def test_get_alphabet_from_selfies(dataset):

    entries, (vocab_stoi, _, _) = dataset

    selfies = [entry.selfies for entry in entries]

    alphabet = sf.get_alphabet_from_selfies(selfies)

    alphabet.add("[nop]")

    alphabet.add(".")

    assert alphabet == set(vocab_stoi.keys())

def test_selfies_to_encoding(dataset):

    entries, (vocab_stoi, vocab_itos, pad_to_len) = dataset

    for entry in entries:

        label, one_hot = sf.selfies_to_encoding(

            entry.selfies, vocab_stoi, pad_to_len, "both"

        )

        assert label == entry.label

        assert one_hot == entry.one_hot

        # recover original selfies

        selfies = sf.encoding_to_selfies(label, vocab_itos, "label")

        selfies = selfies.replace("[nop]", "")

        assert selfies == entry.selfies

        selfies = sf.encoding_to_selfies(one_hot, vocab_itos, "one_hot")

        selfies = selfies.replace("[nop]", "")

        assert selfies == entry.selfies

def test_selfies_to_flat_hot(dataset, dataset_flat_hots):

    entries, (vocab_stoi, vocab_itos, pad_to_len) = dataset

    batch = [entry.selfies for entry in entries]

    flat_hots = sf.batch_selfies_to_flat_hot(batch, vocab_stoi, pad_to_len)

    assert flat_hots == dataset_flat_hots

    # recover original selfies

    recovered = sf.batch_flat_hot_to_selfies(flat_hots, vocab_itos)

    assert batch == [s.replace("[nop]", "") for s in recovered]