import pytest


@pytest.mark.torch
def test_smiles_call():
    """Test __call__ method for the featurizer, which is inherited from HuggingFace's RobertaTokenizerFast"""
    from deepchem.feat.roberta_tokenizer import RobertaFeaturizer
    smiles = ["Cn1c(=O)c2c(ncn2C)n(C)c1=O", "CC(=O)N1CN(C(C)=O)C(O)C1O"]
    long_molecule_smiles = [
        "CCCCCCCCCCCCCCCCCCCC(=O)OCCCNC(=O)c1ccccc1SSc1ccccc1C(=O)NCCCOC(=O)CCCCCCCCCCCCCCCCCCC"
    ]
    featurizer = RobertaFeaturizer.from_pretrained(
        "seyonec/SMILES_tokenized_PubChem_shard00_160k")
    embedding = featurizer(smiles, add_special_tokens=True, truncation=True)
    embedding_long = featurizer(long_molecule_smiles * 2,
                                add_special_tokens=True,
                                truncation=True)
    for emb in [embedding, embedding_long]:
        assert 'input_ids' in emb.keys() and 'attention_mask' in emb.keys()
        assert len(emb['input_ids']) == 2 and len(emb['attention_mask']) == 2


@pytest.mark.torch
def test_smiles_featurize():
    """Test the .featurize method, which will convert the dictionary output to an array

    Checks that all SMILES are featurized and that each featurization
    contains input_ids and attention_mask
    """
    from deepchem.feat.roberta_tokenizer import RobertaFeaturizer
    smiles = ["Cn1c(=O)c2c(ncn2C)n(C)c1=O", "CC(=O)N1CN(C(C)=O)C(O)C1O"]
    long_molecule_smiles = [
        "CCCCCCCCCCCCCCCCCCCC(=O)OCCCNC(=O)c1ccccc1SSc1ccccc1C(=O)NCCCOC(=O)CCCCCCCCCCCCCCCCCCC"
    ]
    featurizer = RobertaFeaturizer.from_pretrained(
        "seyonec/SMILES_tokenized_PubChem_shard00_160k")
    feats = featurizer.featurize(smiles,
                                 add_special_tokens=True,
                                 truncation=True)
    assert (len(feats) == 2)
    assert (all([len(f) == 2 for f in feats]))
    long_feat = featurizer.featurize(long_molecule_smiles,
                                     add_special_tokens=True,
                                     truncation=True)
    assert (len(long_feat) == 1)
    assert (len(long_feat[0] == 2))
