import os

import pytest
import numpy as np

from astropy.io import fits

from pynpoint.core.pypeline import Pypeline
from pynpoint.readwrite.fitsreading import FitsReadingModule
from pynpoint.util.tests import create_config, create_star_data, remove_test_data


class TestFitsReading:

    def setup_class(self) -> None:

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

        create_star_data(self.test_dir+'fits')
        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=['fits'], files=['filenames.dat'])

    def test_fits_reading(self) -> None:

        module = FitsReadingModule(name_in='read1',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=False,
                                   check=True)

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

        data = self.pipeline.get_data('input')
        assert np.sum(data) == pytest.approx(105.54278879805277, rel=self.limit, abs=0.)
        assert data.shape == (10, 11, 11)

    def test_fits_reading_overwrite(self) -> None:

        module = FitsReadingModule(name_in='read2',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=True,
                                   check=True)

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

        data = self.pipeline.get_data('input')
        assert np.sum(data) == pytest.approx(105.54278879805277, rel=self.limit, abs=0.)
        assert data.shape == (10, 11, 11)

    def test_static_not_found(self) -> None:

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

        module = FitsReadingModule(name_in='read3',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=True,
                                   check=True)

        self.pipeline.add_module(module)

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

        assert len(warning) == 2

        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 DIT', static=True)

    def test_static_changing(self) -> None:

        with fits.open(self.test_dir+'fits/images_0.fits') as hdu:
            header = hdu[0].header
            header['HIERARCH ESO DET DIT'] = 0.1
            hdu.writeto(self.test_dir+'fits/images_0.fits', overwrite=True)

        with fits.open(self.test_dir+'fits/images_1.fits') as hdu:
            header = hdu[0].header
            header['HIERARCH ESO DET DIT'] = 0.2
            hdu.writeto(self.test_dir+'fits/images_1.fits', overwrite=True)

        module = FitsReadingModule(name_in='read4',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=True,
                                   check=True)

        self.pipeline.add_module(module)

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

        assert len(warning) == 1

        assert warning[0].message.args[0] == f'Static attribute ESO DET DIT has changed. ' \
                                             f'Possibly the current file {self.test_dir}fits/' \
                                             f'images_1.fits does not belong to the data set ' \
                                             f'\'input\'. Attribute value is updated.'

    def test_header_attribute(self) -> None:

        with fits.open(self.test_dir+'fits/images_0.fits') as hdu:
            header = hdu[0].header
            header['PARANG'] = 1.0
            hdu.writeto(self.test_dir+'fits/images_0.fits', overwrite=True)

        with fits.open(self.test_dir+'fits/images_1.fits') as hdu:
            header = hdu[0].header
            header['PARANG'] = 2.0
            header['HIERARCH ESO DET DIT'] = 0.1
            hdu.writeto(self.test_dir+'fits/images_1.fits', overwrite=True)

        module = FitsReadingModule(name_in='read5',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=True,
                                   check=True)

        self.pipeline.add_module(module)
        self.pipeline.run_module('read5')

    def test_non_static_not_found(self) -> None:

        self.pipeline.set_attribute('config', 'DIT', 'None', static=True)

        for i in range(2):
            with fits.open(f'{self.test_dir}/fits/images_{i}.fits') as hdu:
                header = hdu[0].header
                del header['HIERARCH ESO DET DIT']
                del header['HIERARCH ESO DET EXP NO']
                hdu.writeto(f'{self.test_dir}/fits/images_{i}.fits', overwrite=True)

        module = FitsReadingModule(name_in='read6',
                                   input_dir=self.test_dir+'fits',
                                   image_tag='input',
                                   overwrite=True,
                                   check=True)

        self.pipeline.add_module(module)

        with pytest.warns(UserWarning) as warning:
            self.pipeline.run_module('read6')

        assert len(warning) == 2

        for item in warning:
            assert item.message.args[0] == 'Non-static attribute EXP_NO (=ESO DET EXP NO) not ' \
                                           'found in the FITS header.'

    def test_fits_read_files(self) -> None:

        module = FitsReadingModule(name_in='read7',
                                   input_dir=None,
                                   image_tag='files',
                                   overwrite=False,
                                   check=True,
                                   filenames=[self.test_dir+'fits/images_0.fits',
                                              self.test_dir+'fits/images_1.fits'])

        self.pipeline.add_module(module)

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

        assert len(warning) == 2

        for item in warning:
            assert item.message.args[0] == 'Non-static attribute EXP_NO (=ESO DET EXP NO) not ' \
                                           'found in the FITS header.'

        data = self.pipeline.get_data('files')
        assert np.sum(data) == pytest.approx(105.54278879805277, rel=self.limit, abs=0.)
        assert data.shape == (10, 11, 11)

    def test_fits_read_textfile(self) -> None:

        with open(self.test_dir+'filenames.dat', 'w') as file_obj:
            file_obj.write(self.test_dir+'fits/images_0.fits\n')
            file_obj.write(self.test_dir+'fits/images_1.fits\n')

        module = FitsReadingModule(name_in='read8',
                                   input_dir=None,
                                   image_tag='files',
                                   overwrite=True,
                                   check=True,
                                   filenames=self.test_dir+'filenames.dat')

        self.pipeline.add_module(module)

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

        assert len(warning) == 2

        for item in warning:
            assert item.message.args[0] == 'Non-static attribute EXP_NO (=ESO DET EXP NO) not ' \
                                           'found in the FITS header.'

        data = self.pipeline.get_data('files')
        assert np.sum(data) == pytest.approx(105.54278879805277, rel=self.limit, abs=0.)
        assert data.shape == (10, 11, 11)

    def test_fits_read_files_exists(self) -> None:

        module = FitsReadingModule(name_in='read9',
                                   input_dir=None,
                                   image_tag='files',
                                   overwrite=True,
                                   check=True,
                                   filenames=[f'{self.test_dir}fits/images_0.fits',
                                              f'{self.test_dir}fits/images_2.fits'])

        self.pipeline.add_module(module)

        with pytest.raises(ValueError) as error:
            self.pipeline.run_module('read9')

        assert str(error.value) == f'The file {self.test_dir}fits/images_2.fits does not exist. ' \
                                   f'Please check that the path is correct.'

    def test_fits_read_textfile_exists(self) -> None:

        with open(self.test_dir+'filenames.dat', 'w') as file_obj:
            file_obj.write(self.test_dir+'fits/images_0.fits\n')
            file_obj.write(self.test_dir+'fits/images_2.fits\n')

        module = FitsReadingModule(name_in='read10',
                                   input_dir=None,
                                   image_tag='files',
                                   overwrite=True,
                                   check=True,
                                   filenames=self.test_dir+'filenames.dat')

        self.pipeline.add_module(module)

        with pytest.raises(ValueError) as error:
            self.pipeline.run_module('read10')

        assert str(error.value) == f'The file {self.test_dir}fits/images_2.fits does not exist. ' \
                                   f'Please check that the path is correct.'
