#!/usr/env/bin python
# -*- coding: utf-8 -*-
"""
unit testing for MolVS steps for PubChem Substances
tests include
molvs.standardize_smiles
Standardizer().normalize
Standardizer().disconnect_metals
Standardizer().reionize
molvs.standardize.canonicalize_tautomer_smiles
molvs.validate.Validator()
Standardizer().fragment_parent
"""
from collections import namedtuple
import unittest
import gzip
import os.path

from rdkit import Chem
from rdkit import RDConfig
import molvs
from molvs import Standardizer, validate

doLong = False
TestData = namedtuple('TestData', 'lineNo,smiles,mol,expected')

class TestCase(unittest.TestCase):
    dataPCS_standardize_smiles100k = os.path.join(RDConfig.RDBaseDir, 'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_standardize_sm.csv.gz')
    dataPCS_standardize_smiles1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_standardize_sm.csv.gz')
    dataPCS_nomralized1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_normalized.csv.gz')
    dataPCS_nomralized100k = os.path.join(RDConfig.RDBaseDir, 'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_normalized.csv.gz')
    dataPCS_metal100k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_metals.csv.gz')
    dataPCS_metal1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_metals.csv.gz')
    dataPCS_reionize100k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_reionize.csv.gz')
    dataPCS_reionize1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_reionize.csv.gz')
    dataPCS_fragment100k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_fragment.csv.gz')
    dataPCS_fragmnet1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_fragment.csv.gz')
    dataPCS_tautomer100k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_tautomer.csv.gz')
    dataPCS_tautomer1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_tautomer.csv.gz')
    dataPCS_validate100k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '100kPCS_validate.csv.gz')
    dataPCS_validate1k = os.path.join(RDConfig.RDBaseDir,'rdkit', 'Chem', 'MolStandardize', 'test_data', '1kPCS_validate.csv.gz')

    @classmethod
    def readPCSdata(cls, datafile):
        for data in readPCStestData(datafile):
            yield data

    def testStandardizeSmShort(self):
        for data in self.readPCSdata(self.dataPCS_standardize_smiles1k):
            ss = molvs.standardize_smiles(data.smiles)
            self.assertEqual(ss, data.expected)

    def testStandardizeSmLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_standardize_smiles100k):
            try:
                ss = molvs.standardize_smiles(data.smiles)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS standardization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(ss, data.expected)

    def testNormalizeShort(self):
        for data in self.readPCSdata(self.dataPCS_nomralized1k):
            n = Standardizer()
            nm = n.normalize(data.mol)
            ns = Chem.MolToSmiles(nm)
            self.assertEqual(ns, data.expected)

    def testNormalizeLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_nomralized100k):
            try:
                n = Standardizer()
                nm = n.normalize(data.mol)
                ns = Chem.MolToSmiles(nm)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(ns, data.expected)

    def testMetalShort(self):
        for data in self.readPCSdata(self.dataPCS_metal1k):
            n = Standardizer()
            nm = n.disconnect_metals(data.mol)
            ns = Chem.MolToSmiles(nm)
            self.assertEqual(ns, data.expected)

    def testMetalLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_metal100k):
            try:
                n = Standardizer()
                nm = n.disconnect_metals(data.mol)
                ns = Chem.MolToSmiles(nm)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(ns, data.expected)

    def testReionizeShort(self):
        for data in self.readPCSdata(self.dataPCS_reionize1k):
            n = Standardizer()
            nm = n.reionize(data.mol)
            ns = Chem.MolToSmiles(nm)
            self.assertEqual(ns, data.expected)

    def testReionizeLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_reionize100k):
            try:
                n = Standardizer()
                nm = n.reionize(data.mol)
                ns = Chem.MolToSmiles(nm)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(ns, data.expected)

    def testFragmentShort(self):
        for data in self.readPCSdata(self.dataPCS_fragmnet1k):
            s = Standardizer()
            frag = s.fragment_parent(data.mol)
            ns = Chem.MolToSmiles(frag)
            self.assertEqual(ns, data.expected)

    def testFragmentLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_fragment100k):
            try:
                s = Standardizer()
                frag = s.fragment_parent(data.mol)
                ns = Chem.MolToSmiles(frag)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(ns, data.expected)

    def testTautomerShort(self):
        for data in self.readPCSdata(self.dataPCS_tautomer1k):
            canon_taut = molvs.standardize.canonicalize_tautomer_smiles(data.smiles)
            self.assertEqual(canon_taut, data.expected)

    def testTautomerLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for data in self.readPCSdata(self.dataPCS_tautomer100k):
            try:
                canon_taut = molvs.standardize.canonicalize_tautomer_smiles(data.smiles)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(canon_taut, data.expected)

    def testValidateShort(self):
        for data in self.readPCSdata(self.dataPCS_validate1k):
            v = validate.Validator()
            stdout = v(data.mol)
            if stdout:
                stdout = ' '.join(stdout)
            else:
                stdout = str(stdout)
            self.assertEqual(stdout, data.expected)

    def testValidateLong(self):
        if not doLong:
            raise unittest.SkipTest('long test')
        for i, data in enumerate(self.readPCSdata(self.dataPCS_validate100k)):
            try:
                v = validate.Validator()
                stdout = v(data.mol)
                if stdout:
                    stdout = ','.join(stdout)
                else:
                    stdout = str(stdout)
                # print(i, stdout)
            except Exception:
                raise AssertionError(
                    'Line {0.lineNo}: MolVS normalization failed for SMILES {0.smiles}'.format(data))
            self.assertEqual(stdout, data.expected)

def readPCStestData(filename):
    """ Read test data for MolVS from file """
    with gzip.open(filename, 'rt') as f:
        for lineNo, line in enumerate(f, 1):
            if line[0] == '#':
                continue
            smiles = line.strip().split(',')[0]
            expected = ",".join(line.strip().split(',')[1:])
            mol = Chem.MolFromSmiles(smiles)
            if not mol:
                raise AssertionError('molecule construction failed on line %d' % lineNo)
            yield TestData(lineNo, smiles, mol, expected)

if __name__ == '__main__':
    import argparse
    import sys

    parser = argparse.ArgumentParser()
    parser.add_argument('-l', default=False, action='store_true', dest='doLong')
    args = parser.parse_args()
    doLong = args.doLong

    # Remove the -l flag if present so that it isn't interpreted by unittest.main()
    if '-l' in sys.argv:
        sys.argv.remove('-l')
    unittest.main()
