import unittest
from modeller import Alignment, Model
from mdt_test import MDTTest
import mdt
import mdt.features
import math
import os


class TableTests(MDTTest):

    restyps = "ACDEFGHIKLMNPQRSTVWY"

    def __mdt_restyp(self, a):
        """Convert a one-letter residue code into the MDT bin number"""
        ind = self.restyps.find(a)
        # Map to 'undefined' bin:
        if ind == -1:
            ind = 21
        return ind

    def test_version(self):
        """Check version number"""
        self.assertGreater(len(mdt.version), 0)
        self.assertGreater(len(mdt.__version__), 0)
        self.assertGreater(len(mdt.version_info), 0)

    def test_add(self):
        """Che
ck adding MDTs"""
        mlib = self.get_mdt_library()
        bins = mdt.uniform_bins(3, -1.0, 1.5)
        xray0 = mdt.features.XRayResolution(mlib, bins, protein=0)
        xray1 = mdt.features.XRayResolution(mlib, bins, protein=1)
        m1 = mdt.Table(mlib, features=xray0)
        for (n, val) in enumerate((1, 2, 3, 4)):
            m1[n] = val
        m2 = mdt.Table(mlib, features=xray0)
        for (n, val) in en
umerate((10, 20, 30, 40)):
            m2[n] = val
        m3 = m1 + m2
        m1 += m2
        self.assertMDTsEqual(m1, m3)
        for (n, val) in enumerate((11, 22, 33, 44)):
            self.assertEqual(m3[n], val)
        # Cannot add if numbers of features are different
        badmdt = mdt.Table(mlib, features=(xray0, xray1))
        self.assertRaises(ValueError, m1.__add__, badmdt)

        # Cannot add if feature types are different
        badmdt = mdt.Table(mlib, features=xray1)
        self.assertRaises(ValueError, m1.__add__, badmdt)

        # Cannot add if starts are different
        badmdt = m2.reshape(features=xray0, offset=1, shape=0)
        self.assertRaises(ValueError, m1.__add__, badmdt)

        # Cannot add if nbins are different
        badmdt = m2.reshape(features=xray0, offset=0, shape=-1)
        self.assertRaises(ValueError, m1.__add__, badmdt)

    def test_features_same_library(self):
        """Features must all come from the same Library"""
        mlib1 = self.get_mdt_library()
        mlib2 = self.get_mdt_library()
        restyp1 = mdt.features.ResidueType(mlib1, protein=0)
        restyp2 = mdt.features.ResidueType(mlib2, protein=1)
        _ = mdt.Table(mlib1, features=restyp1)
        _ = mdt.Table(mlib2, features=restyp2)
        for mlib in (mlib1, mlib2):
            self.assertRaises(ValueError, mdt.Table, mlib,
                              features=(restyp1, restyp2))

    def test_make_outrange(self):
        """Check handling of out-of-range features in mdt.make()"""
        class OutOfRangeFeature(object):
            def _get_ifeat(self, mlib):
                return 0
        mlib = self.get_mdt_library()
        f1 = OutOfRangeFeature()
        self.assertRaises(IndexError, mdt.Table, mlib, features=f1)

    def test_make_shape(self):
        """Test Table.make() with shape"""
        mlib = self.get_mdt_library()
        restyp = mdt.features.ResidueType(mlib)
        for shape in [[0, 0], [-22], [23]]:
            self.assertRaises(ValueError, mdt.Table, mlib, features=restyp,
                              shape=shape)
        for shape in [[], [0], [22]]:
            t = mdt.Table(mlib, features=restyp, shape=shape)
            self.assertEqual(t.shape, (22,))
        t = mdt.Table(mlib, features=restyp, shape=[10])
        self.assertEqual(t.shape, (10,))
        t = mdt.Table(mlib, features=restyp, shape=[-4])
        self.assertEqual(t.shape, (18,))

    def test_clear(self):
        """Check Table.clear()"""
        mlib = self.get_mdt_library()
        restyp = mdt.features.ResidueType(mlib)
        t = mdt.Table(mlib, features=restyp)
        t[0] = 5
        t.clear()
        self.assertEqual(t[0], 0)

    def test_bad_read(self):
        """Check error handling in Table.read()"""
        mlib = self.get_mdt_library()
        _ = mdt.features.ResidueType(mlib)
        feat = """
  #  FEATURE NBINS NAME
  1        1    22 Residue type of protein 0

"""

        with open('bad.mdt', 'w') as fh:
            fh.write("%-37s: garbage\n" % "Number of alignments")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write("%-37s: garbage\n" % "Sample size")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write("  #  FEATURE NBINS NAME\ngarbage\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "  # ISTART   IEND\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "  # ISTART   IEND\ngarbage\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "MDT TABLE START: garbage\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "MDT TABLE START:       1\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "MDT TABLE START:       1\ngarbage\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "MDT TABLE START:       1\n1.0\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        with open('bad.mdt', 'w') as fh:
            fh.write(feat + "MDT TABLE START:       1\n1.0\nf\n")
        self.assertRaises(mdt.FileFormatError, mdt.Table, mlib, file='bad.mdt')

        os.unlink('bad.mdt')

    def test_mdt_formats(self):
        """Make sure we can read and write MDT files"""
        env = self.get_environ()
        mlib = self.get_mdt_library()
        restyp0 = mdt.features.ResidueType(mlib, protein=0)
        restyp1 = mdt.features.ResidueType(mlib, protein=1)
        m = mdt.Table(mlib, features=(restyp0, restyp1))
        aln = Alignment(env)
        aln.append_sequence('AFVVTDNCIKCKYTDCVEVXPVDCFYEG')
        aln.append_sequence('CVEVCPVDCFYEGAFVVTDNCIKCKYTX')
        m.add_alignment(aln)
        m.write('test.mdt')
        self.assertRaises(IOError, m.write, '/does/not/exist/foo.mdt')
        m2 = m.copy()
        self.assertMDTsEqual(m, m2)
        m2 = mdt.Table(mlib, file='test.mdt')
        self.assertMDTsEqual(m, m2)
        m.write_hdf5('test.hdf5')
        m.write_hdf5('test.hdf5', gzip=True)
        m.write_hdf5('test.hdf5', gzip=4)
        self.assertRaises(ValueError, m.write_hdf5, 'test.hdf5',
                          gzip=True, chunk_size=(2, 3, 4))
        m.write_hdf5('test.hdf5', gzip=9, chunk_size=(4, 4))
        self.assertRaises(mdt.MDTError, m.write_hdf5,
                          '/does/not/exist/foo.hdf5')
        m2 = mdt.Table(mlib, file='test.hdf5')
        self.assertMDTsEqual(m, m2)
        m2 = mdt.Table(mlib)
        # Trying to read an HDF5 file in text format should fail gracefully:
        self.assertRaises(mdt.MDTError, m2.read, 'test.hdf5')
        # Same problem should occur starting with a non-empty MDT:
        m2 = mdt.Table(mlib, features=(restyp0, restyp1))
        self.assertRaises(mdt.MDTError, m2.read, 'test.hdf5')
        os.unlink('test.hdf5')

    def test_feature_check(self):
        """When rereading MDTs, features should be the same"""
        mlib = self.get_mdt_library()
        feat = mdt.features.AtomDistance(mlib, mdt.uniform_bins(10, 0, 1))
        m = mdt.Table(mlib, features=feat)
        m.write(file='test.mdt')
        m.write_hdf5(file='test.hdf5')
        # Different feature type, same bins should fail:
        mlib = self.get_mdt_library()
        feat = mdt.features.ResidueDistance(mlib, mdt.uniform_bins(10, 0, 1))
        for f in ('test.mdt', 'test.hdf5'):
            self.assertRaises(mdt.MDTError, mdt.Table, mlib