"""Script for testing selfies against large datasets.

"""

import argparse

import pathlib

import pandas as pd

from rdkit import Chem

from tqdm import tqdm

import selfies as sf

parser = argparse.ArgumentParser()

parser.add_argument("--data_path", type=str, default="version.smi.gz")

parser.add_argument("--col_name", type=str, default="isosmiles")

parser.add_argument("--sep", type=str, default=r"\s+")

parser.add_argument("--start_from", type=int, default=0)

args = parser.parse_args()

TEST_DIR = pathlib.Path(__file__).parent

TEST_SET_PATH = TEST_DIR / "test_sets" / args.data_path

ERROR_LOG_DIR = TEST_DIR / "error_logs"

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

def make_reader():

    return pd.read_csv(TEST_SET_PATH, sep=args.sep, chunksize=10000)

def roundtrip_translation():

    sf.set_semantic_constraints("hypervalent")

    n_entries = 0

    for chunk in make_reader():

        n_entries += len(chunk)

    pbar = tqdm(total=n_entries)

    reader = make_reader()

    error_log = open(ERROR_LOG_DIR / f"{TEST_SET_PATH.stem}.txt", "a+")

    curr_idx = 0

    for chunk_idx, chunk in enumerate(reader):

        for in_smiles in chunk[args.col_name]:

            pbar.update(1)

            curr_idx += 1

            if curr_idx < args.start_from:

                continue

            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_log.write(in_smiles + "\n")

                tqdm.write(in_smiles)

                continue

            if not is_same_mol(in_smiles, out_smiles):

                error_log.write(in_smiles + "\n")

                tqdm.write(in_smiles)

    error_log.close()

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

if __name__ == "__main__":

    roundtrip_translation()