import faulthandler
import random

import pytest
from rdkit.Chem import MolFromSmiles

import selfies as sf

faulthandler.enable()


@pytest.fixture()
def max_selfies_len():
    return 1000


@pytest.fixture()
def large_alphabet():
    alphabet = sf.get_semantic_robust_alphabet()
    alphabet.update([
        "[#Br]", "[#Branch1]", "[#Branch2]", "[#Branch3]", "[#C@@H1]",
        "[#C@@]", "[#C@H1]", "[#C@]", "[#C]", "[#Cl]", "[#F]", "[#H]", "[#I]",
        "[#NH1]", "[#N]", "[#O]", "[#P]", "[#Ring1]", "[#Ring2]", "[#Ring3]",
        "[#S]", "[/Br]", "[/C@@H1]", "[/C@@]", "[/C@H1]", "[/C@]", "[/C]",
        "[/Cl]", "[/F]", "[/H]", "[/I]", "[/NH1]", "[/N]", "[/O]", "[/P]",
        "[/S]", "[=Br]", "[=Branch1]", "[=Branch2]", "[=Branch3]", "[=C@@H1]",
        "[=C@@]", "[=C@H1]", "[=C@]", "[=C]", "[=Cl]", "[=F]", "[=H]", "[=I]",
        "[=NH1]", "[=N]", "[=O]", "[=P]", "[=Ring1]", "[=Ring2]", "[=Ring3]",
        "[=S]", "[Br]", "[Branch1]", "[Branch2]", "[Branch3]", "[C@@H1]",
        "[C@@]", "[C@H1]", "[C@]", "[C]", "[Cl]", "[F]", "[H]", "[I]", "[NH1]",
        "[N]", "[O]", "[P]", "[Ring1]", "[Ring2]", "[Ring3]", "[S]", "[\\Br]",
        "[\\C@@H1]", "[\\C@@]", "[\\C@H1]", "[\\C@]", "[\\C]", "[\\Cl]",
        "[\\F]", "[\\H]", "[\\I]", "[\\NH1]", "[\\N]", "[\\O]", "[\\P]",
        "[\\S]", "[nop]"
    ])
    return list(alphabet)


def test_random_selfies_decoder(trials, max_selfies_len, large_alphabet):
    """Tests that SELFIES that are generated by randomly stringing together
    symbols from the SELFIES alphabet are decoded into valid SMILES.
    """

    alphabet = tuple(large_alphabet)

    for _ in range(trials):

        # create random SELFIES and decode
        rand_len = random.randint(1, max_selfies_len)
        rand_selfies = "".join(random_choices(alphabet, k=rand_len))
        smiles = sf.decoder(rand_selfies)

        # check if SMILES is valid
        try:
            is_valid = MolFromSmiles(smiles, sanitize=True) is not None
        except Exception:
            is_valid = False

        err_msg = "SMILES: {}\n\t SELFIES: {}".format(smiles, rand_selfies)
        assert is_valid, err_msg


def test_nop_symbol_decoder(max_selfies_len, large_alphabet):
    """Tests that the '[nop]' symbol is always skipped over.
    """

    alphabet = list(large_alphabet)
    alphabet.remove("[nop]")

    for _ in range(100):

        # create random SELFIES with and without [nop]
        rand_len = random.randint(1, max_selfies_len)
        rand_mol = random_choices(alphabet, k=rand_len)
        rand_mol.extend(["[nop]"] * (max_selfies_len - rand_len))
        random.shuffle(rand_mol)

        with_nops = "".join(rand_mol)
        without_nops = with_nops.replace("[nop]", "")

        assert sf.decoder(with_nops) == sf.decoder(without_nops)


def test_get_semantic_constraints():
    constraints = sf.get_semantic_constraints()
    assert constraints is not sf.get_semantic_constraints()  # not alias
    assert "?" in constraints


def test_change_constraints_cache_clear():
    alphabet = sf.get_semantic_robust_alphabet()
    assert alphabet == sf.get_semantic_robust_alphabet()
    assert sf.decoder("[C][#C]") == "C#C"

    new_constraints = sf.get_semantic_constraints()
    new_constraints["C"] = 1
    sf.set_semantic_constraints(new_constraints)

    new_alphabet = sf.get_semantic_robust_alphabet()
    assert new_alphabet != alphabet
    assert sf.decoder("[C][#C]") == "CC"

    sf.set_semantic_constraints()  # re-set alphabet


def test_invalid_or_unsupported_smiles_encoder():
    malformed_smiles = [
        "",
        "(",
        "C(Cl)(Cl)CC[13C",
        "C(CCCOC",
        "C=(CCOC",
        "CCCC)",
        "C1CCCCC",
        "C(F)(F)(F)(F)(F)F",  # violates bond constraints
        "C=C1=CCCCCC1",  # violates bond constraints
        "CC*CC",  # uses wildcard
        "C$C",  # uses $ bond
        "S[As@TB1](F)(Cl)(Br)N",  # unrecognized chirality,
        "SOMETHINGWRONGHERE",
        "1243124124",
    ]

    for smiles in malformed_smiles:
        with pytest.raises(sf.EncoderError):
            sf.encoder(smiles)


def test_malformed_selfies_decoder():
    with pytest.raises(sf.DecoderError):
        sf.decoder("[O][=C][O][C][C][C][C][O][N][Branch2_3")


def random_choices(population, k):  # random.choices was new in Python v3.6
    return [random.choice(population) for _ in range(k)]


def test_decoder_attribution():
    sm, am = sf.decoder(
        "[C][N][C][Branch1][C][P][C][C][Ring1][=Branch1]", attribute=True)
    # check that P lined up
    for ta in am:
        if ta.token == 'P':
            for a in ta.attribution:
                if a.token == '[P]':
                    return
    raise ValueError('Failed to find P in attribution map')


def test_encoder_attribution():
    smiles = "C1([O-])C=CC=C1Cl"
    indices = [0, 3, 3, 3, 5, 7, 8, 10, None, None, 12]
    s, am = sf.encoder(smiles, attribute=True)
    for i, ta in enumerate(am):
        if ta.attribution:
            assert indices[i] == ta.attribution[0].index, \
                f'found {ta[1]}; should be {indices[i]}'
        if ta.token == '[Cl]':
            assert 'Cl' in [
                a.token for a in ta.attribution],\
                'Failed to find Cl in attribution map'