import unittest

import numpy as np

from tart.imaging import antenna_model, correlator
from tart.operation import observation, settings
from tart.simulation import antennas, radio, simulation_source
from tart.util import angle, utc

TEST_CONFIG = "./tart/test/test_telescope_config.json"
TEST_ANTENNA_POSITIONS = "./tart/test/test_calibrated_antenna_positions.json"


def relative_diff(l1, l2):
    return (l1 - l2) / np.sqrt(np.power(l1, 2) + np.power(l2, 2))


class TestSimulatedVisibility(unittest.TestCase):
    def setUp(self):
        self.utc_date = utc.now()
        self.config = settings.from_file(TEST_CONFIG)
        self.config.load_antenna_positions(
            cal_ant_positions_file=TEST_ANTENNA_POSITIONS
        )

        self.rad = radio.Max2769B(
            noise_level=[0.0 for _ in range(self.config.get_num_antenna())]
        )
        self.sources = [
            simulation_source.SimulationSource(
                r=1e9,
                amplitude=1.0,
                azimuth=angle.from_dms(2.0),
                elevation=angle.from_dms(50.0),
                sample_duration=self.rad.sample_duration,
            )
        ]
        self.ants = [
            antennas.Antenna(self.config.get_loc(), pos)
            for pos in self.config.get_antenna_positions()
        ]
        self.ant_models = [
            antenna_model.GpsPatchAntenna()
            for i in range(self.config.get_num_antenna())
        ]
        self.cor = correlator.Correlator()
        self.timebase = np.arange(
            0, self.rad.sample_duration, 1.0 / self.rad.sampling_rate
        )

        ant_sigs_full = antennas.antennas_signal(
            self.ants, self.ant_models, self.sources, self.timebase
        )
        obs = self.rad.get_full_obs(
            ant_sigs_full, self.utc_date, self.config, self.timebase
        )
        vis = self.cor.correlate(obs)
        self.full_vis = np.array(vis.v)

        ant_sigs_simp = antennas.antennas_simplified_signal(
            self.ants,
            self.ant_models,
            self.sources,
            self.rad.baseband_timebase,
            self.rad.int_freq,
        )
        obs = self.rad.get_simplified_obs(
            ant_sigs_simp, self.utc_date, config=self.config
        )
        vis2 = self.cor.correlate(obs)
        self.sim_vis = np.array(vis2.v)

#        plt.figure(figsize=(4, 4))
#        plt.scatter(self.full_vis.real, self.full_vis.imag, label="full")
#        plt.scatter(self.sim_vis.real, self.sim_vis.imag, color="red", label="simpl")
#        plt.legend()
#        plt.show()

    # def test_full_vis(self):
    #   self.assertGreater(self.full_vis.max(), 0.8)
    #   self.assertLess(self.full_vis.min(), -0.8)

    # def test_simp_vis(self):
    #   self.assertGreater(self.sim_vis.max(), 0.8)
    #   self.assertLess(self.sim_vis.min(), -0.8)

    def test_compare_vis(self):
        d = self.full_vis - self.sim_vis
        print(d)
        print("relative diffs", relative_diff(self.full_vis, self.sim_vis))
        self.assertLess(abs(d.max()), 0.3)
        self.assertLess(d.std(), 0.1)


class TestTelescope(unittest.TestCase):
    def setUp(self):

        # rad = radio.Max2769B(noise_level=0.)
        utc_date = utc.now()
        self.config = settings.from_file(TEST_CONFIG)
        self.config.load_antenna_positions(
            cal_ant_positions_file=TEST_ANTENNA_POSITIONS
        )

        rad = radio.Max2769B(
            noise_level=[0.0 for _ in range(self.config.get_num_antenna())]
        )
        sources = [
            simulation_source.SimulationSource(
                r=1e9,
                amplitude=1.0,
                azimuth=angle.from_dms(2.0),
                elevation=angle.from_dms(50.0),
                sample_duration=rad.sample_duration,
            )
        ]
        ant_models = [
            antenna_model.GpsPatchAntenna()
            for i in range(self.config.get_num_antenna())
        ]
        ants = [
            antennas.Antenna(self.config.get_loc(), pos)
            for pos in self.config.get_antenna_positions()
        ]
        self.timebase = np.arange(0, rad.sample_duration, 1.0 / rad.sampling_rate)

        ant_sigs = antennas.antennas_signal(ants, ant_models, sources, self.timebase)
        self.obs = rad.get_full_obs(ant_sigs, utc_date, self.config, self.timebase)

    def test_data_limits(self):
        for i in range(1, self.obs.config.get_num_antenna()):
            data = self.obs.get_antenna(i)

            self.assertGreater(data.max(), 0.99)
            self.assertLess(data.min(), -0.99)

            self.assertGreater(data.mean(), -0.1)
            self.assertLess(data.mean(), 0.1)

    def test_settings(self):
        self.assertEqual(self.config.get_num_antenna(), 24)

    def test_simulated_observation(self):
        fname = "test_observation.pkl"
        self.obs.save(fname)
        nobs = observation.Observation_Load(fname)
        self.assertEqual(self.obs.timestamp, nobs.timestamp)
        self.assertTrue((self.obs.get_antenna(1) == nobs.get_antenna(1)).all())
