import os

import h5py
import pytest
import numpy as np

from pynpoint.core.pypeline import Pypeline
from pynpoint.readwrite.fitsreading import FitsReadingModule
from pynpoint.processing.psfpreparation import PSFpreparationModule
from pynpoint.processing.psfsubtraction import (
    PcaPsfSubtractionModule,
    ClassicalADIModule,
)
from pynpoint.util.tests import create_config, create_fake_data, remove_test_data


class TestPsfSubtractionAdi:

    def setup_class(self) -> None:

        self.limit = 1e-10
        self.test_dir = os.path.dirname(__file__) + "/"

        create_fake_data(self.test_dir + "science")
        create_fake_data(self.test_dir + "reference")
        create_config(self.test_dir + "PynPoint_config.ini")

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

    def teardown_class(self) -> None:

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

    def test_read_data(self) -> None:

        module = FitsReadingModule(
            name_in="read1", image_tag="science", input_dir=self.test_dir + "science"
        )

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

        data = self.pipeline.get_data("science")
        assert np.sum(data) == pytest.approx(
            11.012854046962481, rel=self.limit, abs=0.0
        )
        assert data.shape == (10, 21, 21)

        self.pipeline.set_attribute(
            "science", "PARANG", np.linspace(0.0, 180.0, 10), static=False
        )

        module = FitsReadingModule(
            name_in="read2",
            image_tag="reference",
            input_dir=self.test_dir + "reference",
        )

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

        data = self.pipeline.get_data("reference")
        assert np.sum(data) == pytest.approx(
            11.012854046962481, rel=self.limit, abs=0.0
        )
        assert data.shape == (10, 21, 21)

    def test_psf_preparation(self) -> None:

        module = PSFpreparationModule(
            name_in="prep1",
            image_in_tag="science",
            image_out_tag="science_prep",
            mask_out_tag=None,
            norm=False,
            resize=None,
            cent_size=0.05,
            edge_size=1.0,
        )

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

        data = self.pipeline.get_data("science_prep")
        assert np.sum(data) == pytest.approx(5.029285028467547, rel=self.limit, abs=0.0)
        assert data.shape == (10, 21, 21)

        module = PSFpreparationModule(
            name_in="prep2",
            image_in_tag="reference",
            image_out_tag="reference_prep",
            mask_out_tag=None,
            norm=False,
            resize=None,
            cent_size=0.05,
            edge_size=1.0,
        )

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

        data = self.pipeline.get_data("reference_prep")
        assert np.sum(data) == pytest.approx(5.029285028467547, rel=self.limit, abs=0.0)
        assert data.shape == (10, 21, 21)

    def test_classical_adi(self) -> None:

        module = ClassicalADIModule(
            threshold=None,
            nreference=None,
            residuals="mean",
            extra_rot=0.0,
            name_in="cadi1",
            image_in_tag="science_prep",
            res_out_tag="cadi_res",
            stack_out_tag="cadi_stack",
        )

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

        data = self.pipeline.get_data("cadi_res")
        assert np.sum(data) == pytest.approx(
            0.8381625719865213, rel=self.limit, abs=0.0
        )
        assert data.shape == (10, 21, 21)

        data = self.pipeline.get_data("cadi_stack")
        assert np.sum(data) == pytest.approx(
            0.08395606034388256, rel=self.limit, abs=0.0
        )
        assert data.shape == (1, 21, 21)

    def test_classical_adi_threshold(self) -> None:

        module = ClassicalADIModule(
            threshold=(0.1, 0.03, 1.0),
            nreference=5,
            residuals="median",
            extra_rot=0.0,
            name_in="cadi2",
            image_in_tag="science_prep",
            res_out_tag="cadi_res",
            stack_out_tag="cadi_stack",
        )

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

        data = self.pipeline.get_data("cadi_res")
        assert np.sum(data) == pytest.approx(0.7158207863548083, rel=1e-2, abs=0.0)
        assert data.shape == (10, 21, 21)

        data = self.pipeline.get_data("cadi_stack")
        assert np.sum(data) == pytest.approx(0.07448334552227256, rel=1e-2, abs=0.0)
        assert data.shape == (1, 21, 21)

    def test_psf_subtraction_pca_single(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_single",
            images_in_tag="science",
            reference_in_tag="science",
            res_mean_tag="res_mean_single",
            res_median_tag="res_median_single",
            res_weighted_tag="res_weighted_single",
            res_rot_mean_clip_tag="res_clip_single",
            res_arr_out_tag="res_arr_single",
            basis_out_tag="basis_single",
            extra_rot=45.0,
            subtract_mean=True,
        )

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

        data = self.pipeline.get_data("res_mean_single")
        assert np.sum(data) == pytest.approx(
            -0.00011857022709778602, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_median_single")
        assert np.sum(data) == pytest.approx(
            -0.002184868916566093, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_weighted_single")
        assert np.sum(data) == pytest.approx(
            0.08102176735226937, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        # data = self.pipeline.get_data('res_clip_single')
        # assert np.sum(data) == pytest.approx(7.09495495339349e-05, rel=self.limit, abs=0.)
        # assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_arr_single1")
        assert np.sum(data) == pytest.approx(
            -0.0002751385418691618, rel=self.limit, abs=0.0
        )
        assert data.shape == (10, 21, 21)

        data = self.pipeline.get_data("basis_single")
        assert np.sum(data) == pytest.approx(
            0.09438697731322143, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_no_mean(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_no_mean",
            images_in_tag="science",
            reference_in_tag="science",
            res_mean_tag="res_mean_no_mean",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_no_mean",
            extra_rot=0.0,
            subtract_mean=False,
        )

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

        data = self.pipeline.get_data("res_mean_no_mean")
        assert np.sum(data) == pytest.approx(
            0.0006081272007585688, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_no_mean")
        assert np.sum(data) == pytest.approx(5.118005177367776, rel=self.limit, abs=0.0)
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_ref(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_ref",
            images_in_tag="science",
            reference_in_tag="reference",
            res_mean_tag="res_mean_ref",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_ref",
            extra_rot=0.0,
            subtract_mean=True,
        )

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

        data = self.pipeline.get_data("res_mean_ref")
        assert np.sum(data) == pytest.approx(
            0.0006330226118859073, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_ref")
        assert np.sum(data) == pytest.approx(
            0.0943869773132221, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_ref_no_mean(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_ref_no_mean",
            images_in_tag="science",
            reference_in_tag="reference",
            res_mean_tag="res_mean_ref_no_mean",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_ref_no_mean",
            extra_rot=0.0,
            subtract_mean=False,
        )

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

        data = self.pipeline.get_data("res_mean_ref_no_mean")
        assert np.sum(data) == pytest.approx(
            0.0006081272007585764, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_ref_no_mean")
        assert np.sum(data) == pytest.approx(5.118005177367774, rel=self.limit, abs=0.0)
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_pca_single_mask(self) -> None:

        pca = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_single_mask",
            images_in_tag="science_prep",
            reference_in_tag="science_prep",
            res_mean_tag="res_mean_single_mask",
            res_median_tag="res_median_single_mask",
            res_weighted_tag="res_weighted_single_mask",
            res_rot_mean_clip_tag="res_clip_single_mask",
            res_arr_out_tag="res_arr_single_mask",
            basis_out_tag="basis_single_mask",
            extra_rot=45.0,
            subtract_mean=True,
        )

        self.pipeline.add_module(pca)
        self.pipeline.run_module("pca_single_mask")

        data = self.pipeline.get_data("res_mean_single_mask")
        assert np.sum(data) == pytest.approx(
            0.00010696166038626307, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_median_single_mask")
        assert np.sum(data) == pytest.approx(
            -0.0021005307611346156, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_weighted_single_mask")
        assert np.sum(data) == pytest.approx(
            0.06014309988789256, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_clip_single_mask")
        # assert np.sum(data) == pytest.approx(9.35120662148806e-05, rel=self.limit, abs=0.)
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("res_arr_single_mask1")
        assert np.sum(data) == pytest.approx(
            0.0006170872862547557, rel=self.limit, abs=0.0
        )
        assert data.shape == (10, 21, 21)

        data = self.pipeline.get_data("basis_single_mask")
        assert np.sum(data) == pytest.approx(
            0.08411251293842359, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_no_mean_mask(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_no_mean_mask",
            images_in_tag="science_prep",
            reference_in_tag="science_prep",
            res_mean_tag="res_mean_no_mean_mask",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_no_mean_mask",
            extra_rot=0.0,
            subtract_mean=False,
        )

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

        data = self.pipeline.get_data("res_mean_no_mean_mask")
        assert np.sum(data) == pytest.approx(
            2.3542359949502915e-05, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_no_mean_mask")
        assert np.sum(data) == pytest.approx(5.655460951633232, rel=self.limit, abs=0.0)
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_ref_mask(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_ref_mask",
            images_in_tag="science_prep",
            reference_in_tag="reference_prep",
            res_mean_tag="res_mean_ref_mask",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_ref_mask",
            extra_rot=0.0,
            subtract_mean=True,
        )

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

        data = self.pipeline.get_data("res_mean_ref_mask")
        assert np.sum(data) == pytest.approx(
            9.400558926815758e-06, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_ref_mask")
        assert np.sum(data) == pytest.approx(
            0.08411251293842326, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_ref_no_mean_mask(self) -> None:

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_ref_no_mean_mask",
            images_in_tag="science_prep",
            reference_in_tag="reference_prep",
            res_mean_tag="res_mean_ref_no_mean_mask",
            res_median_tag=None,
            res_weighted_tag=None,
            res_rot_mean_clip_tag=None,
            res_arr_out_tag=None,
            basis_out_tag="basis_ref_no_mean_mask",
            extra_rot=0.0,
            subtract_mean=False,
        )

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

        data = self.pipeline.get_data("res_mean_ref_no_mean_mask")
        assert np.sum(data) == pytest.approx(
            2.354235994950671e-05, rel=self.limit, abs=0.0
        )
        assert data.shape == (2, 21, 21)

        data = self.pipeline.get_data("basis_ref_no_mean_mask")
        assert np.sum(data) == pytest.approx(5.655460951633233, rel=self.limit, abs=0.0)
        assert data.shape == (2, 21, 21)

    def test_psf_subtraction_pca_multi(self) -> None:

        with h5py.File(self.test_dir + "PynPoint_database.hdf5", "a") as hdf_file:
            hdf_file["config"].attrs["CPU"] = 4

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_multi",
            images_in_tag="science",
            reference_in_tag="science",
            res_mean_tag="res_mean_multi",
            res_median_tag="res_median_multi",
            res_weighted_tag="res_weighted_multi",
            res_rot_mean_clip_tag="res_clip_multi",
            res_arr_out_tag=None,
            basis_out_tag="basis_multi",
            extra_rot=45.0,
            subtract_mean=True,
        )

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

        data_single = self.pipeline.get_data("res_mean_single")
        data_multi = self.pipeline.get_data("res_mean_multi")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("res_median_single")
        data_multi = self.pipeline.get_data("res_median_multi")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("res_weighted_single")
        data_multi = self.pipeline.get_data("res_weighted_multi")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("basis_single")
        data_multi = self.pipeline.get_data("basis_multi")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

    def test_psf_subtraction_pca_multi_mask(self) -> None:

        with h5py.File(self.test_dir + "PynPoint_database.hdf5", "a") as hdf_file:
            hdf_file["config"].attrs["CPU"] = 4

        module = PcaPsfSubtractionModule(
            pca_numbers=range(1, 3),
            name_in="pca_multi_mask",
            images_in_tag="science_prep",
            reference_in_tag="science_prep",
            res_mean_tag="res_mean_multi_mask",
            res_median_tag="res_median_multi_mask",
            res_weighted_tag="res_weighted_multi_mask",
            res_rot_mean_clip_tag="res_clip_multi_mask",
            res_arr_out_tag=None,
            basis_out_tag="basis_multi_mask",
            extra_rot=45.0,
            subtract_mean=True,
        )

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

        data_single = self.pipeline.get_data("res_mean_single_mask")
        data_multi = self.pipeline.get_data("res_mean_multi_mask")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("res_median_single_mask")
        data_multi = self.pipeline.get_data("res_median_multi_mask")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("res_weighted_single_mask")
        data_multi = self.pipeline.get_data("res_weighted_multi_mask")
        assert data_single.shape == data_multi.shape

        assert data_single[data_single > 1e-12] == pytest.approx(
            data_multi[data_multi > 1e-12], rel=self.limit, abs=0.0
        )

        data_single = self.pipeline.get_data("basis_single_mask")
        data_multi = self.pipeline.get_data("basis_multi_mask")
        assert data_single.shape == data_multi.shape
        assert data_single == pytest.approx(data_multi, rel=self.limit, abs=0.0)

    def test_psf_subtraction_len_parang(self) -> None:

        with h5py.File(self.test_dir + "PynPoint_database.hdf5", "a") as hdf_file:
            hdf_file["config"].attrs["CPU"] = 1

        parang = self.pipeline.get_data("header_science/PARANG")
        self.pipeline.set_attribute(
            "science_prep", "PARANG", np.append(parang, 0.0), static=False
        )

        module = PcaPsfSubtractionModule(
            pca_numbers=[
                1,
            ],
            name_in="pca_len_parang",
            images_in_tag="science_prep",
            reference_in_tag="science_prep",
            res_mean_tag="res_mean_len_parang",
            extra_rot=0.0,
        )

        self.pipeline.add_module(module)

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

        assert (
            str(error.value)
            == "The number of images (10) is not equal to the number of "
            "parallactic angles (11)."
        )
