from __future__ import division, print_function

import os
import unittest
from uuid import uuid1
import numpy as np
from redrock import io
import redrock.test.util

class TestIO(unittest.TestCase):
    
    #- Create unique test filename in a subdirectory
    @classmethod
    def setUpClass(cls):
        cls.testfile = 'test-{uuid}.h5'.format(uuid=uuid1())

    #- Cleanup test files if they exist
    @classmethod
    def tearDownClass(cls):
        if os.path.exists(cls.testfile):
            os.remove(cls.testfile)
    
    def setUp(self):
        pass

    def test_endian(self):
        x1 = np.arange(5, dtype='>f')
        x2 = np.arange(5, dtype='<f')
        self.assertTrue(io.native_endian(x1).dtype.isnative)
        self.assertTrue(io.native_endian(x2).dtype.isnative)
        if x1.dtype.isnative:
            self.assertTrue(x1 is io.native_endian(x1))
        else:
            self.assertTrue(x2 is io.native_endian(x2))
                
    ### @unittest.skipIf('RR_TEMPLATE_DIR' not in os.environ, '$RR_TEMPLATE_DIR not set')
    def test_find_templates(self):
        templates = io.find_templates()
        self.assertTrue(len(templates) > 0)
        template_dir = os.path.dirname(templates[0])
        templates = io.find_templates(template_dir = template_dir)
        self.assertTrue(len(templates) > 0)
        
    ### @unittest.skipIf('RR_TEMPLATE_DIR' not in os.environ, '$RR_TEMPLATE_DIR not set')
    def test_read_templates(self):
        for template in io.read_templates():
            self.assertIn('wave', template.__dict__)
            self.assertIn('flux', template.__dict__)
            self.assertIn('type', template.__dict__)
            wave = template.wave
            flux = template.flux
            self.assertEqual(wave.shape[0], flux.shape[1])
            self.assertEqual(wave.ndim, 1)
            self.assertEqual(flux.ndim, 2)
        
    def test_zscan_io(self):
        t1 = redrock.test.util.get_target(0.2)
        t1.id = 111
        t2 = redrock.test.util.get_target(0.5)
        t2.id = 222
        template = redrock.test.util.get_template(subtype='BLAT')
        zscan1, zfit1 = redrock.zfind([t1,t2], [template,], ncpu=1)

        io.write_zscan(self.testfile, zscan1, zfit1)
        io.write_zscan(self.testfile, zscan1, zfit1, clobber=True)
        zscan2, zfit2 = io.read_zscan(self.testfile)

        self.assertEqual(zfit1.colnames, zfit2.colnames)
        self.assertTrue(np.all(zfit1 == zfit2))

        for targetid in zscan1:
            for spectype in zscan1[targetid]:
                for key in zscan1[targetid][spectype]:
                    d1 = zscan1[targetid][spectype][key]
                    d2 = zscan2[targetid][spectype][key]
                    self.assertTrue(np.all(d1==d2), 'data mismatch {}/{}/{}'.format(targetid, spectype, key))
                
if __name__ == '__main__':
    unittest.main()
