#########################################################################
# test_task_zenithdelaycal.py
# Copyright (C) 2018
# Associated Universities, Inc. Washington DC, USA.
#
# This script is free software; you can redistribute it and/or modify it
# under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This library is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Library General Public
# License for more details.
#
# CASR ticket: https://open-jira.nrao.edu/browse/CASR-715
# CAS ticket: https://open-jira.nrao.edu/browse/CAS-14763
#
# Based on the requirements listed in casadocs found here:
# https://casadocs.readthedocs.io/en/stable/api/tt/casatasks.calibration.zenithdelaycal.html
#
##########################################################################
import os
import sys
import shutil
import unittest
import numpy as np

from casatools import ms, ctsys, table, msmetadata
from casatasks import zenithdelaycal, applycal, rmtables

tblocal = table()
mslocal = ms()
msmdlocal = msmetadata()

datapath = ctsys.resolve("unittest/zenithdelaycal/")


class ZenithDelayCal_tests(unittest.TestCase):
    prefix = "ba160d"
    msfile = prefix + ".split.test.ms"
    mbdtable = prefix + ".mbd"
    caltable = prefix + ".zen"
    refant = "PT"
    nant = 10  # Number of antennas in msfile.

    reference_caltable = os.path.join(
        datapath, f"zenithdelaycal_reference/{prefix}.zen"
    )

    def setUp(self):
        shutil.copytree(os.path.join(datapath, self.msfile), self.msfile)
        shutil.copytree(os.path.join(datapath, self.mbdtable), self.mbdtable)

    def tearDown(self):
        rmtables([self.msfile, self.mbdtable, self.caltable])

    def test_defaults(self):
        """
        Test solver with default arguments.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        # Convert to ns for testing.
        fp *= 1e9
        # Load expected caltable.
        reference_caltable = os.path.join(
            datapath, f"zenithdelaycal_reference/{self.prefix}.zen"
        )
        tblocal.open(reference_caltable)
        fp_expected = tblocal.getcol("FPARAM")
        tblocal.close()
        fp_expected *= 1e9
        # Compare tables.
        self.assertTrue(
            np.all(np.isclose(fp, fp_expected)), "FPARAM are not as expected."
        )
        # Second order atmos term is zero.
        self.assertTrue(
            np.all(fp[2, :, :] == 0.0), "Second order atmospheric terms are not zero."
        )
        # First and second order clock terms are zero.
        self.assertTrue(
            np.all(fp[[4, 5, 7, 8], :, :] == 0.0),
            "First and second order clock terms are not zero.",
        )
        # Clock for pol1 is the same as for pol0.
        # This might break once we solve for both pols.
        self.assertTrue(
            np.all(fp[[3, 4, 5], 0, :] == fp[[6, 7, 8], 0, :]),
            "clock0 and clock1 are not the same.",
        )

    def test_invalid_poly_order(self):
        """
        Test for invalid polynomial order.
        """
        with self.assertRaises(
            ValueError, msg="No exception raised for invalid polynomial order."
        ):
            zenithdelaycal(
                self.msfile,
                caltable=self.caltable,
                mbdtable=self.mbdtable,
                refant=self.refant,
                clock_polynomial_order=3,
            )
        with self.assertRaises(
            ValueError, msg="No exception raised for invalid polynomial order."
        ):
            zenithdelaycal(
                self.msfile,
                caltable=self.caltable,
                mbdtable=self.mbdtable,
                refant=self.refant,
                atmospheric_polynomial_order=3,
            )

    def test_no_clock(self):
        """
        Test for solutions without clock terms.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            clock_polynomial_order=-1,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        self.assertTrue(
            np.all(fp[[3, 4, 5, 6, 7, 8]] == 0), "Clock terms are not zero."
        )

    def test_no_atmos(self):
        """
        Test for solutions without atmosphere.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            atmospheric_polynomial_order=-1,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        self.assertTrue(np.all(fp[[0, 1, 2]] == 0), "Atmospheric terms are not zero.")

    def test_clock_linear(self):
        """
        Linear polynomial for clock term.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            clock_polynomial_order=1,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        # Load expected caltable.
        reference_caltable = os.path.join(
            datapath, f"zenithdelaycal_reference/{self.prefix}.zen.c1"
        )
        tblocal.open(reference_caltable)
        fp_expected = tblocal.getcol("FPARAM")
        tblocal.close()
        fp_expected *= 1e9
        # Compare tables.
        self.assertTrue(
            np.all(np.isclose(fp, fp_expected)), "FPARAM are not as expected."
        )
        # Second order atmos term is zero.
        self.assertTrue(
            np.all(fp[2, :, :] == 0.0), "Second order atmospheric terms are not zero."
        )
        # Second order clock terms are zero.
        self.assertTrue(
            np.all(fp[[5, 8], :, :] == 0.0), "Second order clock terms are not zero."
        )
        # Clock for pol1 is the same as for pol0.
        self.assertTrue(
            np.all(fp[[3, 4, 5], 0, :] == fp[[6, 7, 8], 0, :]),
            "clock0 and clock1 are not the same.",
        )

    def test_clock_quadratic(self):
        """
        Quadratic polynomial for clock term.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            clock_polynomial_order=2,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        # Load expected caltable.
        reference_caltable = os.path.join(
            datapath, f"zenithdelaycal_reference/{self.prefix}.zen.c2"
        )
        tblocal.open(reference_caltable)
        fp_expected = tblocal.getcol("FPARAM")
        tblocal.close()
        fp_expected *= 1e9
        # Compare tables.
        self.assertTrue(
            np.all(np.isclose(fp, fp_expected)), "FPARAM are not as expected."
        )
        # Second order atmos term is zero.
        self.assertTrue(
            np.all(fp[2, :, :] == 0.0), "Second order atmospheric terms are not zero."
        )
        # Clock for pol1 is the same as for pol0.
        self.assertTrue(
            np.all(fp[[3, 4, 5], 0, :] == fp[[6, 7, 8], 0, :]),
            "clock0 and clock1 are not the same.",
        )

    def test_atmos_const(self):
        """
        Constant offset for atmospheric term.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            atmospheric_polynomial_order=0,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        # Load expected caltable.
        reference_caltable = os.path.join(
            datapath, f"zenithdelaycal_reference/{self.prefix}.zen.a0"
        )
        tblocal.open(reference_caltable)
        fp_expected = tblocal.getcol("FPARAM")
        tblocal.close()
        fp_expected *= 1e9
        # Compare tables.
        self.assertTrue(
            np.all(np.isclose(fp, fp_expected)), "FPARAM are not as expected."
        )
        # First and second order atmos terms are zero.
        self.assertTrue(
            np.all(fp[[1, 2], :, :] == 0.0),
            "First and second order atmospheric terms are not zero.",
        )
        # First and second order clock terms are zero.
        self.assertTrue(
            np.all(fp[[4, 5, 7, 8], :, :] == 0.0),
            "First and second order clock terms are not zero.",
        )
        # Clock for pol1 is the same as for pol0.
        self.assertTrue(
            np.all(fp[[3, 4, 5], 0, :] == fp[[6, 7, 8], 0, :]),
            "clock0 and clock1 are not the same.",
        )

    def test_per_pol_clock(self):
        """
        Per polarization clock solution.
        """
        r = zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            clock_polynomial_order=2,
            per_pol_clock=True,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        fp *= 1e9
        # Clock for pol1 is not the same as for pol0, except for the reference antenna.
        i = r[f"ant_{self.refant}"]["index"]
        a = [j for j in range(3) if j != i]
        self.assertTrue(
            np.all(fp[[3, 4, 5], 0, :][:, a] != fp[[6, 7, 8], 0, :][:, a]),
            "clock0 and clock1 are not the same.",
        )

    def test_plot_files(self):
        """
        Check that plot files are generated when requested.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            plot_delay_solution=True,
        )

        self.assertTrue(
            os.path.exists("zenithdelaycal_delay_plot_ant_PT.png"),
            "Plots were not writen.",
        )

        # Cleanup.
        for filename in os.listdir("."):
            if filename.endswith(".png"):
                full_path = os.path.join(".", filename)
                os.remove(full_path)

    def test_antenna_selection(self):
        """
        Check that antenna selection restricts the output.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            antenna="0," + self.refant,
        )

        self.assertTrue(os.path.exists(self.caltable), "No caltable writen.")
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        tblocal.close()
        self.assertEqual(fp.shape, (9, 1, 2), "FPARAM shape is not as expected.")

        # Using an antenna not in the MS raises an exception.
        with self.assertRaises(
            ValueError, msg="No exception raised for invalid antenna selection."
        ):
            zenithdelaycal(
                self.msfile,
                caltable=self.caltable,
                mbdtable=self.mbdtable,
                refant=self.refant,
                antenna="12," + self.refant,
            )

    def test_airmass_model(self):
        """
        Tests for airmass model.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            atmospheric_delay_model="secant cubed",
        )

        self.assertTrue(os.path.exists(self.caltable))
        tblocal.open(self.caltable)
        fp = tblocal.getcol("FPARAM")
        am = tblocal.getkeyword("AMassMod")
        tblocal.close()
        fp *= 1e9
        # Load expected caltable.
        reference_caltable = os.path.join(
            datapath, f"zenithdelaycal_reference/{self.prefix}.zen.secant_cubed"
        )
        tblocal.open(reference_caltable)
        fp_expected = tblocal.getcol("FPARAM")
        tblocal.close()
        fp_expected *= 1e9
        # Compare tables.
        self.assertTrue(np.all(np.isclose(fp, fp_expected)))
        self.assertEqual(fp.shape, (9, 1, self.nant))
        self.assertTrue(np.all(fp[2, 0, :] == 0.0))
        self.assertTrue(am == "secant cubed", "Airmass model is not secant cubed")

        # Using an unrecognized airmass model raises an exception.
        with self.assertRaises(
            ValueError, msg="No exception raised for invalid atmospheric_delay_model"
        ):
            zenithdelaycal(
                self.msfile,
                caltable=self.caltable,
                mbdtable=self.mbdtable,
                refant=self.refant,
                atmospheric_delay_model="secant tripled",
            )

    def test_field_selection(self):
        """
        Tests for field selection argument.
        """
        field = "J2253+16"
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            field=field,
        )

        # Using a non-existing field raises an exception.
        with self.assertRaises(
            ValueError, msg="No exception raised for invalid field."
        ):
            zenithdelaycal(
                self.msfile,
                caltable=self.caltable,
                mbdtable=self.mbdtable,
                refant=self.refant,
                field="definitely not here",
            )


class ZenithDelayCal_applycal_tests(unittest.TestCase):
    """
    Tests for applycal.
    """

    prefix = "ba160d"
    msfile = prefix + ".split.test.ms"
    mbdtable = prefix + ".mbd"
    caltable = prefix + ".zen"
    refant = "PT"

    def setUp(self):
        shutil.copytree(os.path.join(datapath, self.msfile), self.msfile)
        shutil.copytree(
            os.path.join(datapath, f"zenithdelaycal_reference/{self.caltable}"),
            self.caltable,
        )
        shutil.copytree(os.path.join(datapath, self.mbdtable), self.mbdtable)

    def tearDown(self):
        rmtables([self.msfile, self.caltable, self.mbdtable])

    def test_applycal_zenithdelaycal(self):

        applycal(self.msfile, gaintable=self.caltable, spwmap=[0] * 4, flagbackup=False)

        selection = {"field_id": 25, "antenna1": 4, "antenna2": 7, "data_desc_id": 0}

        msmdlocal.open(self.msfile)
        nspw = msmdlocal.nspw(includewvr=False)
        nchn = msmdlocal.nchan(
            0
        )  # All spectral windows should have the same number of channels.
        msmdlocal.close()

        x = np.arange(nspw * (nchn - 4))
        # Ignore edge channels.
        s = slice(2, -2)
        data_phase = np.zeros((2, nspw, nchn - 4), dtype=float)
        corr_phase = np.zeros((2, nspw, nchn - 4), dtype=float)
        mslocal.open(self.msfile)
        for i in range(nspw):
            mslocal.reset()
            mslocal.selectinit(datadescid=i)
            selection.update({"data_desc_id": i})
            mslocal.select(selection)
            ms_data = mslocal.getdata(["phase", "corrected_phase"], average=True)
            data_phase[0, i, :] = ms_data["phase"][0, s]  # First pol.
            corr_phase[0, i, :] = ms_data["corrected_phase"][0, s]
            data_phase[1, i, :] = ms_data["phase"][-1, s]  # Second pol.
            corr_phase[1, i, :] = ms_data["corrected_phase"][-1, s]
            mslocal.selectinit(reset=True)
        mslocal.close()
        # Fit a straight line to compare the slopes.
        pfit_data = np.polyfit(x, data_phase[0].flatten(), 1)
        pfit_corr = np.polyfit(x, corr_phase[0].flatten(), 1)

        # Corrected delay is less than in data.
        self.assertTrue(
            abs(pfit_corr[0]) < abs(pfit_data[0]),
            "Corrected delay is larger than raw for the first polarization.",
        )
        # Test for an absolute value of the delay.
        self.assertTrue(
            abs(pfit_corr[0]) < 2e-2,
            "Corrected delay is greater than 2e-2 deg for the first polarization.",
        )
        # Same for the second polarization.
        pfit_data = np.polyfit(x, data_phase[1].flatten(), 1)
        pfit_corr = np.polyfit(x, corr_phase[1].flatten(), 1)
        # Corrected delay is less than in data.
        self.assertTrue(
            abs(pfit_corr[0]) < abs(pfit_data[0]),
            "Corrected delay is larger than raw for the second polarization.",
        )
        # Test for an absolute value of the delay.
        self.assertTrue(
            abs(pfit_corr[0]) < 2e-2,
            "Corrected delay is greater than 2e-2 deg for the second polarization.",
        )
        # Orthogonal polarizations had the same solution applied.
        data_diff = np.unwrap(data_phase[0] - data_phase[1])
        corr_diff = np.unwrap(corr_phase[0] - corr_phase[1])
        self.assertTrue(
            np.all(np.isclose(data_diff, corr_diff, atol=1e-3)),
            "Orthogonal polarizations had different phase solutions applied.",
        )

    def test_applycal_zenithdelaycal_per_pol(self):
        """
        Test that orthogonal polarizations get different phase corrections applied
        whe using `per_pol_clock=True` in `zenithdelaycal`.
        """
        zenithdelaycal(
            self.msfile,
            caltable=self.caltable,
            mbdtable=self.mbdtable,
            refant=self.refant,
            per_pol_clock=True,
        )
        applycal(self.msfile, gaintable=self.caltable, spwmap=[0] * 4, flagbackup=False)

        selection = {"field_id": 25, "antenna1": 4, "antenna2": 7, "data_desc_id": 0}

        mslocal.open(self.msfile)
        mslocal.selectinit(datadescid=0)
        mslocal.select(selection)
        ms_data = mslocal.getdata(["phase", "corrected_phase"], average=True)
        mslocal.close()
        # Check that the difference between orthogonal polarizations is not the same in
        # phase and corrected phase.
        data_diff = ms_data["phase"][0, :] - ms_data["phase"][-1, :]
        corr_diff = ms_data["corrected_phase"][0, :] - ms_data["corrected_phase"][-1, :]
        self.assertTrue(
            np.all(~np.isclose(data_diff, corr_diff, atol=1e-5)),
            "Phase difference did not change after correction.",
        )
        # Check that the phase is flatter.
        # Will fit a straight line to the phase as a function of channel,
        # and compare its slope between phase and corrected_phase.
        msmdlocal.open(self.msfile)
        nspw = msmdlocal.nspw(includewvr=False)
        nchn = msmdlocal.nchan(
            0
        )  # All spectral windows should have the same number of channels.
        msmdlocal.close()
        x = np.arange(nspw * (nchn - 4))
        # Ignore edge channels.
        s = slice(2, -2)
        data_phase = np.zeros((2, nspw, nchn - 4), dtype=complex)
        corr_phase = np.zeros((2, nspw, nchn - 4), dtype=complex)
        mslocal.open(self.msfile)
        for i in range(nspw):
            mslocal.reset()
            mslocal.selectinit(datadescid=i)
            selection.update({"data_desc_id": i})
            mslocal.select(selection)
            ms_data = mslocal.getdata(["phase", "corrected_phase"], average=True)
            data_phase[0, i, :] = ms_data["phase"][0, s]  # First pol.
            corr_phase[0, i, :] = ms_data["corrected_phase"][0, s]
            data_phase[1, i, :] = ms_data["phase"][-1, s]  # Second pol.
            corr_phase[1, i, :] = ms_data["corrected_phase"][-1, s]
            mslocal.selectinit(reset=True)
        mslocal.close()
        pfit_data = np.polyfit(x, data_phase[0].flatten(), 1)
        pfit_corr = np.polyfit(x, corr_phase[0].flatten(), 1)
        # Corrected delay is less than in data.
        self.assertTrue(
            abs(pfit_corr[0]) < abs(pfit_data[0]),
            "Corrected delay is larger than raw for the first polarization.",
        )
        # Test for an absolute value of the delay.
        self.assertTrue(
            abs(pfit_corr[0]) < 2e-2,
            "Corrected delay is greater than 2e-2 deg for the first polarization.",
        )
        # Same for the second polarization.
        pfit_data = np.polyfit(x, data_phase[1].flatten(), 1)
        pfit_corr = np.polyfit(x, corr_phase[1].flatten(), 1)
        # Corrected delay is less than in data.
        self.assertTrue(
            abs(pfit_corr[0]) < abs(pfit_data[0]),
            "Corrected delay is larger than raw for the second polarization.",
        )
        # Test for an absolute value of the delay.
        self.assertTrue(
            abs(pfit_corr[0]) < 2e-2,
            "Corrected delay is greater than 2e-2 deg for the second polarization.",
        )


unittest.addModuleCleanup(ZenithDelayCal_tests.tearDown)
unittest.addModuleCleanup(ZenithDelayCal_applycal_tests.tearDown)

"""
# To generate the test data the following CASA6 commands were used
# on the VLBA dataset BA160D as downloaded from the NRAO archive.

importfitsidi(fitsidifile="VLBA_BA160D_ba160d_BIN0_SRC0_0_230710T193428.idifits",
              vis="ba160d.ms",
              constobsid=True,
              scanreindexgap_s=15,
              specframe='GEO')
fringefit(vis="ba160d.ms", 
          caltable="ba160d.sbd", 
          field="13",
          scan="101", 
          solint="inf", 
          zerorates=True, 
          refant="PT",
          minsnr=10,
          parang=True
          )
fringefit(vis="ba160d.ms",
          caltable="ba160d.mbd",
          solint="50s",
          combine="spw",
          field='1,2,3,4,5,6,7,8,9,10,11,14,15,16,17,19,20,21,22,23,24,25,26,27,18,19,30,31',
          zerorates=False,
          refant="PT,FD,OV,NL,HN,MK",
          minsnr=10,
          gaintable=["ba160d.sbd"],
          interp=["nearest"],
          parang=True
          )
clearcal(vis="ba160d.ms")
applycal(vis="ba160d.ms",
         gaintable=["ba160d.sbd"],
         interp=["nearest"],
         parang=True
         )
split(vis="ba160d.ms",
      outputvis="ba160d.split.test.ms",
      scan="1,233",
      width="16",
      datacolumn="corrected",
      keepflags=False,
     )
"""


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