from __future__ import absolute_import
import shutil
import os
import logging
from packaging import version

import unittest
import numpy as np
import bilby
import scipy
from scipy.stats import ks_2samp, kstest


def ks_2samp_wrapper(data1, data2):
    if version.parse(scipy.__version__) >= version.parse("1.3.0"):
        return ks_2samp(data1, data2, alternative="two-sided", mode="asymp")
    else:
        return ks_2samp(data1, data2)


class Test(unittest.TestCase):
    outdir = 'outdir_for_tests'

    @classmethod
    def setUpClass(self):
        if os.path.isdir(self.outdir):
            try:
                shutil.rmtree(self.outdir)
            except OSError:
                logging.warning(
                    "{} not removed prior to tests".format(self.outdir))

    @classmethod
    def tearDownClass(self):
        if os.path.isdir(self.outdir):
            try:
                shutil.rmtree(self.outdir)
            except OSError:
                logging.warning(
                    "{} not removed prior to tests".format(self.outdir))

    def test_fifteen_dimensional_cbc(self):
        duration = 4.
        sampling_frequency = 2048.
        label = 'full_15_parameters'
        np.random.seed(88170235)

        waveform_arguments = dict(waveform_approximant='IMRPhenomPv2',
                                  reference_frequency=50., minimum_frequency=20.)
        waveform_generator = bilby.gw.WaveformGenerator(
            duration=duration, sampling_frequency=sampling_frequency,
            frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
            parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters,
            waveform_arguments=waveform_arguments)

        ifos = bilby.gw.detector.InterferometerList(['H1', 'L1'])
        ifos.set_strain_data_from_power_spectral_densities(
            sampling_frequency=sampling_frequency, duration=duration,
            start_time=0)

        priors = bilby.gw.prior.BBHPriorDict()
        priors.pop('mass_1')
        priors.pop('mass_2')
        priors['chirp_mass'] = bilby.prior.Uniform(
            name='chirp_mass', latex_label='$M$', minimum=10.0, maximum=100.0,
            unit='$M_{\\odot}$')
        priors['mass_ratio'] = bilby.prior.Uniform(
            name='mass_ratio', latex_label='$q$', minimum=0.5, maximum=1.0)
        priors['geocent_time'] = bilby.core.prior.Uniform(
            minimum=-0.1, maximum=0.1)

        likelihood = bilby.gw.GravitationalWaveTransient(
            interferometers=ifos, waveform_generator=waveform_generator,
            priors=priors, distance_marginalization=False,
            phase_marginalization=False, time_marginalization=False)

        likelihood = bilby.core.likelihood.ZeroLikelihood(likelihood)

        result = bilby.run_sampler(
            likelihood=likelihood, priors=priors, sampler='dynesty',
            npoints=1000, walks=100, outdir=self.outdir, label=label)
        pvalues = [ks_2samp_wrapper(result.priors[key].sample(10000),
                                    result.posterior[key].values).pvalue
                   for key in priors.keys()]
        print("P values per parameter")
        for key, p in zip(priors.keys(), pvalues):
            print(key, p)
        self.assertGreater(kstest(pvalues, "uniform").pvalue, 0.01)


if __name__ == '__main__':
    unittest.main()
