import faulthandler

import pathlib

import random

import pandas as pd

import pytest

from rdkit import Chem

import selfies as sf

faulthandler.enable()

TEST_SET_DIR = pathlib.Path(__file__).parent / "test_sets"

ERROR_LOG_DIR = pathlib.Path(__file__).parent / "error_logs"

ERROR_LOG_DIR.mkdir(exist_ok=True, parents=True)

datasets = list(TEST_SET_DIR.glob("**/*.csv"))

@pytest.mark.parametrize("test_path", datasets)

def test_roundtrip_translation(test_path, dataset_samples):

    """Tests SMILES -> SELFIES -> SMILES translation on various datasets.

    """

    # very relaxed constraints

    constraints = sf.get_preset_constraints("hypervalent")

    constraints.update({"P": 7, "P-1": 8, "P+1": 6, "?": 12})

    sf.set_semantic_constraints(constraints)

    error_path = ERROR_LOG_DIR / "{}.csv".format(test_path.stem)

    with open(error_path, "w+") as error_log:

        error_log.write("In, Out\n")

    error_data = []

    error_found = False

    n_lines = sum(1 for _ in open(test_path)) - 1

    n_keep = dataset_samples if (0 < dataset_samples <= n_lines) else n_lines

    skip = random.sample(range(1, n_lines + 1), n_lines - n_keep)

    reader = pd.read_csv(test_path, chunksize=10000, header=0, skiprows=skip)

    for chunk in reader:

        for in_smiles in chunk["smiles"]:

            in_smiles = in_smiles.strip()

            mol = Chem.MolFromSmiles(in_smiles, sanitize=True)

            if (mol is None) or ("*" in in_smiles):

                continue

            try:

                selfies = sf.encoder(in_smiles, strict=True)

                out_smiles = sf.decoder(selfies)

            except (sf.EncoderError, sf.DecoderError):

                error_data.append((in_smiles, ""))

                continue

            if not is_same_mol(in_smiles, out_smiles):

                error_data.append((in_smiles, out_smiles))

        with open(error_path, "a") as error_log:

            for entry in error_data:

                error_log.write(",".join(entry) + "\n")

        error_found = error_found or error_data

        error_data = []

    sf.set_semantic_constraints()  # restore constraints

    assert not error_found

def is_same_mol(smiles1, smiles2):

    try:

        can_smiles1 = Chem.CanonSmiles(smiles1)

        can_smiles2 = Chem.CanonSmiles(smiles2)

        return can_smiles1 == can_smiles2

    except Exception:

        return False