# -*- coding: utf-8 -*-
"""
Tests to ensure that the new GW proposal is equivalent to the old method
"""
import copy
import logging

import numpy as np

import pytest
from conftest import requires_dependency

from nessai.model import Model as BaseModel
from nessai.gw.proposal import GWFlowProposal, LegacyGWFlowProposal
from nessai.flowsampler import FlowSampler
from nessai.livepoint import dict_to_live_points

try:
    import bilby
except ImportError:
    print('Could not import bilby, some tests will not work')


logger = logging.getLogger(__name__)


@pytest.fixture(scope='module')
def injection_parameters():
    injection_parameters = dict(
        mass_1=36., mass_2=29., a_1=0.4, a_2=0.3, tilt_1=0.5, tilt_2=1.0,
        phi_12=1.7, phi_jl=0.3, luminosity_distance=2000., theta_jn=0.4,
        psi=2.659, phase=1.3, geocent_time=1126259642.413, ra=1.375,
        dec=-1.2108)
    return injection_parameters


@pytest.fixture(scope='function')
def legacy_kwargs():
    kwargs = {
        "boundary_inversion": True,
        "update_bounds": True,
        "rescale_parameters": True,
        "analytic_priors": True,
        "reparameterisations": {
            "rescale_angles": ["psi"],
            "inversion": True,
            "convert_to_angle": {"a_1": "split", "a_2": "duplicate"},
            "use_x_prime_prior": True,
            "uniform_distance_parameter": False,
            "spin_conversion": False
        },
        "detect_edges": True,
        "inversion_type": {"mass_ratio": "duplicate", "dc3": "duplicate"},
    }
    return kwargs


@pytest.fixture(scope='function')
def kwargs(legacy_kwargs):
    kwargs = copy.deepcopy(legacy_kwargs)
    kwargs['reparameterisations'] = {
        "mass_ratio": {"reparameterisation": "mass_ratio", "prior": "uniform"},
        "chirp_mass": {"reparameterisation": "mass", "prior": "uniform"},
        "luminosity_distance": {"reparameterisation": "distance",
                                "prior": "uniform-comoving-volume"},
        "geocent_time": {"reparameterisation": "time", "prior": "uniform",
                         "update_bounds": False},
        "sky-ra-dec": {"parameters": ["ra", "dec"], "prior": "isotropic"},
        "psi": {"reparameterisation": "angle-pi"},
        "theta_jn": {"reparameterisation": "angle-sine"},
        "tilt_1": {"reparameterisation": "angle-sine"},
        "tilt_2": {"reparameterisation": "angle-sine"},
        "phi_12": {"reparameterisation": "angle-2pi"},
        "phi_jl": {"reparameterisation": "angle-2pi"},
        "a_1": {"reparameterisation": "to-cartesian", "prior": "uniform"},
        "a_2": {"reparameterisation": "to-cartesian", "prior": "uniform",
                "mode": "duplicate"}
    }
    kwargs.pop('inversion_type')
    return kwargs


@requires_dependency('bilby')
@requires_dependency('lal')
@requires_dependency('astropy')
@pytest.mark.parametrize('parameters',
                         ['psi',
                          'theta_jn',
                          'tilt_1',
                          'tilt_2',
                          'phi_12',
                          'phi_jl',
                          'a_1',
                          'a_2',
                          ['ra', 'dec'],
                          ['chirp_mass', 'geocent_time'],
                          ['chirp_mass', 'mass_ratio'],
                          ['chirp_mass', 'luminosity_distance'],
                          ])
def test_parameter(parameters, injection_parameters, kwargs, legacy_kwargs,
                   tmpdir):
    """
    Test that the two proposal methods are equivalent for a set of parameters.

    Checks if the resulting x prime parameters are the same and their
    log-Jacobian determinants are also the same.
    """
    outdir = str(tmpdir.mkdir('test'))
    priors = bilby.gw.prior.BBHPriorDict()
    fixed_params = ["chirp_mass", "mass_ratio", "phi_12", "phi_jl", "a_1",
                    "a_2", "tilt_1", "tilt_2", "ra", "dec",
                    "luminosity_distance", "geocent_time", "theta_jn", "psi"]
    try:
        fixed_params.remove(parameters)
    except ValueError:
        for p in parameters:
            fixed_params.remove(p)
    priors['geocent_time'] = bilby.core.prior.Uniform(
            minimum=injection_parameters['geocent_time'] - 0.1,
            maximum=injection_parameters['geocent_time'] + 0.1,
            name='geocent_time', latex_label='$t_c$', unit='$s$')
    for key in fixed_params:
        if key in injection_parameters:
            priors[key] = injection_parameters[key]
        else:
            priors.pop(key)

    priors['mass_1'] = injection_parameters['mass_1']
    priors['mass_2'] = injection_parameters['mass_2']

    waveform_generator = bilby.gw.WaveformGenerator(
        duration=1, sampling_frequency=256,
        frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole)

    likelihood = bilby.gw.GravitationalWaveTransient(
        interferometers=['H1'],
        waveform_generator=waveform_generator,
        priors=priors,
        phase_marginalization=True,
        distance_marginalization=False,
        time_marginalization=False, reference_frame='sky'
        )

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

    def log_prior(theta):
        return priors.ln_prob(theta, axis=0)

    search_parameter_keys = []
    for key in priors:
        if isinstance(priors[key], bilby.core.prior.Prior) \
                and priors[key].is_fixed is False:
            search_parameter_keys.append(key)

    def log_likelihood(theta):
        params = {
            key: t for key, t in zip(search_parameter_keys, theta)}

        likelihood.parameters.update(params)
        return likelihood.log_likelihood_ratio()

    class Model(BaseModel):
        def __init__(self, names, priors):
            self.names = names
            self.priors = priors
            self._update_bounds()

        @staticmethod
        def log_likelihood(x, **kwargs):
            theta = [x[n].item() for n in search_parameter_keys]
            return log_likelihood(theta)

        @staticmethod
        def log_prior(x, names=None, **kwargs):
            if names is None:
                names = search_parameter_keys
            theta = {n: x[n] for n in names}
            return log_prior(theta)

        def _update_bounds(self):
            self.bounds = \
                {key: [self.priors[key].minimum, self.priors[key].maximum]
                 for key in self.names}

        def new_point(self, N=1):
            prior_samples = self.priors.sample(size=N)
            samples = {n: prior_samples[n] for n in self.names}
            self._update_bounds()
            return dict_to_live_points(samples)

        def new_point_log_prob(self, x):
            return self.log_prior(x)

    reparameterisations = kwargs.pop('reparameterisations')
    for k in list(reparameterisations.keys()):
        if k == 'sky-ra-dec':
            if any(p in fixed_params for p in ['ra', 'dec']):
                del reparameterisations['sky-ra-dec']
        elif k not in search_parameter_keys:
            del reparameterisations[k]

    if 'luminosity_distance' in parameters:
        legacy_kwargs['reparameterisations']['uniform_distance_parameter'] = \
                True

    model = Model(search_parameter_keys, priors)
    sampler = FlowSampler(model, output=outdir,
                          flow_class=LegacyGWFlowProposal, **legacy_kwargs)

    new_sampler = FlowSampler(model, output=outdir,
                              reparameterisations=reparameterisations,
                              flow_class=GWFlowProposal, **kwargs)

    sampler.ns._flow_proposal.initialise()
    new_sampler.ns._flow_proposal.initialise()

    sampler.ns._uninformed_proposal.populate(N=100)
    x = sampler.ns._uninformed_proposal.samples

    orig_proposal = sampler.ns._flow_proposal
    new_proposal = new_sampler.ns._flow_proposal

    tests = ['lower', 'upper', False]
    test_results = {t: False for t in tests}
    for test in tests:
        orig_proposal.check_state(x)
        np.random.seed(1234)
        x_prime, lj = orig_proposal.rescale(x, test=test)
        log_p = orig_proposal.x_prime_log_prior(x_prime)

        new_proposal.check_state(x)
        np.random.seed(1234)
        x_prime_new, lj_new = new_proposal.rescale(x, test=test)
        log_p_new = new_proposal.x_prime_log_prior(x_prime_new)

        np.testing.assert_array_equal(lj, lj_new)
        np.testing.assert_array_equal(log_p, log_p_new)
        try:
            np.testing.assert_array_equal(x_prime, x_prime_new)
        except AssertionError:
            flag = {n: False for n in x_prime.dtype.names}
            for n in x_prime.dtype.names:
                for nn in x_prime_new.dtype.names:
                    if np.allclose(x_prime[n], x_prime_new[nn]):
                        logger.critical(f'{n} is equivalent to {nn}')
                        flag[n] = True

            assert all(v for v in flag.values()), print(f'Flags: {flag}')
        test_results[test] = True

    assert all(t for t in test_results.values())

    if getattr(orig_proposal, '_rescaled_min'):
        logger.info('Checking prime prior bounds')
        test_results = {t: False for t in tests}
        for test in tests:
            logger.info(f'Testing with inversion and test={test}')
            orig_proposal.check_state(x)
            np.random.seed(1234)
            t, lj = orig_proposal.rescale(x, test=test)
            new_proposal.check_state(x)
            np.random.seed(1234)
            t_new, lj_new = new_proposal.rescale(x, test=test)
            min_vals = \
                {n: False for n in new_proposal._reparameterisation.keys()}
            max_vals = \
                {n: False for n in new_proposal._reparameterisation.keys()}
            for k, r in new_proposal._reparameterisation.items():
                for pb in r.prime_prior_bounds.values():
                    for vmin in orig_proposal._rescaled_min.values():
                        logger.warning([k, pb, vmin])
                        if np.isclose(pb[0], vmin):
                            min_vals[k] = True
                    for vmax in orig_proposal._rescaled_max.values():
                        logger.warning([k, pb, vmax])
                        if np.isclose(pb[1], vmax):
                            max_vals[k] = True

            assert all(t for t in min_vals.values()), print(min_vals)
            assert all(t for t in max_vals.values()), print(max_vals)
            test_results[test] = True

        assert all(t for t in test_results.values())
