import numpy as np
import pytest
import stingray.spectroscopy as spec
from stingray.filters import Window1D
from stingray import Lightcurve, AveragedPowerspectrum, AveragedCrossspectrum, Powerspectrum, Crossspectrum
from stingray.modeling import fit_lorentzians, fit_crossspectrum, fit_powerspectrum
from stingray.filters import Optimal1D, Window1D
from scipy.fft import fft
from astropy.modeling import models
from astropy.table import Table, Column

from os import remove


def random_walk(n, step):
    """
    Parameters
    ----------
    n : int
        number of points
    step : float
        maximum step
    """

    return np.cumsum(np.random.uniform(-step, step, n))


def normalize_angle(angle):
    """Refer angle to 0-pi."""
    while angle < 0:
        angle += np.pi
    while angle > np.pi:
        angle -= np.pi
    return angle


def waveform_simple(phases, dph=np.pi / 5, dampl=2.,
                    initial_phase=np.pi / 3):
    """Phases are in radians!"""
    fine_phases = np.arange(0, 2 * np.pi, 0.002)
    fine_shape = np.cos(fine_phases) + 1 / dampl * np.cos(
        2 * (fine_phases + dph))
    mean = np.mean(fine_shape)
    ampl = (np.max(fine_shape) - np.min(fine_shape)) / 2

    shape = np.cos(phases + initial_phase) + 1 / dampl * np.cos(
        2 * (initial_phase + phases + dph))
    shape -= mean
    shape /= ampl

    return shape


def fake_qpo(t, f0=1., waveform=None, timescale=10, astep=0, phstep=0,
             waveform_opts=None, rms=0.1):
    """
    Parameters
    ----------
    t : array
        Times at which the lightcurve is calculated
    f0: float
        mean frequency of the QPO
    waveform : function
        Function that accepts phases from 0 to pi as input and returns an
        waveform with mean 0 and total amplitude 2 (typically from -1 to 1).
        Default is ``np.sin``
    timescale : float
        In this timescale, expressed in seconds, the waveform changes its phase
        or amplitude or both by an entire step
    astep : float
        Steps of variation of amplitude
    phstep : float
        Steps of variation of phase
    """

    if waveform is None:
        waveform = np.sin

    n = len(t)
    phase = 2 * np.pi * f0 * t + random_walk(n, phstep / timescale)
    amp = 1 + random_walk(n, astep / timescale)
    qpo_lc = amp * (1 + rms * waveform(phase, **waveform_opts))

    return phase, qpo_lc


class TestCCF(object):

    @classmethod
    def setup_class(cls):
        total_length = 10000
        f_qpo = 1.5
        cls.dt = 1 / f_qpo / 40
        approx_Q = 10
        q_len = approx_Q / f_qpo
        sigma = 0.1
        astep = 0.01
        phstep = 1
        real_dphi = 0.4 * np.pi
        cls.n_seconds = 500
        cls.n_seg = int(total_length / cls.n_seconds)
        cls.n_bins = cls.n_seconds/cls.dt

        times = np.arange(0, total_length, cls.dt)
        _, cls.ref_counts = fake_qpo(times, f0=f_qpo, astep=astep, rms=sigma,
                                     waveform=waveform_simple, phstep=phstep,
                                     timescale=q_len,
                                     waveform_opts={'dph': real_dphi})
        _, ci_counts = fake_qpo(times, f0=f_qpo, astep=astep, rms=sigma,
                                waveform=waveform_simple, phstep=phstep,
                                timescale=q_len,
                                waveform_opts={'dph': real_dphi})
        cls.ci_counts = np.array([ci_counts])

        cls.ref_times = np.arange(0, cls.n_seconds * cls.n_seg, cls.dt)
        cls.ref_lc = Lightcurve(cls.ref_times, cls.ref_counts, dt=cls.dt)
        ref_aps = AveragedPowerspectrum(cls.ref_lc, segment_size=cls.n_seconds,
                                        norm='abs')
        df = ref_aps.freq[1] - ref_aps.freq[0]
        amplitude_0 = np.max(ref_aps.power)
        x_0_0 = ref_aps.freq[np.argmax(ref_aps.power)]
        amplitude_1 = amplitude_0 / 2
        x_0_1 = x_0_0 * 2
        fwhm = df

        cls.model = models.Lorentz1D(amplitude=amplitude_0, x_0=x_0_0,
                                     fwhm=fwhm) + \
            models.Lorentz1D(amplitude=amplitude_1, x_0=x_0_1,
                             fwhm=fwhm)
        cls.ref_aps = ref_aps

    def test_ccf(self):
        # to make testing faster, fitting is not done.
        ref_ps = Powerspectrum(self.ref_lc, norm='abs')

        ci_counts_0 = self.ci_counts[0]
        ci_times = np.arange(0, self.n_seconds * self.n_seg, self.dt)
        ci_lc = Lightcurve(ci_times, ci_counts_0, dt=self.dt)

        # rebinning factor used in `rebin_log`
        rebin_log_factor = 0.4

        acs = AveragedCrossspectrum(lc1=ci_lc, lc2=self.ref_lc,
                                    segment_size=self.n_seconds, norm='leahy',
                                    power_type="absolute")
        acs = acs.rebin_log(rebin_log_factor)

        # parest, res = fit_crossspectrum(acs, self.model, fitmethod="CG")
        acs_result_model = self.model

        # using optimal filter
        optimal_filter = Optimal1D(acs_result_model)
        optimal_filter_freq = optimal_filter(acs.freq)
        filtered_acs_power = optimal_filter_freq * np.abs(acs.power)

        # rebinning power spectrum
        new_df = spec.get_new_df(ref_ps, self.n_bins)
        ref_ps_rebinned = ref_ps.rebin(df=new_df)

        # parest, res = fit_powerspectrum(ref_ps_rebinned, self.model)
        ref_ps_rebinned_result_model = self.model

        # calculating rms from power spectrum
        ref_ps_rebinned_rms = spec.compute_rms(ref_ps_rebinned,
                                               ref_ps_rebinned_result_model,
                                               criteria="optimal")

        # calculating normalized ccf
        ccf_norm = spec.ccf(filtered_acs_power, ref_ps_rebinned_rms,
                            self.n_bins)

        # calculating ccf error
        meta = {'N_SEG': self.n_seg, 'NSECONDS': self.n_seconds, 'DT': self.dt,
                'N_BINS': self.n_bins}
        error_ccf, avg_seg_ccf = spec.ccf_error(self.ref_counts, ci_counts_0,
                                                acs_result_model,
                                                rebin_log_factor,
                                                meta, ref_ps_rebinned_rms,
                                                filter_type="optimal")

        assert np.all(np.isclose(ccf_norm, avg_seg_ccf, atol=0.01))
        assert np.all(np.isclose(error_ccf, np.zeros(shape=error_ccf.shape),
                                 atol=0.01))

        # using window function
        tophat_filter = Window1D(acs_result_model)
        tophat_filter_freq = tophat_filter(acs.freq)
        filtered_acs_power = tophat_filter_freq * np.abs(acs.power)

        ref_ps_rebinned_rms = spec.compute_rms(ref_ps_rebinned,
                                               ref_ps_rebinned_result_model,
                                               criteria="window")

        ccf_norm = spec.ccf(filtered_acs_power, ref_ps_rebinned_rms,
                            self.n_bins)

        error_ccf, avg_seg_ccf = spec.ccf_error(self.ref_counts, ci_counts_0,
                                                acs_result_model,
                                                rebin_log_factor,
                                                meta, ref_ps_rebinned_rms,
                                                filter_type="window")

        assert np.all(np.isclose(ccf_norm, avg_seg_ccf, atol=0.01))
        assert np.all(np.isclose(error_ccf, np.zeros(shape=error_ccf.shape),
                                 atol=0.01))

    def test_get_mean_phase_difference(self):
        _, a, b, _ = spec.get_parameters(self.ref_aps.lc1.counts,
                                         self.ref_aps.lc1.dt, self.model)
        assert a == b


def test_load_lc_fits():
    dt = 0.1
    n_seg = 5
    n_seconds = 100
    output_file = 'out.fits'

    ref_lc = np.arange(0, n_seconds, dt)
    ci_lc = np.array(np.array_split(np.arange(0, n_seconds*n_seg, dt), n_seg))

    n_bins = len(ref_lc)

    lightcurves = Table()
    lightcurves.add_column(Column(name='REF', data=ref_lc.T/dt))
    lightcurves.add_column(Column(name='CI', data=ci_lc.T/dt))
    lightcurves.meta['N_BINS'] = n_bins
    lightcurves.meta['DT'] = dt
    lightcurves.meta['N_SEG'] = n_seg
    lightcurves.meta['NSECONDS'] = n_seconds
    lightcurves.write(output_file, format='fits', overwrite=True)

    ref, ci, meta = spec.load_lc_fits(output_file, counts_type=False)
    remove(output_file)

    assert np.allclose(ref_lc, ref)
    assert np.allclose(ci_lc, ci)
    assert meta['N_BINS'] == n_bins
    assert meta['DT'] == dt
    assert meta['N_SEG'] == n_seg
    assert meta['NSECONDS'] == n_seconds


def test_psi_distance():
    a = np.arange(-10, 11)
    dist = spec.psi_distance(0, a)
    assert np.allclose(np.abs(a), dist)


def test_get_new_df():
    np.random.seed(150)

    amplitude_0 = 200.0
    amplitude_1 = 100.0
    amplitude_2 = 50.0

    x_0_0 = 0.5
    x_0_1 = 2.0
    x_0_2 = 7.5

    fwhm_0 = 0.1
    fwhm_1 = 1.0
    fwhm_2 = 0.5

    whitenoise = 100.0

    model = models.Lorentz1D(amplitude_0, x_0_0, fwhm_0) + \
        models.Lorentz1D(amplitude_1, x_0_1, fwhm_1) + \
        models.Lorentz1D(amplitude_2, x_0_2, fwhm_2) + \
        models.Const1D(whitenoise)

    freq = np.linspace(0.01, 10.0, 1000)
    p = model(freq)
    noise = np.random.exponential(size=len(freq))

    power = p * noise
    cs = Crossspectrum()
    cs.freq = freq
    cs.power = power
    cs.df = cs.freq[1] - cs.freq[0]
    cs.n = len(freq)
    cs.m = 1

    assert np.isclose(cs.df, spec.get_new_df(cs, cs.n), rtol=0.001)


def test_waveform():
    mu = 1
    avg_sigma_1 = 1
    avg_sigma_2 = 2
    cap_phi_1 = np.pi/2
    cap_phi_2 = np.pi/2
    x = np.linspace(0, 2 * 2 * np.pi, 100)

    y1 = spec.waveform(x, mu, avg_sigma_1, avg_sigma_2, cap_phi_1, cap_phi_2)
    y2 = mu * (1 + np.sqrt(2) * (avg_sigma_1 * np.cos(x - cap_phi_1) +
                                 avg_sigma_2 * np.cos(2 * x - cap_phi_2)))

    assert np.allclose(y1, y2)


def test_x_2_function():
    a = [-2, -1, 0, 1, 2]
    y1 = spec.x_2_function(0, a)
    y2 = np.sum(np.array(a) ** 2)
    assert np.allclose(y1, y2)


def test_compute_rms():
    np.random.seed(150)

    amplitude_0 = 200.0
    amplitude_1 = 100.0
    amplitude_2 = 50.0

    x_0_0 = 0.5
    x_0_1 = 2.0
    x_0_2 = 7.5

    fwhm_0 = 0.1
    fwhm_1 = 1.0
    fwhm_2 = 0.5

    whitenoise = 100.0

    model = models.Lorentz1D(amplitude_0, x_0_0, fwhm_0) + \
        models.Lorentz1D(amplitude_1, x_0_1, fwhm_1) + \
        models.Lorentz1D(amplitude_2, x_0_2, fwhm_2) + \
        models.Const1D(whitenoise)

    freq = np.linspace(-10.0, 10.0, 1000)
    p = model(freq)
    noise = np.random.exponential(size=len(freq))

    power = p * noise
    cs = Crossspectrum()
    cs.freq = freq
    cs.power = power
    cs.df = cs.freq[1] - cs.freq[0]
    cs.n = len(freq)
    cs.m = 1

    rms = np.sqrt(np.sum(model(cs.freq) * cs.df)).mean()

    assert rms == spec.compute_rms(cs, model, criteria="all")

    rms_pos = np.sqrt(np.sum(model(cs.freq[cs.freq > 0]) * cs.df)).mean()

    assert rms_pos == spec.compute_rms(cs, model, criteria="posfreq")

    optimal_filter = Window1D(model)
    optimal_filter_freq = optimal_filter(cs.freq)
    filtered_cs_power = optimal_filter_freq * np.abs(model(cs.freq))

    rms = np.sqrt(np.sum(filtered_cs_power * cs.df)).mean()
    assert rms == spec.compute_rms(cs, model, criteria="window")

    with pytest.raises(ValueError):
        spec.compute_rms(cs, model, criteria="filter")
