from tart.imaging import location
from tart.imaging import sun
from tart.util import angle
from tart.imaging import tart_util

from tart.util import utc

import datetime
import numpy as np
import unittest

from astropy import coordinates as coord
from astropy import units as u
from astropy import time
from astropy.time import Time

Dunedin = location.Dunedin
Aachen = location.Aachen


class TestLocation(unittest.TestCase):
    # http://www.heavens-above.com/whattime.aspx?lat=-45.86391200&lng=170.51348452&loc=TART&alt=46.5&tz=NZST
    def setUp(self):
        self.utc_date = utc.utc_datetime(
            year=2009, month=7, day=22, hour=5, minute=9, second=50.3
        )

    def astropy_get_AltAz(self, r, el, az, loc, utc_date):
        """Get an astropy AltAz object
        r - distance from observer to object (m)
        el - elevation angle
        az - azimuth angle
        """
        obstime = time.Time(utc_date, scale="utc")
        eloc = coord.EarthLocation(
            lat=loc.latitude_deg() * u.deg,
            lon=loc.longitude_deg() * u.deg,
            height=loc.alt * u.m,
        )

        # altaz_frame = coord.AltAz(obstime=obstime, location=eloc,
        # temperature=15 * u.deg_C, pressure=1.010 * u.bar)
        altaz_frame = coord.AltAz(obstime=obstime, location=eloc)

        elaz = coord.SkyCoord(
            alt=el * u.deg, az=az * u.deg, distance=r * u.m, frame=altaz_frame
        )
        return elaz

    def astropy_horizontal_to_ECI(self, r, el, az, loc, utc_date):
        obstime = time.Time(utc_date, scale="utc")
        elaz = self.astropy_get_AltAz(r, el, az, loc, utc_date)
        gcrs = elaz.transform_to(coord.GCRS(obstime=obstime))
        return [gcrs.cartesian.x.value, gcrs.cartesian.y.value, gcrs.cartesian.z.value]

    def astropy_horizontal_to_ECEF(self, r, el, az, loc, utc_date):
        obstime = time.Time(utc_date, scale="utc")
        elaz = self.astropy_get_AltAz(r, el, az, loc, utc_date)
        itrf = elaz.transform_to(coord.ITRS)
        return [itrf.cartesian.x.value, itrf.cartesian.y.value, itrf.cartesian.z.value]

    """
    http://www.ngs.noaa.gov/cgi-bin/xyz_getxyz.prl
    =============================================================
                        Latitude                     Longitude            Ellip_Ht     Ellipsoid
    INPUT =     S455150.0832             E1703048.54427 46.5             GRS80        
    =============================================================


        X (Meters)        Y (Meters)        Z (Meters)        ELLIPSOID
        ------------- ------------- ------------- ---------

        -4388375.1765     733300.4705 -4554760.4738 GRS80
    """

    def test_ecef(self):
        x, y, z = Dunedin.get_ecef()
        self.assertAlmostEqual(x, -4388375.1765, 3)
        self.assertAlmostEqual(y, 733300.4705, 3)
        self.assertAlmostEqual(z, -4554760.4738, 3)

        x, y, z = Dunedin.get_ecef_wgs84()
        self.assertAlmostEqual(x, -4388375.1765, 3)
        self.assertAlmostEqual(y, 733300.4705, 3)
        self.assertAlmostEqual(z, -4554760.4738, 3)

    def test_XYX(self):
        x, y, z = Dunedin.get_XYZ(0.0, 0.0, 0.0)
        self.assertAlmostEqual(x, 0.0)
        self.assertAlmostEqual(y, 0.0)
        self.assertAlmostEqual(z, 0.0)
        x, y, z = Dunedin.get_XYZ(0.0, 10.0, 0.0)  # 100m north
        print(x, y, z)
        dx, dy, dz = Dunedin.get_ecef_delta_from_enu(0.0, 10.0, 0.0)
        print(dx, dy, dz)
        self.assertAlmostEqual(x, dx, 1)
        self.assertAlmostEqual(y, dy, 1)
        self.assertAlmostEqual(z, dz, 1)

    """ 
    The test vectors for this are generated by 
    http://xjubier.free.fr/en/site_pages/astronomy/coordinatesConverter.html
    """

    def test_get_ecef_delta_from_enu(self):
        # print('Testing ECEF2ENU2ECEF')
        x, y, z = Dunedin.get_ecef()
        e, n, u = Dunedin.ecef_to_enu(x, y, z)
        x2, y2, z2 = Dunedin.get_ecef_delta_from_enu(e, n, u)
        self.assertAlmostEqual(x, x2, 3)
        self.assertAlmostEqual(y, y2, 3)
        self.assertAlmostEqual(z, z2, 3)

    def test_horizontal_to_equatorial(self):
        # el_t = angle.from_dms(90.)
        # az_t = angle.from_dms(0.0)
        # ra, decl = Dunedin.horizontal_to_equatorial(self.utc_date, el_t, az_t)
        # st = Dunedin.LST(self.utc_date)
        # self.assertAlmostEqual(st.to_degrees(), ra.to_degrees(), 5)

        import ephem

        dnd = ephem.Observer()
        dnd.lon = str(Dunedin.lon.to_degrees())
        dnd.lat = str(Dunedin.lat.to_degrees())
        dnd.elevation = Dunedin.alt
        dnd.pressure = 0.0
        dnd.date = self.utc_date

        for az in np.arange(0, 360, 10):
            for el in np.arange(-90, 90, 5):
                az_i = angle.from_dms(az)
                el_i = angle.from_dms(el)

                ra, dec = Dunedin.horizontal_to_equatorial(self.utc_date, el_i, az_i)
                el_f, az_f = Dunedin.equatorial_to_horizontal(self.utc_date, ra, dec)

                # print(Dunedin)
                # print('INIT: el az:', el_i, az_i)
                # print('RA, DEC', ra, dec)
                # print('FINAL: el az:', el_f, az_f)

                self.assertTrue((el_i - el_f) < angle.from_dms(0.01))
                if abs(el) != 90:
                    self.assertTrue((az_i - az_f) < angle.from_dms(0.01))

    def test_equatorial_to_horizontal_dunedin(self):

        # http://www.ccdimages.com/Planets.aspx
        # get RA,dec for planets
        # verify RA and dec with ephem                                                [OK]
        # get el and az with ephem for dunedin today
        # calc el and az from RA and dec with OUR function

        import ephem

        dnd = ephem.Observer()
        dnd.lon = str(Dunedin.lon.to_degrees())
        dnd.lat = str(Dunedin.lat.to_degrees())
        dnd.elevation = Dunedin.alt
        dnd.pressure = 0.0
        datee = datetime.datetime.utcnow()
        dnd.date = datee

        objectlist = [
            ephem.Sun(),
            ephem.Mercury(),
            ephem.Venus(),
            ephem.Mars(),
            ephem.Jupiter(),
            ephem.Saturn(),
            ephem.Uranus(),
            ephem.Neptune(),
        ]
        for j in objectlist:
            j.compute(dnd)
            ra = angle.from_rad(j.ra.real)
            dec = angle.from_rad(j.dec.real)
            el, az = Dunedin.equatorial_to_horizontal(datee, ra, dec)
            self.assertAlmostEqual(el.to_rad(), j.alt.real, 3)
            self.assertAlmostEqual(az.to_rad(), j.az.real, 3)

    def test_equatorial_to_horizontal_aachen(self):

        # http://www.ccdimages.com/Planets.aspx
        # get RA,dec for planets
        # verify RA and dec with ephem                                                [OK]
        # get el and az with ephem for dunedin today
        # calc el and az from RA and dec with OUR function

        import ephem

        dnd = ephem.Observer()
        dnd.lon = str(Aachen.lon.to_degrees())
        dnd.lat = str(Aachen.lat.to_degrees())
        dnd.elevation = Aachen.alt
        dnd.pressure = 0.0
        datee = datetime.datetime.utcnow()

        dnd.date = datee
        # print(dnd.date)

        objectlist = [
            ephem.Sun(),
            ephem.Mercury(),
            ephem.Venus(),
            ephem.Mars(),
            ephem.Jupiter(),
            ephem.Saturn(),
            ephem.Uranus(),
            ephem.Neptune(),
        ]
        for j in objectlist:
            j.compute(dnd)
            ra = angle.from_rad(j.ra.real)
            dec = angle.from_rad(j.dec.real)
            el, az = Aachen.equatorial_to_horizontal(datee, ra, dec)
            self.assertAlmostEqual(el.to_rad(), j.alt.real, 3)
            self.assertAlmostEqual(az.to_rad(), j.az.real, 3)

    def test_horizontal_to_equatorial_astropy(self):
        from astropy import coordinates as coord
        from astropy import units as u
        from astropy import time
        from astropy.time import Time

        utc_date = datetime.datetime.utcnow()
        loc = Dunedin

        obstime = time.Time(utc_date, scale="utc")
        eloc = coord.EarthLocation(
            lat=loc.latitude_deg() * u.deg,
            lon=loc.longitude_deg() * u.deg,
            height=loc.alt * u.m,
        )
        altaz_frame = coord.AltAz(obstime=obstime, location=eloc)

        for el, az in zip(np.linspace(0, 90, 10), np.linspace(0, 259, 10)):

            elaz = coord.SkyCoord(alt=el * u.deg, az=az * u.deg, frame=altaz_frame)
            radec = elaz.transform_to(coord.ICRS)

            ra, dec = loc.horizontal_to_equatorial(
                utc_date, angle.from_dms(el), angle.from_dms(az)
            )

            self.assertAlmostEqual(
                radec.ra.degree, ra.to_degrees(), -1
            )  # TODO better agreement should be possible.
            self.assertAlmostEqual(
                radec.dec.degree, dec.to_degrees(), 1
            )  # TODO better agreement should be possible.

    def test_LST(self):
        """ the Right Ascension of the Zenith: it is exactly the same as your Local Sidereal Time. """
        st = Dunedin.LST(self.utc_date)
        check = angle.from_hours(12.53767)
        self.assertAlmostEqual(st.to_degrees(), check.to_degrees(), 2)

        utc_datetime = datetime.datetime(2013, 9, 12, 9, 10, 0)

        LST = Dunedin.LST(utc_datetime).to_hours() % 24

        lst_h = int(LST)
        lst_m = int((LST - int(LST)) * 60.0)
        lst_s = int(((LST - int(LST)) * 60 - int((LST - int(LST)) * 60)) * 60)

        self.assertAlmostEqual(lst_h, 19, 2)
        self.assertAlmostEqual(lst_m, 58, 2)
        self.assertAlmostEqual(lst_s, 13, 2)

    def test_LST_astropy(self):
        start_date = datetime.datetime(2013, 9, 12, 9, 10, 0)

        loc = Dunedin

        eloc = coord.EarthLocation(
            lat=loc.latitude_deg() * u.deg,
            lon=loc.longitude_deg() * u.deg,
            height=loc.alt * u.m,
        )

        for t in np.linspace(0, 1440, 50):
            utc_date = start_date + datetime.timedelta(seconds=t)
            LST = loc.LST(utc_date)
            h, m, s = LST.to_hms()
            print(LST.to_hours(), h, m, s)
            obstime = time.Time(utc_date, scale="utc")
            st = obstime.sidereal_time("mean", longitude=eloc.lon)
            print(st)
            self.assertAlmostEqual(h, st.hms.h, 5)
            self.assertAlmostEqual(m, st.hms.m, 5)
            self.assertTrue(abs(s - st.hms.s) < 0.9)

    def test_GST(self):
        utc_datetime = datetime.datetime(2013, 9, 12, 9, 10, 0)
        GST = Dunedin.GST(utc_datetime).to_hours() % 24
        gst_h = int(GST)
        gst_m = int((GST - int(GST)) * 60.0)
        gst_s = int(((GST - int(GST)) * 60 - int((GST - int(GST)) * 60)) * 60)

        self.assertAlmostEqual(gst_h, 8, 2)
        self.assertAlmostEqual(gst_m, 36, 2)
        self.assertAlmostEqual(gst_s, 9, 2)

    def test_GST_astropy(self):
        utc_date = datetime.datetime(2013, 9, 12, 9, 10, 0)
        GST = Dunedin.GST(utc_date).to_hours() % 24
        gst_h = int(GST)
        gst_m = int((GST - int(GST)) * 60.0)
        gst_s = int(((GST - int(GST)) * 60 - int((GST - int(GST)) * 60)) * 60)

        obstime = time.Time(utc_date, scale="utc")
        st = obstime.sidereal_time("apparent", "greenwich")
        print(st)
        self.assertAlmostEqual(gst_h, st.hms.h, 5)
        self.assertAlmostEqual(gst_m, st.hms.m, 5)
        self.assertAlmostEqual(gst_s, int(st.hms.s), 5)

    def test_julian_date(self):
        jd = tart_util.JulianDay(self.utc_date)
        self.assertAlmostEqual(jd, 2455034.71516, 5)

    def test_solar_elevation(self):
        t = utc.now()
        s = sun.Sun()
        ra, decl = s.radec(t)

        el, az = Dunedin.equatorial_to_horizontal(t, ra, decl)
        el1, az1 = s.to_horizontal(Dunedin, t)
        self.assertAlmostEqual(el.to_degrees(), el1.to_degrees(), 5)
        self.assertAlmostEqual(az.to_degrees(), az1.to_degrees(), 5)

    def test_eci_to_ecef(self):
        utc_date = utc.now()
        n_tests = 50
        x_arr = np.random.rand(n_tests) * 12345e3
        y_arr = np.random.rand(n_tests) * 23456e3
        z_arr = np.random.rand(n_tests) * 12345e3

        for x_in, y_in, z_in in zip(x_arr, y_arr, z_arr):
            ecef = location.eci_to_ecef(utc_date, x_in, y_in, z_in)

            x, y, z = location.ecef_to_eci(utc_date, ecef[0], ecef[1], ecef[2])
            self.assertAlmostEqual(x_in, x)
            self.assertAlmostEqual(y_in, y)
            self.assertAlmostEqual(z_in, z)

    def test_horizontal_to_ecef(self):
        theta = 90.0  # Straight Up
        phi = 0.0

        x, y, z = Dunedin.horizontal_to_ecef(
            0.0, angle.from_dms(theta), angle.from_dms(phi)
        )

        ecef = Dunedin.get_ecef()

        self.assertAlmostEqual(x, ecef[0])
        self.assertAlmostEqual(y, ecef[1])
        self.assertAlmostEqual(z, ecef[2])

    def test_horizontal_to_ecef_and_back(self):
        n_tests = 50

        theta = np.random.rand(n_tests) * np.pi / 2
        phi = np.random.rand(n_tests) * np.pi * 2
        r = 100.0
        for i in range(n_tests):
            x, y, z = Dunedin.horizontal_to_ecef(
                r, angle.from_rad(theta[i]), angle.from_rad(phi[i])
            )

            r2, theta2, phi2 = Dunedin.ecef_to_horizontal(x, y, z)

            self.assertAlmostEqual(theta2.to_rad(), theta[i])
            self.assertAlmostEqual(phi2.to_rad(), phi[i])

    def test_horizontal_to_ecef_vs_astropy(self):
        loc = Dunedin

        n_tests = 50
        el_arr = np.random.rand(n_tests) * np.pi / 2
        az_arr = np.random.rand(n_tests) * np.pi * 2
        r_arr = np.random.rand(n_tests) * 1e8

        utc_date = utc.now()

        for r, el, az in zip(r_arr, el_arr, az_arr):
            x, y, z = self.astropy_horizontal_to_ECEF(r, el, az, loc, utc_date)
            xi, yi, zi = Dunedin.horizontal_to_ecef(
                r, angle.from_dms(el), angle.from_dms(az)
            )

            self.assertAlmostEqual(x / xi, 1.0, 3)
            self.assertAlmostEqual(y / yi, 1.0, 3)
            self.assertAlmostEqual(z / zi, 1.0, 3)

    def test_horizontal_to_eci_and_back(self):

        n_tests = 50
        theta_arr = np.random.rand(n_tests) * np.pi / 2
        phi_arr = np.random.rand(n_tests) * np.pi * 2
        r = 1.3e7
        utc_date = utc.now()

        for theta, phi in zip(theta_arr, phi_arr):
            xi, yi, zi = Dunedin.horizontal_to_eci(
                r, angle.from_rad(theta), angle.from_rad(phi), utc_date
            )
            r2, theta2, phi2 = Dunedin.eci_to_horizontal(xi, yi, zi, utc_date)

            self.assertAlmostEqual(theta2.to_rad(), theta)
            self.assertAlmostEqual(phi2.to_rad(), phi)
            self.assertAlmostEqual(r2, r)

    def test_horizontal_to_eci_vs_astropy(self):
        utc_date = datetime.datetime.utcnow()
        loc = Dunedin

        n_tests = 50
        el_arr = np.random.rand(n_tests) * np.pi / 2
        az_arr = np.random.rand(n_tests) * np.pi * 2
        r_arr = np.random.rand(n_tests) * 1e8

        # utc_date = utc.now()

        for r, el, az in zip(r_arr, el_arr, az_arr):
            x, y, z = self.astropy_horizontal_to_ECI(r, el, az, loc, utc_date)
            xi, yi, zi = Dunedin.horizontal_to_eci(
                r, angle.from_dms(el), angle.from_dms(az), utc_date
            )
            delta = np.sqrt((x - xi) ** 2 + (y - yi) ** 2 + (z - zi) ** 2)
            print(f"loc_eci : {xi,yi,zi}, astropy: {x,y,z}")
            # TODO these are failing
            self.assertAlmostEqual(delta / r, 0.0, 3)
            self.assertAlmostEqual(x / xi, 1.0, 3)
            self.assertAlmostEqual(y / yi, 1.0, 3)
            self.assertAlmostEqual(z / zi, 1.0, 3)
