"""
Tests for cotrans functions.

These tests include the function in the following files:
    utilites/contrans.py
    utilities/cotrans_lib.py
    utilites/igrf.py
    utilites/j2000.py
    themis/cotrans/dsl2gse.py

"""
import unittest
import pyspedas
import logging
from pyspedas.themis.cotrans.dsl2gse import dsl2gse
from pyspedas.cotrans.cotrans import cotrans
from pyspedas.cotrans.fac_matrix_make import fac_matrix_make
from pytplot import get_data, store_data, del_data
from pyspedas import cotrans_get_coord, cotrans_set_coord


class CotransTestCases(unittest.TestCase):
    """Tests for cotrans."""
    def test_fac_matrix_make(self):
        doesntexist = fac_matrix_make('doesnt_exist')

    def test_get_set_coord_wrappers(self):
        """ Test for cotrans_set_coord/cotrans_get_coord wrappers """
        del_data()
        doesntexist = cotrans_get_coord('test_coord')
        self.assertTrue(doesntexist is None)
        store_data('test_coord', data={'x': [1, 2, 3, 4, 5], 'y': [1, 1, 1, 1, 1]})
        cotrans(name_in='test_coord', coord_out="geo")
        before = cotrans_get_coord('test_coord')
        self.assertTrue(before is None)
        setcoord = cotrans_set_coord('test_coord', 'GSE')
        self.assertTrue(setcoord)
        after = cotrans_get_coord('test_coord')
        self.assertTrue(after == 'GSE')
        md = get_data('test_coord',metadata=True)
        md['data_att']['units'] = 'km'
        setcoord = cotrans_set_coord('test_coord', 'GSM')
        self.assertTrue(setcoord)
        md_after = get_data('test_coord',metadata=True)
        after = cotrans_get_coord('test_coord')
        self.assertTrue(after == 'GSM')
        self.assertTrue(md_after['data_att']['units'] == 'km')
        setcoord = cotrans_set_coord('doesnt_exist', 'GSM')

    def test_get_set_coords(self):
        """ Test for pytplot.set_coords/get_coords """
        from pytplot import set_coords,get_coords

        del_data()
        doesntexist = get_coords('test_coord')
        self.assertTrue(doesntexist is None)
        store_data('test_coord', data={'x': [1, 2, 3, 4, 5], 'y': [1, 1, 1, 1, 1]})
        cotrans(name_in='test_coord', coord_out="geo")
        before = get_coords('test_coord')
        self.assertTrue(before is None)
        setcoord = set_coords('test_coord', 'GSE')
        self.assertTrue(setcoord)
        after = get_coords('test_coord')
        self.assertTrue(after == 'GSE')
        md = get_data('test_coord',metadata=True)
        md['data_att']['units'] = 'km'
        setcoord = set_coords('test_coord', 'GSM')
        self.assertTrue(setcoord)
        md_after = get_data('test_coord',metadata=True)
        after = get_coords('test_coord')
        self.assertTrue(after == 'GSM')
        self.assertTrue(md_after['data_att']['units'] == 'km')
        setcoord = set_coords('doesnt_exist', 'GSM')

    def test_get_set_units(self):
        """ Test for pytplot.set_coords/get_coords """
        from pytplot import set_units,get_units, set_coords, get_coords

        del_data()
        doesntexist = get_units('test_units')
        self.assertTrue(doesntexist is None)
        store_data('test_units', data={'x': [1, 2, 3, 4, 5], 'y': [1, 1, 1, 1, 1]})
        before = get_units('test_units')
        self.assertTrue(before is None)
        setunits = set_units('test_units', 'Km')
        self.assertTrue(setunits)
        after = get_units('test_units')
        self.assertTrue(after == 'Km')
        set_coords('test_units','GEO')
        setunits = set_units('test_units', 'mm')
        self.assertTrue(setunits)
        coords_after=get_coords('test_units')
        units_after=get_units('test_units')
        self.assertTrue(coords_after == 'GEO')
        self.assertTrue(units_after == 'mm')

    def test_dsl2gse(self):
        """Test themis.cotrans.dsl2gse."""
        del_data()
        # Try with missing variables. It should exit without problems.
        dsl2gse('tha_fgl_dsl','tha_fgl_gse')
        # Now load the needed variables.
        time_range = ['2017-03-23 00:00:00', '2017-03-23 23:59:59']
        pyspedas.themis.state(probe='a', trange=time_range,
                              get_support_data=True,
                              varnames=['tha_spinras', 'tha_spindec'])
        pyspedas.themis.fgm(probe='a', trange=time_range,
                            varnames=['tha_fgl_dsl'])

        fac_matrix_make('tha_fgl_dsl')

        dsl2gse('tha_fgl_dsl', 'tha_fgl_gse')

        t, d = get_data('tha_fgl_gse')
        # Now test the inverse.
        dsl2gse('tha_fgl_dsl', 'tha_fgl_gse',
                isgsetodsl=True)

        self.assertTrue(abs(d[0].tolist()[0]-15.905078404701147) <= 1e-6)
        self.assertTrue(abs(d[0].tolist()[1]--13.962618931740064) <= 1e-6)
        self.assertTrue(abs(d[0].tolist()[2]-16.392516225582813) <= 1e-6)

        self.assertTrue(abs(d[50000].tolist()[0]-16.079111468932435) <= 1e-6)
        self.assertTrue(abs(d[50000].tolist()[1]--18.858874541698583) <= 1e-6)
        self.assertTrue(abs(d[50000].tolist()[2]-14.75796300561617) <= 1e-6)

    def test_cotrans(self):
        """Test cotrans.py GEI->GEO with Themis data."""
        del_data()
        trange = ['2010-02-25/00:00:00', '2010-02-25/23:59:59']
        probe = 'a'
        name_in = "tha_pos"
        name_out = "tha_pos_new_geo"
        pyspedas.themis.state(probe=probe, trange=trange,
                              time_clip=True, varnames=[name_in])
        cotrans(name_in=name_in, name_out=name_out,
                coord_in="gei", coord_out="geo")
        din = get_data(name_in)
        in_len = len(din[0])
        dout = get_data(name_out)
        out_len = len(dout[0])
        cotrans(name_in=name_in, coord_in="gei", coord_out="geo")
        self.assertTrue(out_len == in_len)

    def test_cotrans_coord_mismatch(self):
        """Test that cotrans rejects a request where in_coord does not match the system from the variable metadata."""
        del_data()
        trange = ['2010-02-25/00:00:00', '2010-02-25/23:59:59']
        probe = 'a'
        name_in = "tha_pos"
        name_out = "tha_pos_new_geo"
        pyspedas.themis.state(probe=probe, trange=trange,
                              time_clip=True, varnames=[name_in])
        # Metadata coordinate system is GEI, but requesting GSM->GEO transform.  This should generate an error message
        # and return failure.
        result = cotrans(name_in=name_in, name_out=name_out,
                coord_in="gsm", coord_out="geo")
        self.assertTrue(result == 0)

    def test_cotrans_igrf(self):
        """Test GSE->GSM and IGRF."""
        del_data()
        d = [[245.0, -102.0, 251.0], [775.0, 10.0, -10],
             [121.0, 545.0, -1.0], [304.65, -205.3, 856.1],
             [464.34, -561.55, -356.22]]
        t = [1577112800, 1577308800, 1577598800, 1577608800, 1577998800]
        name_out = "tname_out"
        cotrans(name_out=name_out, time_in=t, data_in=d,
                coord_in="gse", coord_out="gsm")
        in_len = len(t)
        dout = get_data(name_out)
        out_len = len(dout[0])
        gsm = dout[1][1]
        res = [775.0, 11.70325822, -7.93937951]
        self.assertTrue(out_len == in_len)
        self.assertTrue(abs(gsm[0]-res[0]) <= 1e-6)
        self.assertTrue(abs(gsm[1]-res[1]) <= 1e-6)
        self.assertTrue(abs(gsm[2]-res[2]) <= 1e-6)

    def test_cotrans_j2000(self):
        """Test GEI->J2000 and J2000 params."""
        del_data()
        d = [[245.0, -102.0, 251.0], [775.0, 10.0, -10],
             [121.0, 545.0, -1.0], [304.65, -205.3, 856.1],
             [464.34, -561.55, -356.22]]
        t = [1577112800, 1577308800, 1577598800, 1577608800, 1577998800]
        name_out = "tname_out"
        cotrans(name_out=name_out, time_in=t, data_in=d,
                coord_in="gei", coord_out="j2000")
        in_len = len(t)
        dout = get_data(name_out)
        out_len = len(dout[0])
        j2000 = dout[1][1]
        res = [775.01595209, 6.59521058, -11.47942543]
        self.assertTrue(out_len == in_len)
        self.assertTrue(abs(j2000[0]-res[0]) <= 1e-6)
        self.assertTrue(abs(j2000[1]-res[1]) <= 1e-6)
        self.assertTrue(abs(j2000[2]-res[2]) <= 1e-6)

    def test_cotrans_geigsm(self):
        """Test daisy chain of multiple transofrmations, GEI->GSE->GSM."""
        del_data()
        d = [[245.0, -102.0, 251.0], [775.0, 10.0, -10],
             [121.0, 545.0, -1.0], [304.65, -205.3, 856.1],
             [464.34, -561.55, -356.22]]
        t = [1577112800, 1577308800, 1577598800, 1577608800, 1577998800]
        name_out = "tname_out"
        cotrans(name_out=name_out, time_in=t, data_in=d,
                coord_in="gei", coord_out="gsm")
        in_len = len(t)
        dout = get_data(name_out)
        out_len = len(dout[0])
        gsm = dout[1][1]
        res = [4.59847513e+01, 7.62303493e+02, 1.32679267e+02]
        self.assertTrue(out_len == in_len)
        self.assertTrue(abs(gsm[0]-res[0]) <= 1e-6)
        self.assertTrue(abs(gsm[1]-res[1]) <= 1e-6)
        self.assertTrue(abs(gsm[2]-res[2]) <= 1e-6)

    def test_all_cotrans(self):
        """Test all cotrans pairs.

        Apply transformation, then inverse transformation and compare.
        """
        cotrans()
        all_cotrans = ['gei', 'geo', 'j2000', 'gsm', 'mag', 'gse', 'sm']
        d = [[245.0, -102.0, 251.0], [775.0, 10.0, -10],
             [121.0, 545.0, -1.0], [304.65, -205.3, 856.1],
             [464.34, -561.55, -356.22]]
        dd1 = d[1]
        t = [1577112800, 1577308800, 1577598800, 1577608800, 1577998800]
        in_len = len(t)
        name1 = "name1"
        name2 = "name2"
        count = 0
        # Test non-existent systems.
        result = cotrans(name_out=name1, time_in=t, data_in=d,
                coord_in="badcoord", coord_out="gei")
        self.assertTrue(result == 0)
        result = cotrans(name_out=name1, time_in=t, data_in=d,
                         coord_in="gei", coord_out="badcoord")
        self.assertTrue(result == 0)

        # Test empty data.
        cotrans(name_out=name1, time_in=t, data_in=[],
                coord_in="gei", coord_out="geo")
        cotrans(time_in=t, data_in=d, coord_in="gse", coord_out="gsm")
        # Test all combinations.
        for coord_in in all_cotrans:
            for coord_out in all_cotrans:
                count += 1
                del_data()
                cotrans(name_out=name1, time_in=t, data_in=d,
                        coord_in=coord_in, coord_out=coord_out)
                dout = get_data(name1)
                out_len1 = len(dout[0])
                self.assertTrue(out_len1 == in_len)
                # Now perform inverse transformation.
                cotrans(name_in=name1, name_out=name2,
                        coord_in=coord_out, coord_out=coord_in)
                dout2 = get_data(name2)
                out_len2 = len(dout2[0])
                dd2 = dout2[1][1]
                logging.info("%d --- in: %s out: %s", count, coord_in, coord_out)
                # print(dout[1][1])
                # print(dd2)
                self.assertTrue(out_len2 == in_len)
                self.assertTrue(abs(dd1[0]-dd2[0]) <= 1e-6)
                self.assertTrue(abs(dd1[1]-dd2[1]) <= 1e-6)
                self.assertTrue(abs(dd1[2]-dd2[2]) <= 1e-6)


if __name__ == '__main__':
    unittest.main()
