import iopro
import unittest
import warnings
import numpy


try:
    import pymongo
    pymongo_installed = True
except ImportError:
    pymongo_installed = False


class TestMongoAdapter(unittest.TestCase):

    hostname = 'localhost'
    port = 27017
    db_name = 'MongoAdapter_tests'


    def test_ints(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port, TestMongoAdapter.db_name, 'test_ints')
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))
        sorted = numpy.sort(array)
        for i in range(1000):
            self.assertEqual(sorted[i].item(), (i*5, i*5+1, i*5+2, i*5+3, i*5+4))


    def test_floats(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port, TestMongoAdapter.db_name, 'test_floats')
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('f8,f8,f8,f8,f8'))
        sorted = numpy.sort(array)
        for i in range(1000):
            self.assertEqual(sorted[i].item(), (i*5, i*5+1, i*5+2, i*5+3, i*5+4))


    def test_type_inference(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port, TestMongoAdapter.db_name, 'test_type_inference')
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('f8,O,u8,u8,u8'))
        sorted = numpy.sort(array)
        for i in range(999):
            self.assertEqual(sorted[i].item(), (float(i*5), str(i*5+1), i*5+2, i*5+3, i*5+4))
        self.assertEqual(sorted[-1].item(), (float(4995.5), '4996', 4997, 4998, 4999))


    def test_field_subset(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port, TestMongoAdapter.db_name, 'test_ints')

        array = adapter['f0'][:]
        array.sort()
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype([('f0', 'u8')]))
        for i in range(1000):
            self.assertEqual(array[i].item(), (i*5,))

        array = adapter[['f0', 'f4']][:]
        array.sort()
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype([('f0', 'u8'),('f4', 'u8')]))
        for i in range(1000):
            self.assertEqual(array[i].item(), (i*5, i*5+4))

        array = adapter[('f0', 'f1', 'f2', 'f3', 'f4')][:]
        array.sort()
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))
        for i in range(1000):
            self.assertEqual(array[i].item(), (i*5, i*5+1, i*5+2, i*5+3, i*5+4))


    def test_array_slicing(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port, TestMongoAdapter.db_name, 'test_ints')
        control = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]

        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][0:1000]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))
        for i in range(1000):
            self.assertEqual(array[i].item(), control[i].item())
       
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][-1]
        self.assertEqual(array.size, 1)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))
        self.assertEqual(array[0].item(), control[-1].item())
 
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][::2]
        '''self.assertEqual(array.size, 500)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))
        for i in range(0, 1000, 2):
            self.assertEqual(array[i].item(), control[i].item())'''


    @unittest.skip('test_converter')
    def test_converter(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port,
            TestMongoAdapter.db_name, 'test_ints')
        adapter.set_converter('f0', lambda x: x * 2, use_numba=False)
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        for i in range(1000):
            self.assertEqual(array[i], (i*5*2, i*5+1, i*5+2, i*5+3, i*5+4))

        adapter.set_converter('f0', lambda x: x * 2, use_numba=True)
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        for i in range(1000):
            self.assertEqual(array[i], (i*5*2, i*5+1, i*5+2, i*5+3, i*5+4))


    @unittest.skip('test_set_field_types')
    def test_set_field_types(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port,
            TestMongoAdapter.db_name, 'test_ints', infer_types=False)
        adapter.set_field_types({'f0':'f8', 'f4':'O'})
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('f8,S10,u8,u8,O'))
        for i in range(1000):
            self.assertEqual(array[i], (float(i*5), i*5+1, i*5+2, i*5+3, str(i*5+4)))


    @unittest.skip('test_missing_fill_values')
    def test_missing_fill_values(self):
        adapter = iopro.MongoAdapter(TestMongoAdapter.hostname, TestMongoAdapter.port,
            TestMongoAdapter.db_name, 'test_missing_ints', infer_types=False)

        # Force f1 type to integer
        adapter.set_field_types({'f1':'u8'})

        adapter.set_missing_values({0:['NA']})
        adapter.set_fill_values({0:99, 1:999, 4:9999}, loose=True)
        
        array = adapter[['f0', 'f1', 'f2', 'f3', 'f4']][:]
        self.assertEqual(array.size, 1000)
        self.assertEqual(array.dtype, numpy.dtype('u8,u8,u8,u8,u8'))

        # Records 0,5,10... contains 'NA' and 'NaN' values in f0
        # that should convert to 99.
        # Records 1,6,11... contains values in f1 that can't be converted to int;
        # these should convert to 999.
        # Records 4,9,14... contains missing values in f4
        # that should convert to 9999.
        for i in range(1000):
            if i % 5 == 0:
                self.assertEqual(array[i], (99, i*5+1, i*5+2, i*5+3, i*5+4))
            elif i % 5 == 1:
                self.assertEqual(array[i], (i*5, 999, i*5+2, i*5+3, i*5+4))
            elif i % 5 == 4:
                self.assertEqual(array[i], (i*5, i*5+1, i*5+2, i*5+3, 9999))
            else:
                self.assertEqual(array[i], (i*5, i*5+1, i*5+2, i*5+3, i*5+4))


def run(verbosity=2, hostname='localhost', port=27017):

    TestMongoAdapter.hostname = hostname
    TestMongoAdapter.port = port

    if not pymongo_installed:
        warnings.warn('Could not run Mongo tests: pymongo is not installed.')
        return

    try:
        mongo_conn = pymongo.MongoClient(TestMongoAdapter.hostname, TestMongoAdapter.port)
    except:
        warnings.warn('Could not run Mongo tests: could not connect to Mongo database.')
        return

    mongo_db = mongo_conn['MongoAdapter_tests']

    col = mongo_db['test_ints']
    if col.count() > 0:
        col.remove()
    for i in range(1000):
        record = {'f0': i * 5,
                  'f1': i * 5 + 1,
                  'f2': i * 5 + 2,
                  'f3': i * 5 + 3,
                  'f4': i * 5 + 4}
        col.insert(record)

    col = mongo_db['test_floats']
    if col.count() > 0:
        col.remove()
    for i in range(1000):
        record = {'f0': float(i * 5),
                  'f1': float(i * 5 + 1),
                  'f2': float(i * 5 + 2),
                  'f3': float(i * 5 + 3),
                  'f4': float(i * 5 + 4)}
        col.insert(record)

    col = mongo_db['test_strings']
    if col.count() > 0:
        col.remove()
    for i in range(1000):
        record = {'f0': str(i * 5),
                  'f1': str(i * 5 + 1),
                  'f2': str(i * 5 + 2),
                  'f3': str(i * 5 + 3),
                  'f4': str(i * 5 + 4)}
        col.insert(record)

    col = mongo_db['test_missing_ints']
    if col.count() > 0:
        col.remove()
    for i in range(1000):
        if i % 5 == 0:
            record = {'f0': 'NA',
                      'f1': i * 5 + 1,
                      'f2': i * 5 + 2,
                      'f3': i * 5 + 3,
                      'f4': i * 5 + 4}
        elif i % 5 == 1:
            record = {'f0': i * 5,
                      'f1': 'xxx',
                      'f2': i * 5 + 2,
                      'f3': i * 5 + 3,
                      'f4': i * 5 + 4}
        elif i % 5 == 4:
            record = {'f0': i * 5,
                      'f1': i * 5 + 1,
                      'f2': i * 5 + 2,
                      'f3': i * 5 + 3}
        else:
            record = {'f0': i * 5,
                      'f1': i * 5 + 1,
                      'f2': i * 5 + 2,
                      'f3': i * 5 + 3,
                      'f4': i * 5 + 4}
        col.insert(record)

    col = mongo_db['test_type_inference']
    if col.count() > 0:
        col.remove()
    for i in range(999):
        record = {'f0': i * 5,
                  'f1': i * 5 + 1,
                  'f2': i * 5 + 2,
                  'f3': i * 5 + 3,
                  'f4': i * 5 + 4}
        col.insert(record)
    col.insert({'f0':4995.5, 'f1':'4996', 'f2':4997, 'f3':4998, 'f4':4999})
       
    suite = unittest.TestLoader().loadTestsFromTestCase(TestMongoAdapter)
    return unittest.TextTestRunner(verbosity=verbosity).run(suite)


if __name__ == '__main__':
    run()
