import os
import numpy as np
import pytest
import warnings
import matplotlib.pyplot as plt
import scipy.special
from astropy.io import fits
from stingray import Lightcurve
from stingray import Crossspectrum, AveragedCrossspectrum
from stingray.crossspectrum import cospectra_pvalue
from stingray.crossspectrum import (
    normalize_crossspectrum, normalize_crossspectrum_gauss)
from stingray.crossspectrum import coherence, time_lag
from stingray import StingrayError
from stingray.simulator import Simulator
from stingray.fourier import poisson_level, raw_coherence

from stingray.events import EventList
import copy

np.random.seed(20160528)
curdir = os.path.abspath(os.path.dirname(__file__))
datadir = os.path.join(curdir, "data")


def avg_cdf_two_spectra(x):

    prefac = 0.25

    if x >= 0:
        fac1 = 2 * scipy.special.gamma(2) - scipy.special.gammaincc(2, 2 * x)
        fac2 = 2. * scipy.special.gamma(1) - scipy.special.gammaincc(1, 2 * x)
    else:
        fac1 = scipy.special.gammaincc(2, -2 * x)
        fac2 = scipy.special.gammaincc(1, -2 * x)

    return prefac * (fac1 + fac2)


class TestClassicalPvalue(object):

    def test_pval_returns_float_when_float_input(self):
        power = 1.0
        nspec = 1.0
        pval = cospectra_pvalue(power, nspec)
        assert isinstance(pval, float)

    def test_pval_returns_iterable_when_iterable_input(self):
        power = [0, 1, 2]
        nspec = 1.0
        pval = cospectra_pvalue(power, nspec)
        assert isinstance(pval, np.ndarray)
        assert len(pval) == len(power)

    def test_pval_returns_iterable_when_iterable_input_nspec2(self):
        power = [0, 1, 2]
        nspec = 2
        # It will use the formulation by Huppenkothen & Bachetti
        pval = cospectra_pvalue(power, nspec)
        assert isinstance(pval, np.ndarray)
        assert len(pval) == len(power)

    def test_pval_fails_if_single_power_infinite(self):
        power = np.inf
        nspec = 1
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_single_power_nan(self):
        power = np.nan
        nspec = 1
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_multiple_powers_nan(self):
        power = [1, np.nan, 2.0]
        nspec = 1
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_multiple_powers_inf(self):
        power = [1, 2.0, np.inf]
        nspec = 1
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_nspec_zero(self):
        power = 1.0
        nspec = 0
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_nspec_negative(self):
        power = 1.0
        nspec = -10
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_pval_fails_if_nspec_not_integer(self):
        power = 1.0
        nspec = 1.5
        with pytest.raises(ValueError):
            pval = cospectra_pvalue(power, nspec)

    def test_single_spectrum(self):
        # the Laplace distribution is symmetric around
        # 0, so a power of 0 should return p=0.5
        power = 0.0
        nspec = 1
        assert cospectra_pvalue(power, nspec) == 0.5

    def test_single_spectrum_with_positive_power(self):
        """
        Because the Laplace distribution is always symmetric
        around zero, let's do a second version where I look
        for a different number.
        """
        power = 0.69314718055
        nspec = 1
        assert np.isclose(cospectra_pvalue(power, nspec), 0.25)

    def test_two_averaged_spectra(self):
        """
        For nspec=2, I can derive this by hand:
        """
        power = 1.0
        nspec = 2
        manual_pval = 1.0 - avg_cdf_two_spectra(power)
        assert np.isclose(cospectra_pvalue(power, nspec), manual_pval)

    def test_sixty_spectra(self):
        power = 1.0
        nspec = 60
        gauss = scipy.stats.norm(0, np.sqrt(2/(nspec+1)))
        pval_theory = gauss.sf(power)
        assert np.isclose(cospectra_pvalue(power, nspec), pval_theory)


class TestAveragedCrossspectrumEvents(object):

    def setup_class(self):
        tstart = 0.0
        tend = 1.0
        self.dt = np.longdouble(0.0001)
        segment_size = 1
        self.segment_size = segment_size
        N = np.rint(segment_size / self.dt).astype(int)
        # adjust dt
        self.dt = segment_size / N

        times1 = np.sort(np.random.uniform(tstart, tend, 1000000))
        times2 = np.sort(np.random.uniform(tstart, tend, 1000000))
        gti = np.array([[tstart, tend]])

        self.events1 = EventList(times1, gti=gti)
        self.events2 = EventList(times2, gti=gti)

        self.cs = Crossspectrum(self.events1, self.events2, dt=self.dt, norm="none")

        self.acs = AveragedCrossspectrum(self.events1.to_lc(self.dt),
                                         self.events2.to_lc(self.dt), silent=True,
                                         segment_size=segment_size, dt=self.dt, norm='none',
                                         power_type="all")
        self.lc1, self.lc2 = self.events1, self.events2

    @pytest.mark.parametrize("norm", ["leahy", "frac", "abs", "none"])
    def test_common_mean_gives_comparable_scatter(self, norm):
        acs = AveragedCrossspectrum(
             self.events1, self.events2, dt=self.dt, silent=True,
             segment_size=self.segment_size, norm=norm, power_type="real",
             use_common_mean=False)
        acs_comm = AveragedCrossspectrum(
             self.events1, self.events2, dt=self.dt, silent=True,
             segment_size=self.segment_size, norm=norm, power_type="real",
             use_common_mean=True)

        assert np.isclose(acs_comm.power.std(), acs.power.std(), rtol=0.1)

    @pytest.mark.parametrize("use_common_mean", [True, False])
    @pytest.mark.parametrize("legacy", [True, False])
    def test_leahy_correct_for_multiple(self, legacy, use_common_mean):
        n = 30
        lc_all = []
        for i in range(n):
            time = np.arange(0.0, 10.0, 10. / 10000)
            counts = np.random.poisson(1000, size=time.size)
            lc = Lightcurve(time, counts)
            lc_all.append(lc)

        ps = AveragedCrossspectrum(lc_all, lc_all, 1.0, norm="leahy", legacy=legacy,
                                   use_common_mean=use_common_mean)

        assert ps.m == 300
        assert np.isclose(np.mean(ps.power), 2.0, atol=1e-2, rtol=1e-2)
        assert np.isclose(np.std(ps.power), 2.0 / np.sqrt(ps.m), atol=0.1,
                          rtol=0.1)

    def test_from_events_works_cs(self):
        lccs = Crossspectrum.from_events(self.events1, self.events2,
                                         dt=self.dt, norm='none', silent=True)
        power1 = lccs.power.real
        power2 = self.cs.power.real
        assert np.allclose(power1, power2, rtol=0.01)
        lag1 = lccs.time_lag()
        lag2 = self.cs.time_lag()
        assert np.allclose(lag1, lag2)
        assert lccs.power_err is not None

    @pytest.mark.parametrize("legacy", [True, False])
    def test_internal_from_events_works_acs(self, legacy):
        lccs = AveragedCrossspectrum(self.events1, self.events2,
                                     segment_size=1, dt=self.dt, norm='none', silent=True, legacy=legacy)
        power1 = lccs.power.real
        power2 = self.acs.power.real
        assert np.allclose(power1, power2, rtol=0.01)
        lag1, lag1_e = lccs.time_lag()
        lag2, lag2_e = self.acs.time_lag()
        assert np.allclose(lag1, lag2)
        good = ~np.isnan(lag2_e)
        assert np.allclose(lag1_e[good], lag2_e[good])
        assert lccs.power_err is not None

    def test_from_events_works_acs(self):
        lccs = AveragedCrossspectrum.from_events(self.events1, self.events2,
                                                 segment_size=1, dt=self.dt, norm='none', silent=True)
        power1 = lccs.power.real
        power2 = self.acs.power.real
        assert np.allclose(power1, power2, rtol=0.01)
        lag1, lag1_e = lccs.time_lag()
        lag2, lag2_e = self.acs.time_lag()
        assert np.allclose(lag1, lag2)
        assert np.allclose(lag1_e, lag2_e, equal_nan=True)
        assert lccs.power_err is not None

    def test_from_lc_iter_works(self):
        lccs = AveragedCrossspectrum.from_lc_iterable(
            self.events1.to_lc_iter(self.dt, self.segment_size),
            self.events2.to_lc_iter(self.dt, self.segment_size),
            segment_size=self.segment_size, dt=self.dt, norm='none', silent=True)
        power1 = lccs.power.real
        power2 = self.acs.power.real
        assert np.allclose(power1, power2, rtol=0.01)

    def test_from_lc_iter_with_err_works(self):
        def iter_lc_with_errs(iter_lc):
            for lc in iter_lc:
                lc._counts_err = np.zeros_like(lc.counts) + lc.counts.mean()**0.5
                yield lc

        lccs = AveragedCrossspectrum.from_lc_iterable(
            iter_lc_with_errs(self.events1.to_lc_iter(self.dt, self.segment_size)),
            iter_lc_with_errs(self.events2.to_lc_iter(self.dt, self.segment_size)),
            segment_size=self.segment_size, dt=self.dt, norm='none', silent=True)
        power1 = lccs.power.real
        power2 = self.acs.power.real
        assert np.allclose(power1, power2, rtol=0.01)

    def test_from_lc_iter_counts_only_works(self):
        def iter_lc_counts_only(iter_lc):
            for lc in iter_lc:
                yield lc.counts

        lccs = AveragedCrossspectrum.from_lc_iterable(
            iter_lc_counts_only(self.events1.to_lc_iter(self.dt, self.segment_size)),
            iter_lc_counts_only(self.events2.to_lc_iter(self.dt, self.segment_size)),
            segment_size=self.segment_size, dt=self.dt, norm='none', silent=True)
        power1 = lccs.power.real
        power2 = self.acs.power.real
        assert np.allclose(power1, power2, rtol=0.01)

    def test_from_time_array_works_with_memmap(self):
        with fits.open(os.path.join(datadir, "monol_testA.evt"), memmap=True) as hdul:
            times1 = hdul[1].data["TIME"]

            gti = np.array([[hdul[2].data["START"][0], hdul[2].data["STOP"][0]]])

            times2 = np.random.uniform(gti[0, 0], gti[0, 1], 1000)

            _ = AveragedCrossspectrum.from_time_array(
                times1, times2, segment_size=128, dt=self.dt, gti=gti, norm='none',
                use_common_mean=False)

    @pytest.mark.parametrize("norm", ["frac", "abs", "none", "leahy"])
    def test_from_lc_with_err_works(self, norm):
        lc1 = self.events1.to_lc(self.dt)
        lc2 = self.events2.to_lc(self.dt)
        lc1._counts_err = np.sqrt(lc1.counts.mean()) + np.zeros_like(lc1.counts)
        lc2._counts_err = np.sqrt(lc2.counts.mean()) + np.zeros_like(lc2.counts)
        pds = AveragedCrossspectrum.from_lightcurve(
            lc1, lc2,
            segment_size=self.segment_size, norm=norm)
        pds_ev = AveragedCrossspectrum.from_events(
            self.events1, self.events2,
            segment_size=self.segment_size, dt=self.dt, norm=norm)
        for attr in ["power", "freq", "m", "n", "nphots1", "nphots2", "segment_size"]:
            assert np.allclose(getattr(pds, attr), getattr(pds_ev, attr))

    def test_it_works_with_events(self):
        lc1 = self.events1.to_lc(self.dt)
        lc2 = self.events2.to_lc(self.dt)
        lccs = Crossspectrum(lc1, lc2, norm="none")
        assert np.allclose(lccs.power, self.cs.power)

    def test_no_segment_size(self):
        with pytest.raises(ValueError):
            cs = AveragedCrossspectrum(self.lc1, self.lc2, dt=self.dt)

    def test_init_with_norm_not_str(self):
        with pytest.raises(TypeError):
            cs = AveragedCrossspectrum(self.lc1, self.lc2, segment_size=1,
                                       norm=1, dt=self.dt)

    def test_init_with_invalid_norm(self):
        with pytest.raises(ValueError):
            cs = AveragedCrossspectrum(self.lc1, self.lc2, segment_size=1,
                                       norm='frabs', dt=self.dt)

    def test_init_with_inifite_segment_size(self):
        with pytest.raises(ValueError):
            cs = AveragedCrossspectrum(self.lc1, self.lc2,
                                       segment_size=np.inf, dt=self.dt)

    def test_coherence(self):
        with pytest.warns(UserWarning) as w:
            coh = self.acs.coherence()
        assert len(coh[0]) == 4999
        assert len(coh[1]) == 4999

    def test_failure_when_normalization_not_recognized(self):
        with pytest.raises(ValueError):
            cs = AveragedCrossspectrum(self.lc1, self.lc2,
                                       segment_size=1,
                                       norm="wrong", dt=self.dt)

    def test_failure_when_power_type_not_recognized(self):
        with pytest.raises(ValueError):