"""
Tests for cotrans_tools 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_tools/dsl2gse.py

"""

import os
import unittest
import logging
import numpy as np
from numpy.testing import assert_allclose
import pyspedas
from pyspedas import (
    sm2mlt,
    tplot_copy,
    set_units,
    get_units,
    time_double,
    get_data,
    store_data,
    del_data,
    get_coords,
    set_coords,
    tplot_restore,
    tplot_names,
    options,
    subtract,
    tplot,
    cotrans,
    fac_matrix_make,
)
from pyspedas.projects.themis.cotrans.dsl2gse import dsl2gse
from pyspedas.utilities.config_testing import TESTING_CONFIG, test_data_download_file

# Whether to display plots during testing
global_display = TESTING_CONFIG["global_display"]
# Directory to save testing output files
output_dir = TESTING_CONFIG["local_testing_dir"]
# Ensure output directory exists
cotrans_dir = "cotrans_tools"
save_dir = os.path.join(output_dir, cotrans_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
# Directory with IDL SPEDAS validation files
validation_dir = TESTING_CONFIG["remote_validation_dir"]


class CotransTestCases(unittest.TestCase):
    """Tests for cotrans."""

    def test_fac_matrix_make(self):
        with self.assertLogs(level="ERROR") as log:
            doesntexist = fac_matrix_make("doesnt_exist")
            self.assertTrue(doesntexist is None)
            self.assertIn("Error reading tplot variable: doesnt_exist", log.output[0])

    def test_get_set_coords(self):
        """Test for pytplot.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]})
        with self.assertLogs(level="ERROR") as log:
            cotrans(name_in="test_coord", coord_out="geo")
            self.assertIn(
                "cotrans error: No input coordinates were provided.", log.output[0]
            )
        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"""

        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_tools.dsl2gse."""
        del_data()
        # Try with missing variables. It should exit without problems.
        with self.assertLogs(level="ERROR") as log:
            dsl2gse("tha_fgl_dsl", "tha_fgl_gse")
            self.assertIn("Variables needed: ['tha_fgl_dsl']", log.output[0])
            self.assertIn("Variables missing: ['tha_fgl_dsl']", log.output[1])
            self.assertIn("Please load missing variables.", log.output[2])
        # Now load the needed variables.
        time_range = ["2017-03-23 00:00:00", "2017-03-23 23:59:59"]
        pyspedas.projects.themis.state(
            probe="a",
            trange=time_range,
            get_support_data=True,
            varnames=["tha_spinras", "tha_spindec"],
        )
        pyspedas.projects.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.projects.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.projects.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.
        with self.assertLogs(level="ERROR"):
            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.70357713, -7.93890939]
        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 = [45.98475132, 762.29816225, 132.7098884]
        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.
        """
        with self.assertLogs(level="ERROR") as log:
            cotrans()
            self.assertIn(
                "cotrans error: No output coordinates were provided.", log.output[0]
            )
        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.
        with self.assertLogs(level="ERROR") as log:
            result = cotrans(
                name_out=name1,
                time_in=t,
                data_in=d,
                coord_in="badcoord",
                coord_out="gei",
            )
            self.assertTrue(result == 0)
            self.assertIn(
                "cotrans error: Requested input coordinate system", log.output[0]
            )
            result = cotrans(
                name_out=name1,
                time_in=t,
                data_in=d,
                coord_in="gei",
                coord_out="badcoord",
            )
            self.assertTrue(result == 0)
            self.assertIn(
                "cotrans error: Requested output coordinate system", log.output[1]
            )

        # Test empty data.
        with self.assertLogs(level="ERROR") as log:
            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")
            self.assertIn("cotrans error: Data is empty.", log.output[0])
            self.assertEqual(len(log.output), 1)
        # Test all combinations.
        del_data()  # Ensure there are no leftover variables.
        with self.assertLogs(level="WARNING") as log:
            for i, coord_in in enumerate(all_cotrans):
                for j, coord_out in enumerate(all_cotrans):
                    count += 1
                    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 coord_out == coord_in:
                        self.assertIn(
                            "Warning: coord_in equal to coord_out.",
                            log.output[i * 2],
                        )
                        self.assertIn(
                            "Warning: coord_in equal to coord_out.",
                            log.output[j * 2 + 1],
                        )
                    del_data()  # Clean up for next test.

    def test_mlt(self):
        """Test sm2mlt.

        Data is from goes18, 2023-01-01. [0, 100, 600, 1000]
        """

        # SM coordinates and MLT values from IDL
        sm_idl = [
            [33199.660, 25815.376, 2995.6487],
            [18814.747, 37614.243, 2990.7136],
            [-40428.563, -11607.452, 2995.3227],
            [15798.372, -38974.509, 2995.0028],
        ]
        mlt_idl = [14.524527, 16.228377, 1.0679544, 7.4710163]

        x = [item[0] for item in sm_idl]
        y = [item[1] for item in sm_idl]
        z = [item[2] for item in sm_idl]

        mlt_python = sm2mlt(x, y, z)

        for i in range(4):
            self.assertTrue(abs(mlt_idl[i] - mlt_python[i]) <= 1e-6)

    def test_units_coords(self):
        """Test setting units and coordinate systems (including wildcards)"""

        test_dat = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        test_times = np.array([1.0, 2.0, 3.0, 4.0, 5.0])

        store_data("v1", data={"x": test_times, "y": test_dat})
        tplot_copy("v1", "v2")
        tplot_copy("v1", "v3")
        tplot_copy("v1", "v4")

        set_coords("v?", "GSE")
        c1 = get_coords("v1")
        c2 = get_coords("v2")
        c3 = get_coords("v3")
        c4 = get_coords("v4")
        self.assertTrue(c1 == "GSE")
        self.assertTrue(c2 == "GSE")
        self.assertTrue(c3 == "GSE")
        self.assertTrue(c4 == "GSE")

        set_units("v?", "nT")
        u1 = get_units("v1")
        u2 = get_units("v2")
        u3 = get_units("v3")
        u4 = get_units("v4")
        self.assertTrue(u1 == "nT")
        self.assertTrue(u2 == "nT")
        self.assertTrue(u3 == "nT")
        self.assertTrue(u4 == "nT")

        res = set_coords("nonexistent", "GSE")
        self.assertFalse(res)

        res = set_units("xxx yyy zzz", "nT")
        self.assertFalse(res)

    def test_geogsm_geopack(self):
        # Originally reported by Bob Weigel -- PySPEDAS GEO->GSM cotrans had weird artifacts at year boundaries
        # compared to Geopack GEO->GSM transform and other Python packages
        # The SPEDAS script that creates the file: general/tools/python_validate/cotrans_gsm_test.pro

        # Load validation variables from the test file
        filename = test_data_download_file(
            validation_dir, cotrans_dir, "geogsm_gp_test.tplot", save_dir
        )

        tplot_restore(filename)
        tplot_names()
        cotrans("zgeo", "zgeo_gsm_py", coord_in="geo", coord_out="gsm")
        options("zgeo_gsm_py", "ytitle", "zgeo_gsm_py")
        subtract("zgeo_gsm_py", "zgeo_gsm", newname="pydiff")
        options("pydiff", "ytitle", "idl_python_diff")
        subtract("zgeo_gsm_gp", "zgeo_gsm", newname="gpdiff")
        options("gpdiff", "ytitle", "idl_geopack_diff")
        local_png = os.path.join(save_dir, "geo_gsm_diffs.png")
        tplot(
            ["zgeo", "zgeo_gsm", "zgeo_gsm_py", "pydiff", "gpdiff"],
            display=global_display,
            save_png=local_png,
        )
        didl = get_data("zgeo_gsm")
        dpy = get_data("zgeo_gsm_py")
        # Max relative difference around 2.7e-07
        assert_allclose(didl.y, dpy.y, atol=3.0e-07)

    def test_contrans_single_point(self):
        """Test cotrans with a single data point input."""
        data_in = np.array([[3.46410162, 3.46410162, 3.46410162]])
        time_in = [time_double("2015-12-30T00:00:00")]
        known_frames = ["GSE", "GSM", "GEI", "SM", "GEO", "J2000"]
        for f1 in known_frames:
            for f2 in known_frames:
                res = cotrans(
                    time_in=time_in, data_in=data_in, coord_in=f1, coord_out=f2
                )
                self.assertEqual(len(res[0]), 3)


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