import os
from os.path import join as pjoin
import numbers
import warnings
import pickle

import pytest
import numpy as np
from numpy.testing import (
    assert_almost_equal,
    assert_,
    assert_equal,
    assert_array_less,
    assert_allclose,
)
import h5py

import refnx.reduce.platypusnexus as plp
from refnx.reduce import (
    PlatypusReduce,
    PlatypusNexus,
    SpatzNexus,
    basename_datafile,
    catalogue,
    SpinChannel,
    SpinSet,
)
from refnx.reduce.peak_utils import gauss
from refnx.reduce.platypusnexus import (
    fore_back_region,
    EXTENT_MULT,
    PIXEL_OFFSET,
    create_detector_norm,
    Catalogue,
    PlatypusCatalogue,
)


class TestSpinSet(object):
    @pytest.mark.usefixtures("no_data_directory")
    @pytest.fixture(autouse=True)
    def setup_method(self, tmpdir, data_directory):
        self.pth = pjoin(data_directory, "reduce", "PNR_files")

        def fpath(f):
            return pjoin(self.pth, f)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)

            self.spinset = SpinSet(
                down_down=fpath("PLP0008864.nx.hdf"),
                up_up=fpath("PLP0008861.nx.hdf"),
                down_up=fpath("PLP0008863.nx.hdf"),
                up_down=fpath("PLP0008862.nx.hdf"),
            )

            self.spinset_3 = SpinSet(
                down_down=fpath("PLP0008864.nx.hdf"),
                up_up=fpath("PLP0008861.nx.hdf"),
                down_up=fpath("PLP0008863.nx.hdf"),
            )
            self.spinset_2 = SpinSet(
                down_down=fpath("PLP0008864.nx.hdf"),
                up_up=fpath("PLP0008861.nx.hdf"),
            )

        self.cwd = os.getcwd()

        self.tmpdir = tmpdir.strpath
        os.chdir(self.tmpdir)
        return 0

    def test_spin_channels(self):
        assert self.spinset.spin_channels == [(0, 0), (0, 1), (1, 0), (1, 1)]
        assert self.spinset_3.spin_channels == [(0, 0), (0, 1), None, (1, 1)]
        assert self.spinset_2.spin_channels == [(0, 0), None, None, (1, 1)]

    def test_process(self):
        self.spinset.process()
        assert self.spinset.dd.processed_spectrum
        assert self.spinset.du.processed_spectrum
        assert self.spinset.ud.processed_spectrum
        assert self.spinset.uu.processed_spectrum

        self.spinset_3.process()
        assert self.spinset_3.dd.processed_spectrum
        assert self.spinset_3.du.processed_spectrum
        assert self.spinset_3.ud is None
        assert self.spinset_3.uu.processed_spectrum

        self.spinset_2.process()
        assert self.spinset_2.dd.processed_spectrum
        assert self.spinset_2.du is None
        assert self.spinset_2.ud is None
        assert self.spinset_2.uu.processed_spectrum

    def test_processing_different_reduction_options(self):
        # Check every combination of spin channel and reduction option that
        # determines the resulting wavelength axis to make sure errors
        # are raised appropriately

        standard_opts = dict(
            lo_wavelength=2.5,
            hi_wavelength=12.5,
            rebin_percent=3,
        )
        for spin_set in [self.spinset, self.spinset_3, self.spinset_2]:
            for sc in ["dd", "du", "ud", "uu"]:
                if spin_set.channels[sc] is None:
                    continue
                else:
                    for option in [
                        "lo_wavelength",
                        "hi_wavelength",
                        "rebin_percent",
                    ]:
                        spin_set.sc_opts[sc].update({option: 5})
                        with pytest.raises(ValueError):
                            spin_set.process()
                        # Reset dd_opts
                        spin_set.sc_opts[sc].update(standard_opts)


class TestPlatypusNexus(object):
    @pytest.mark.usefixtures("no_data_directory")
    @pytest.fixture(autouse=True)
    def setup_method(self, tmpdir, data_directory):
        self.pth = pjoin(data_directory, "reduce")

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            self.f113 = PlatypusNexus(pjoin(self.pth, "PLP0011613.nx.hdf"))
            self.f641 = PlatypusNexus(pjoin(self.pth, "PLP0011641.nx.hdf"))

            # These PNR datasets all have different flipper settings
            self.f8861 = PlatypusNexus(
                pjoin(self.pth, "PNR_files/PLP0008861.nx.hdf")
            )
            self.f8862 = PlatypusNexus(
                pjoin(self.pth, "PNR_files/PLP0008862.nx.hdf")
            )
            self.f8863 = PlatypusNexus(
                pjoin(self.pth, "PNR_files/PLP0008863.nx.hdf")
            )
            self.f8864 = PlatypusNexus(
                pjoin(self.pth, "PNR_files/PLP0008864.nx.hdf")
            )

        self.cwd = os.getcwd()

        self.tmpdir = tmpdir.strpath
        os.chdir(self.tmpdir)
        return 0

    def teardown_method(self):
        os.chdir(self.cwd)

    def test_chod(self):
        flight_length = self.f113.chod()
        assert_almost_equal(flight_length[0], 7141.413818359375)
        assert_almost_equal(flight_length[1], 808)

        flight_length = self.f641.chod(omega=1.8, twotheta=3.6)
        assert_almost_equal(flight_length[0], 7146.3567785516016)
        assert_almost_equal(flight_length[1], 808)

    def test_phase_angle(self):
        # TODO. Add more tests where the phase_angle isn't zero.
        phase_angle, master_opening = self.f641.phase_angle()
        assert_almost_equal(phase_angle, 0)
        assert_almost_equal(master_opening, 60)

        assert self.f641.cat.t_offset is None

    def test_background_subtract_line(self):
        # checked each step of the background subtraction with IGOR
        # so this test background correction should be correct.

        # create some test data
        xvals = np.linspace(-10, 10, 201)
        yvals = np.ceil(gauss(xvals, 0, 100, 0, 1) + 2 * xvals + 30)

        # add some reproducible random noise
        np.random.seed(1)
        yvals += np.sqrt(yvals) * np.random.randn(yvals.size)
        yvals_sd = np.sqrt(yvals)

        mask = np.zeros(201, bool)
        mask[30:70] = True
        mask[130:160] = True

        profile, profile_sd = plp.background_subtract_line(
            yvals, yvals_sd, mask
        )

        verified_data = np.load(pjoin(self.pth, "background_subtract.npy"))

        assert_almost_equal(verified_data, np.c_[profile, profile_sd])

    def test_find_specular_ridge(self):
        xvals = np.linspace(-10, 10, 201)
        yvals = np.ceil(gauss(xvals, 0, 1000, 0, 1))
        detector = np.repeat(yvals[:, np.newaxis], 1000, axis=1).T
        detector_sd = np.sqrt(detector)
        output = plp.find_specular_ridge(
            detector[np.newaxis, :], detector_sd[np.newaxis, :]
        )
        assert_(len(output) == 5)
        assert_almost_equal(output[0][0], 100)

    def test_pickle(self):
        # can we pickle a PlatypusNexus object?
        self.f8864.process()
        pkl = pickle.dumps(self.f8864)
        o = pickle.loads(pkl)
        ops, f8864 = o.processed_spectrum, self.f8864.processed_spectrum
        assert_allclose(ops["m_spec"], f8864["m_spec"])

    def test_background_subtract(self):
        # create some test data
        xvals = np.linspace(-10, 10, 201)
        yvals = np.ceil(gauss(xvals, 0, 100, 0, 1) + 2 * xvals + 30)

        # add some reproducible random noise
        np.random.seed(1)
        yvals += np.sqrt(yvals) * np.random.randn(yvals.size)
        yvals_sd = np.sqrt(yvals)

        # now make an (N, T, Y) detector image
        n_tbins = 10
        detector = np.repeat(yvals, n_tbins).reshape(xvals.size, n_tbins).T
        detector_sd = (
            np.repeat(yvals_sd, n_tbins).reshape(xvals.size, n_tbins).T
        )
        detector = detector.reshape(1, n_tbins, xvals.size)
        detector_sd = detector_sd.reshape(1, n_tbins, xvals.size)

        mask = np.zeros((1, n_tbins, 201), bool)
        mask[:, :, 30:70] = True
        mask[:, :, 130:160] = True

        det_bkg, detSD_bkg = plp.background_subtract(
            detector, detector_sd, mask
        )

        # each of the (N, T) entries should have the same background subtracted
        # entries
        verified_data = np.load(pjoin(self.pth, "background_subtract.npy"))

        it = np.nditer(detector, flags=["multi_index"])
        it.remove_axis(2)
        while not it.finished:
            profile = det_bkg[it.multi_index]
            profile_sd = detSD_bkg[it.multi_index]
            assert_almost_equal(verified_data, np.c_[profile, profile_sd])
            it.iternext()

    def test_calculate_bins(self):
        bins = plp.calculate_wavelength_bins(2.0, 18, 2.0)
        assert_almost_equal(bins[0], 1.98)
        assert_almost_equal(bins[-1], 18.18)

    def test_event(self):
        # When you use event mode processing, make sure the right amount of
        # spectra are created
        out = self.f641.process(eventmode=[0, 900, 1800], integrate=0)
        assert_(np.size(out[1], axis=0) == 2)

    def test_event_folder(self):
        self.f641.process(
            eventmode=[0, 900, 1800], integrate=0, event_folder=self.pth
        )

    def test_multiple_acquisitions(self):
        """
        TODO: add a dataset which has multiple spectra in it, and make sure it
        processes.
        """
        pass

    def test_reduction_runs(self):
        # just check it runs
        self.f641.process()

        # check that event mode reduction gives the same output as non-event
        # mode reduction.
        spectrum0 = self.f641.process(direct=False)
        assert "reduction_options" in self.f641.processed_spectrum

        spectrum1 = self.f641.process(direct=False, eventmode=[], integrate=0)
        assert_allclose(spectrum0[1][0], spectrum1[1][0], rtol=0.001)

        # check that the wavelength resolution is roughly right, between 7 and
        # 8%.
        res = (
            self.f641.processed_spectrum["m_lambda_fwhm"][0]
            / self.f641.processed_spectrum["m_lambda"][0]
        )
        assert_array_less(res, np.ones_like(res) * 0.08)
        assert_array_less(np.ones_like(res) * 0.07, res)

    def test_save_spectrum(self):
        # test saving spectrum
        self.f113.process()

        # can save the spectra by supplying a filename
        self.f113.write_spectrum_xml(pjoin(self.tmpdir, "test.xml"))
        self.f113.write_spectrum_dat(pjoin(self.tmpdir, "test.dat"))

        # can save by supplying file handle:
        with open(pjoin(self.tmpdir, "test.xml"), "wb") as f:
            self.f113.write_spectrum_xml(f)

    def test_accumulate_files(self):
        fnames = ["PLP0000708.nx.hdf", "PLP0000709.nx.hdf"]
        pths = [pjoin(self.pth, fname) for fname in fnames]
        plp.accumulate_HDF_files(pths)
        f8, f9, fadd = None, None, None

        try:
            f8 = h5py.File(pjoin(self.pth, "PLP0000708.nx.hdf"), "r")
            f9 = h5py.File(pjoin(self.pth, "PLP0000709.nx.hdf"), "r")
            fadd = h5py.File(pjoin(self.tmpdir, "ADD_PLP0000708.nx.hdf"), "r")

            f8d = f8["entry1/data/hmm"][0]
            f9d = f9["entry1/data/hmm"][0]
            faddd = fadd["entry1/data/hmm"][0]
            assert_equal(faddd, f8d + f9d)
        finally:
            if f8 is not None:
                f8.close()
            if f9 is not None:
                f9.close()
            if fadd is not None:
                fadd.close()

    def test_accumulate_files_reduce(self):
        # test by adding a file to itself. Should have smaller stats
        fnames = ["PLP0000708.nx.hdf", "PLP0000708.nx.hdf"]
        pths = [pjoin(self.pth, fname) for fname in fnames]
        plp.accumulate_HDF_files(pths)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            # it should be processable
            fadd = PlatypusNexus(pjoin(os.getcwd(), "ADD_PLP0000708.nx.hdf"))
            fadd.process()

            # it should also be reduceable
            reducer = PlatypusReduce(pjoin(self.pth, "PLP0000711.nx.hdf"))

            datasets, reduced = reducer.reduce(
                pjoin(os.getcwd(), "ADD_PLP0000708.nx.hdf")
            )
            assert_("y" in reduced)

            # the error bars should be smaller
            datasets2, reduced2 = reducer.reduce(
                pjoin(self.pth, "PLP0000708.nx.hdf")
            )

            assert_(np.all(reduced["y_err"] < reduced2["y_err"]))

    def test_manual_beam_find(self):
        # you can specify a function that finds where the specular ridge is.
        def manual_beam_find(detector, detector_sd, name):
            beam_centre = np.zeros(len(detector))
            beam_sd = np.zeros(len(detector))
            beam_centre += 50
            beam_sd += 5
            return beam_centre, beam_sd, np.array([40]), np.array([60]), [[]]

        # the manual beam find is only mandatory when peak_pos == -1.
        # the beam_sd is much larger than the beam divergence, so a warning
        # should be raised.
        with pytest.warns(UserWarning):
            self.f113.process(manual_beam_find=manual_beam_find, peak_pos=-1)
        assert_equal(self.f113.processed_spectrum["m_beampos"][0], 50)

        # manual beam finding also specifies the lower and upper pixel of the
        # foreground
        assert_equal(self.f113.processed_spectrum["lopx"][0], 40)
        assert_equal(self.f113.processed_spectrum["hipx"][0], 60)

    def test_fore_back_region(self):
        # calculation of foreground and background regions is done correctly
        centres = np.array([100.0, 90.0])
        sd = np.array([5.0, 11.5])
        lopx, hipx, background_pixels = fore_back_region(centres, sd)

        assert_(len(lopx) == 2)
        assert_(len(hipx) == 2)
        assert_(len(background_pixels) == 2)
        assert_(isinstance(lopx[0], numbers.Integral))
        assert_(isinstance(hipx[0], numbers.Integral))

        calc_lower = np.floor(centres - sd * EXTENT_MULT)
        assert_equal(lopx, calc_lower)
        calc_higher = np.ceil(centres + sd * EXTENT_MULT)
        assert_equal(hipx, calc_higher)

        y1 = np.atleast_1d(np.round(lopx - PIXEL_OFFSET).astype("int"))
        y0 = np.atleast_1d(
            np.round(lopx - PIXEL_OFFSET - EXTENT_MULT * sd).astype("int")
        )

        y2 = np.atleast_1d(np.round(hipx + PIXEL_OFFSET).astype("int"))
        y3 = np.atleast_1d(
            np.round(hipx + PIXEL_OFFSET + EXTENT_MULT * sd).astype("int")
        )

        bp = np.r_[np.arange(y0[0], y1[0] + 1), np.arange(y2[0], y3[0] + 1)]

        assert_equal(bp, background_pixels[0])

    def test_basename_datafile(self):
        # check that the right basename is returned
        pth = "a/b/c.nx.hdf"
        assert_(basename_datafile(pth) == "c")

        pth = "c.nx.hdf"
        assert_(basename_datafile(pth) == "c")

    def test_floodfield_correction(self):
        # check that flood field calculation works
        # the values were worked out by hand on a randomly
        # generated array
        test_norm = np.array(
            [1.12290503, 1.23743017, 0.8603352, 0.70111732, 1.07821229]
        )
        test_norm_sd = np.array(
            [0.05600541, 0.05879208, 0.04902215, 0.04425413, 0.05487956]
        )

        fname = os.path.join(self.pth, "flood.h5")
        with h5py.File(fname, "r") as f:
            norm, norm_sd = create_detector_norm(f, 3.5, -3.5, axis=3)

            assert_almost_equal(norm, test_norm, 6)
            assert_almost_equal(norm_sd, test_norm_sd, 6)

            norm, norm_sd = create_detector_norm(f, -3.5, 3.5, axis=3)
            assert_almost_equal(norm, test_norm, 6)
            assert_almost_equal(norm_sd, test_norm_sd, 6)

    def test_PNR_metadata(self):
        self.f8861.process()
        self.f8862.process()
        self.f8863.process()
        self.f8864.process()

        # Check that we can read all of the flipper statuses in files

        # Flipper 1 on, flipper 2 on
        assert_almost_equal(self.f8861.cat.cat["pol_flip_current"], 5.0)
        assert_almost_equal(self.f8861.cat.cat["anal_flip_current"], 4.5)

        # Flipper 1 on, flipper 2 off
        assert_almost_equal(self.f8862.cat.cat["anal_flip_current"], 0)
        assert_almost_equal(self.f8862.cat.cat["pol_flip_current"], 5.0)

        # Flipper 1 off, flipper 2 on
        assert_almost_equal(self.f8863.cat.cat["anal_flip_current"], 4.5)
        assert_almost_equal(self.f8863.cat.cat["pol_flip_current"], 0)

        # Flipper 1 off, flipper 2 off
        assert_almost_equal(self.f8864.cat.cat["anal_flip_current"], 0)
        assert_almost_equal(self.f8864.cat.cat["pol_flip_current"], 0)

        # Check SpinChannel for each file
        assert self.f8861.spin_state == SpinChannel.UP_UP
        assert self.f8862.spin_state == SpinChannel.UP_DOWN
        assert self.f8863.spin_state == SpinChannel.DOWN_UP
        assert self.f8864.spin_state == SpinChannel.DOWN_DOWN

        # test spin channel setting
        # non spin analysed. mode is POL, not POLANAL
        pn = PlatypusNexus(pjoin(self.pth, "PLP0016427.nx.hdf"))
        assert pn.spin_state == SpinChannel.UP_UP
        pn = PlatypusNexus(pjoin(self.pth, "PLP0016426.nx.hdf"))
        assert pn.spin_state == SpinChannel.DOWN_DOWN

    def test_PNR_magnet_read(self):
        self.f8861.process()
        # Check magnetic field sensors
        assert_almost_equal(self.f8861.cat.cat["magnet_current_set"], 0)
        assert_almost_equal(self.f8861.cat.cat["magnet_output_current"], 0.001)


class TestSpatzNexus:
    @pytest.mark.usefixtures("no_data_directory")
    @pytest.fixture(autouse=True)
    def setup_method(self, tmpdir, data_directory):
        self.pth = pjoin(data_directory, "reduce")

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            self.f342 = SpatzNexus(pjoin(self.pth, "SPZ0000342.nx.hdf"))
        self.cwd = os.getcwd()

        self.tmpdir = tmpdir.strpath
        os.chdir(self.tmpdir)
        return 0

    def teardown_method(self):
        os.chdir(self.cwd)

    def test_chod(self):
        flight_length = self.f342.chod()
        assert_almost_equal(flight_length[0], 8062.0232)
        assert_almost_equal(flight_length[1], 479.9536, decimal=4)

    def test_phase_angle(self):
        assert_allclose(self.f342.cat.master_phase_offset, -25.90)
        assert self.f342.cat.master == 1
        assert self.f342.cat.slave == 2
        assert self.f342.cat.t_offset is None
        assert_allclose(self.f342.cat.frequency, 25)
        assert_allclose(self.f342.cat.phase, 34.22)
        assert_allclose(self.f342.cat.poff_c2_slave_1_master[0], -0.22)

        phase_angle, master_opening = self.f342.phase_angle()
        # chopper 1 has an opening of 26 degrees
        assert_almost_equal(master_opening, 26)
        assert_allclose(phase_angle, 0, atol=1e-5)

        toff = self.f342.time_offset(-25.90, 26, 25, 0.0, None, None, None)
        assert_allclose(toff, 5.5555555555555)

        toff = self.f342.time_offset(
            -25.90,
            26,
            25,
            0.0,
            None,
            None,
            None,
            t_offset=1438.8888888888888,
        )
        assert_allclose(toff, 0, atol=1e-12)


@pytest.mark.usefixtures("no_data_directory")
def test_catalogue(data_directory):
    pth = pjoin(data_directory, "reduce")
    catalogue(0, 10000000, data_folder=pth, prefix="PLP")
    catalogue(0, 10000000, data_folder=pth, prefix="SPZ")


@pytest.mark.usefixtures("no_data_directory")
def test_Catalogue_pickle(data_directory):
    pth = pjoin(data_directory, "reduce")
    f113 = pjoin(pth, "PLP0011613.nx.hdf")
    with h5py.File(f113) as f:
        c = PlatypusCatalogue(f)

    assert c.datafile_number == 11613
    assert c.prefix == "PLP"
    pkl = pickle.dumps(c)
    pickle.loads(pkl)
