import os

import pytest
import numpy as np

from astropy.io import fits

from pynpoint.core.pypeline import Pypeline
from pynpoint.readwrite.nearreading import NearReadingModule
from pynpoint.util.tests import create_config, create_near_data, remove_test_data


class TestNearReading:

    def setup_class(self) -> None:

        self.limit = 1e-8
        self.test_dir = os.path.dirname(__file__) + "/"
        self.fitsfile = self.test_dir + "near/images_0.fits"

        create_near_data(path=self.test_dir + "near")
        create_config(self.test_dir + "PynPoint_config.ini")

        self.pipeline = Pypeline(self.test_dir, self.test_dir, self.test_dir)

        self.pipeline.set_attribute(
            "config", "NFRAMES", "ESO DET CHOP NCYCLES", static=True
        )
        self.pipeline.set_attribute("config", "EXP_NO", "ESO TPL EXPNO", static=True)
        self.pipeline.set_attribute("config", "NDIT", "None", static=True)
        self.pipeline.set_attribute("config", "PARANG_START", "None", static=True)
        self.pipeline.set_attribute("config", "PARANG_END", "None", static=True)
        self.pipeline.set_attribute("config", "DITHER_X", "None", static=True)
        self.pipeline.set_attribute("config", "DITHER_Y", "None", static=True)
        self.pipeline.set_attribute("config", "PIXSCALE", 0.045, static=True)
        self.pipeline.set_attribute("config", "MEMORY", 100, static=True)

        self.positions = ("chopa", "chopb")

    def teardown_class(self) -> None:

        remove_test_data(self.test_dir, folders=["near"])

    def test_near_read(self) -> None:

        module = NearReadingModule(
            name_in="read1a",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)
        self.pipeline.run_module("read1a")

        for item in self.positions:

            data = self.pipeline.get_data(item)
            assert np.mean(data) == pytest.approx(0.060582854, rel=self.limit, abs=0.0)
            assert data.shape == (10, 10, 10)

    def test_near_subtract_crop_mean(self) -> None:

        module = NearReadingModule(
            name_in="read1b",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
            subtract=True,
            crop=(None, None, 0.3),
            combine="mean",
        )

        self.pipeline.add_module(module)
        self.pipeline.run_module("read1b")

        data = self.pipeline.get_data(self.positions[0])
        assert np.mean(data) == pytest.approx(0.0, rel=self.limit, abs=0.0)
        assert data.shape == (2, 7, 7)

        data = self.pipeline.get_data(self.positions[1])
        assert np.mean(data) == pytest.approx(0.0, rel=self.limit, abs=0.0)
        assert data.shape == (2, 7, 7)

    def test_near_median(self) -> None:

        module = NearReadingModule(
            name_in="read1c",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
            combine="median",
        )

        self.pipeline.add_module(module)
        self.pipeline.run_module("read1c")

        data = self.pipeline.get_data(self.positions[0])
        assert np.mean(data) == pytest.approx(0.060582854, rel=self.limit, abs=0.0)
        assert data.shape == (2, 10, 10)

        data = self.pipeline.get_data(self.positions[1])
        assert np.mean(data) == pytest.approx(0.060582854, rel=self.limit, abs=0.0)
        assert data.shape == (2, 10, 10)

    def test_static_not_found(self) -> None:

        self.pipeline.set_attribute("config", "DIT", "Test", static=True)

        module = NearReadingModule(
            name_in="read2",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module("read2")

        assert len(warning) == 4

        for item in warning:
            assert (
                item.message.args[0]
                == "Static attribute DIT (=Test) not found in the FITS "
                "header."
            )

        self.pipeline.set_attribute("config", "DIT", "ESO DET SEQ1 DIT", static=True)

    def test_nonstatic_not_found(self) -> None:

        self.pipeline.set_attribute("config", "NDIT", "Test", static=True)

        module = NearReadingModule(
            name_in="read3",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module("read3")

        assert len(warning) == 4

        for item in warning:
            assert (
                item.message.args[0]
                == "Non-static attribute NDIT (=Test) not found in the "
                "FITS header."
            )

        self.pipeline.set_attribute("config", "NDIT", "None", static=True)

    def test_check_header(self) -> None:

        with fits.open(self.fitsfile) as hdulist:
            hdulist[0].header["ESO DET CHOP ST"] = "F"
            hdulist[0].header["ESO DET CHOP CYCSKIP"] = 1
            hdulist[0].header["ESO DET CHOP CYCSUM"] = "T"
            hdulist.writeto(self.fitsfile, overwrite=True)

        module = NearReadingModule(
            name_in="read4",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module("read4")

        assert len(warning) == 3

        assert warning[0].message.args[0] == "Dataset was obtained without chopping."
        assert warning[1].message.args[0] == "Chop cycles (1) have been skipped."
        assert warning[2].message.args[0] == "FITS file contains averaged images."

        with fits.open(self.fitsfile) as hdulist:
            hdulist[0].header["ESO DET CHOP ST"] = "T"
            hdulist[0].header["ESO DET CHOP CYCSKIP"] = 0
            hdulist[0].header["ESO DET CHOP CYCSUM"] = "F"
            hdulist.writeto(self.fitsfile, overwrite=True)

    def test_frame_type_invalid(self) -> None:

        with fits.open(self.fitsfile) as hdulist:
            hdulist[10].header["ESO DET FRAM TYPE"] = "Test"
            hdulist.writeto(self.fitsfile, overwrite=True)

        module = NearReadingModule(
            name_in="read5",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.raises(ValueError) as error:
            self.pipeline.run_module("read5")

        assert (
            str(error.value)
            == "Frame type (Test) not a valid value. Expecting HCYCLE1 or "
            "HCYCLE2 as value for ESO DET FRAM TYPE."
        )

    def test_frame_type_missing(self) -> None:

        with fits.open(self.fitsfile) as hdulist:
            hdulist[10].header.remove("ESO DET FRAM TYPE")
            hdulist.writeto(self.fitsfile, overwrite=True)

        module = NearReadingModule(
            name_in="read6",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.raises(ValueError) as error:
            self.pipeline.run_module("read6")

        assert (
            str(error.value)
            == "Frame type not found in the FITS header. Image number: 9."
        )

    def test_same_cycle(self) -> None:

        with fits.open(self.fitsfile) as hdulist:

            with pytest.warns(UserWarning) as warning:
                hdulist[10].header["ESO DET FRAM TYPE"] = "HCYCLE1"

            assert len(warning) == 1

            assert (
                warning[0].message.args[0]
                == "Keyword name 'ESO DET FRAM TYPE' is greater "
                "than 8 characters or contains characters not "
                "allowed by the FITS standard; a HIERARCH card "
                "will be created."
            )

            hdulist.writeto(self.fitsfile, overwrite=True)

        module = NearReadingModule(
            name_in="read7",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module("read7")

        assert len(warning) == 2

        assert (
            warning[0].message.args[0]
            == "Previous and current chop position (HCYCLE1) are "
            "the same. Skipping the current image."
        )

        assert (
            warning[1].message.args[0]
            == "The number of images is not equal for chop A and "
            "chop B."
        )

    def test_odd_number_images(self) -> None:

        with fits.open(self.fitsfile) as hdulist:
            del hdulist[11]
            hdulist.writeto(self.fitsfile, overwrite=True)

        module = NearReadingModule(
            name_in="read8",
            input_dir=self.test_dir + "near",
            chopa_out_tag=self.positions[0],
            chopb_out_tag=self.positions[1],
        )

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module("read8")

        assert len(warning) == 2

        assert (
            warning[0].message.args[0] == f"FITS file contains odd number of images: "
            f"{self.fitsfile}"
        )

        assert (
            warning[1].message.args[0]
            == "The number of chop cycles (5) is not equal to "
            "half the number of available HDU images (4)."
        )
