"""
Tests for serial.py.
"""
import cPickle
from cStringIO import StringIO
import gzip
import shutil
import tempfile
import unittest

from rdkit import Chem
from rdkit.Chem import AllChem

from vs_utils.utils.rdkit_utils import conformers, serial


class TestMolIO(unittest.TestCase):
    """
    Base test class for molecule I/O.
    """
    def setUp(self):
        """
        Write SDF and SMILES molecules to temporary files.
        """
        self.temp_dir = tempfile.mkdtemp()

        # aspirin
        self.aspirin = self._get_mol_from_smiles('CC(=O)OC1=CC=CC=C1C(=O)O',
       
                    'aspirin')
        self.aspirin_h = Chem.AddHs(self.aspirin)
        self.aspirin_sodium = self._get_mol_from_smiles(
            'CC(=O)OC1=CC=CC=C1C(=O)[O-].[Na+]', 'aspirin sodium')

        # levalbuterol (chiral)
        self.levalbuterol = self._get_mol_from_smiles(
            'CC(C)(C)NC[C@@H](C1=CC(=C(C=C1)O)CO)O', 'levalbuterol')
        self.levalbuterol_hcl = self._get_mol_from_smiles(
            'CC(C)(C)NC[C@@H](C1=CC(=C(C=C1)O)CO)O.Cl',
            'levalbuterol hydrochloride')

        self.ref_mols = [self.aspirin, self.levalbuterol]
        self.reader = serial.MolReader(compute_2d_coords=False)

    def _get_mol_from_smiles(self, smiles, name=None):
        """
        Construct a molecule from a SMILES string.

        Molecules loaded from SMILES strings have zero conformers, but
        molecules loaded from SDF blocks are treated as 3D and have one
        conformer even if coordinates are not set. This method dumps the
        molecule to SDF and loads it again to obtain a molecule with one
        conformer.

        Parameters
        ----------
        smiles : str
            SMILES string.
        name : str, optional
            Molecule name.
        """
        mol = Chem.MolFromSmiles(smiles)
        if name is not None:
            mol.SetProp('_Name', name)
        AllChem.Compute2DCoords(mol)  # required to preserve stereo
        sdf = Chem.MolToMolBlock(mol, includeStereo=True)
        mol_with_conf = Chem.MolFromMolBlock(sdf)
        return mol_with_conf

    def tearDown(self):
        """
        Clean up temporary files.
        """
        shutil.rmtree(self.temp_dir)

    def test_guess_mol_format(self):
        """
        Test MolIO.guess_mol_format.
        """
        mol_formats = {
            'pkl': ['test.pkl', 'test.pkl.gz', 'test.test.pkl',
                    'test.test.pkl.gz'],
            'sdf': ['test.sdf', 'test.sdf.gz', 'test.test.sdf',
                    'test.test.sdf.gz'],
            'smi': ['test.smi', 'test.smi.gz', 'test.can', 'test.can.gz',
                    'test.ism', 'test.ism.gz', 'test.test.smi',
                    'test.test.smi.gz']
        }
        for mol_format in mol_formats.keys():
            for filename in mol_formats[mol_format]:
                assert self.reader.guess_mol_format(filename) == mol_format

    def test_close_context(self):
        """
        Make sure MolIO closes files it opened.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf', dir=self.temp_dir)
        self.reader.open(filename)
        self.reader.close()
        assert self.reader.f.closed

        # also test the context manager
        with self.reader.open(filename):
            pass
        assert self.reader.f.closed

    def test_not_close_other(self):
        """
        Make sure MolIO doesn't close files it didn't open.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf', dir=self.temp_dir)
        with open(filename) as f:
            reader = serial.MolReader(f, mol_format='sdf')
            reader.close()
            assert not f.closed

        # also test the context manager
        with open(filename) as g:
            with serial.MolReader(g, mol_format='sdf'):
                pass
            assert not g.closed


class TestMolReader(TestMolIO):
    """
    Test MolReader.
    """
    def test_read_sdf(self):
        """
        Read an SDF file.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf', dir=self.temp_dir)
        with open(filename, 'wb') as f:
            f.write(Chem.MolToMolBlock(self.aspirin))
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_sdf_gz(self):
        """
        Read a compressed SDF file.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf.gz', dir=self.temp_dir)
        with gzip.open(filename, 'wb') as f:
            f.write(Chem.MolToMolBlock(self.aspirin))
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_smi(self):
        """
        Read a SMILES file.
        """
        self.aspirin.RemoveAllConformers()  # SMILES are read without confs
        _, filename = tempfile.mkstemp(suffix='.smi', dir=self.temp_dir)
        with open(filename, 'wb') as f:
            f.write(Chem.MolToSmiles(self.aspirin))
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_smi_title(self):
        """
        Read a SMILES file with molecule titles.
        """
        self.aspirin.RemoveAllConformers()  # SMILES are read without confs
        _, filename = tempfile.mkstemp(suffix='.smi', dir=self.temp_dir)
        with open(filename, 'wb') as f:
            f.write('{}\t{}'.format(Chem.MolToSmiles(self.aspirin), 'aspirin'))
        self.reader.open(filename)
        mols = self.reader.get_mols()
        mol = mols.next()
        assert mol.ToBinary() == self.aspirin.ToBinary()
        assert mol.GetProp('_Name') == self.aspirin.GetProp('_Name')

    def test_read_smi_gz(self):
        """
        Read a compressed SMILES file.
        """
        self.aspirin.RemoveAllConformers()  # SMILES are read without confs
        _, filename = tempfile.mkstemp(suffix='.smi.gz', dir=self.temp_dir)
        with gzip.open(filename, 'wb') as f:
            f.write(Chem.MolToSmiles(self.aspirin))
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_pickle(self):
        """
        Read from a pickle.
        """
        _, filename = tempfile.mkstemp(suffix='.pkl', dir=self.temp_dir)
        with open(filename, 'wb') as f:
            cPickle.dump([self.aspirin], f, cPickle.HIGHEST_PROTOCOL)
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_pickle_gz(self):
        """
        Read from a compressed pickle.
        """
        _, filename = tempfile.mkstemp(suffix='.pkl.gz', dir=self.temp_dir)
        with gzip.open(filename, 'wb') as f:
            cPickle.dump([self.aspirin], f, cPickle.HIGHEST_PROTOCOL)
        self.reader.open(filename)
        mols = self.reader.get_mols()
        assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_file_like(self):
        """
        Read from a file-like object.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf', dir=self.temp_dir)
        with open(filename, 'wb') as f:
            f.write(Chem.MolToMolBlock(self.aspirin))
        with open(filename) as f:
            reader = serial.MolReader(f, mol_format='sdf')
            mols = reader.get_mols()
            assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_compressed_file_like(self):
        """
        Read from a file-like object using gzip.
        """
        _, filename = tempfile.mkstemp(suffix='.sdf.gz', dir=self.temp_dir)
        with gzip.open(filename, 'wb') as f:
            f.write(Chem.MolToMolBlock(self.aspirin))
        with gzip.open(filename) as f:
            reader = serial.MolReader(f, mol_format='sdf')
            mols = reader.get_mols()
            assert mols.next().ToBinary() == self.aspirin.ToBinary()

    def test_read_multiple_sdf(self):
        """
        Read a 