import copy
import unittest
import random
import os
from numpy import complex128, real, sqrt, sin, cos, angle, ceil, log
from numpy import zeros, argmax, array
from pycbc import DYN_RANGE_FAC
from pycbc.waveform import get_td_waveform, get_fd_waveform, td_approximants, fd_approximants
from pycbc.pnutils import nearest_larger_binary_number
from pycbc.types import FrequencySeries, TimeSeries, complex_same_precision_as
from pycbc.types import load_frequencyseries
from pycbc.filter import sigmasq, overlap_cplx, matched_filter_core
from pycbc.filter import compute_max_snr_over_sky_loc_stat
from pycbc.filter import compute_max_snr_over_sky_loc_stat_no_phase
from pycbc.filter import compute_u_val_for_sky_loc_stat_no_phase
from pycbc.filter import compute_u_val_for_sky_loc_stat
from pycbc import psd
from pycbc import vetoes
from utils import parse_args_all_schemes, simple_exit

_scheme, _context = parse_args_all_schemes("correlate")

expected_results = {}
for idx in range(4):
    expected_results[idx] = {}
    for jdx in range(4):
        expected_results[idx][jdx] = {}

expected_results[0][0]['Ip_snr'] = 100.0
expected_results[0][0]['Ip_angle'] = -1.88619488652e-18
expected_results[0][0]['Ip_argmax'] = 0
expected_results[0][0]['Ic_snr'] = 98.7349100759
expected_results[0][0]['Ic_angle'] = 1.60960753393
expected_results[0][0]['Ic_argmax'] = 3

expected_results[0][1]['Ip_snr'] = 96.3390579783
expected_results[0][1]['Ip_angle'] = 0.511744420131
expected_results[0][1]['Ip_argmax'] = 1387
expected_results[0][1]['Ic_snr'] = 96.3390579783
expected_results[0][1]['Ic_angle'] = 2.08254074693
expected_results[0][1]['Ic_argmax'] = 1387

expected_results[0][2]['Ip_snr'] = 98.8434423546
expected_results[0][2]['Ip_angle'] = 0.566451407787
expected_results[0][2]['Ip_argmax'] = 1100
expected_results[0][2]['Ic_snr'] = 98.8485523538
expected_results[0][2]['Ic_angle'] = 2.10418718318
expected_results[0][2]['Ic_argmax'] = 1099

expected_results[0][3]['Ip_snr'] = 96.4239530554
expected_results[0][3]['Ip_angle'] = -0.946889162447
expected_results[0][3]['Ip_argmax'] = 1447
expected_results[0][3]['Ic_snr'] = 95.6528566731
expected_results[0][3]['Ic_angle'] = 1.10466400896
expected_results[0][3]['Ic_argmax'] = 1484

expected_results[1][0]['Ip_snr'] = 96.3390579783
expected_results[1][0]['Ip_angle'] = -0.511744420131
expected_results[1][0]['Ip_argmax'] = 260757
expected_results[1][0]['Ic_snr'] = 94.4282712604
expected_results[1][0]['Ic_angle'] = 0.957548192904
expected_results[1][0]['Ic_argmax'] = 260753

expected_results[1][1]['Ip_snr'] = 100.0
expected_results[1][1]['Ip_angle'] = -1.73241699772e-18
expected_results[1][1]['Ip_argmax'] = 0
expected_results[1][1]['Ic_snr'] = 100.0
expected_results[1][1]['Ic_angle'] = 1.57079632679
expected_results[1][1]['Ic_argmax'] = 0

expected_results[1][2]['Ip_snr'] = 97.6397283701
expected_results[1][2]['Ip_angle'] = 0.156578884423
expected_results[1][2]['Ip_argmax'] = 261862
expected_results[1][2]['Ic_snr'] = 97.7584101045
expected_results[1][2]['Ic_angle'] = 1.69481552225
expected_results[1][2]['Ic_argmax'] = 261861

expected_results[1][3]['Ip_snr'] = 99.4434573331
expected_results[1][3]['Ip_angle'] = -1.50330148916
expected_results[1][3]['Ip_argmax'] = 58
expected_results[1][3]['Ic_snr'] = 98.0771994342
expected_results[1][3]['Ic_angle'] = 0.521399663782
expected_results[1][3]['Ic_argmax'] = 94

expected_results[2][0]['Ip_snr'] = 98.8434423546
expected_results[2][0]['Ip_angle'] = -0.566451407787
expected_results[2][0]['Ip_argmax'] = 261044
expected_results[2][0]['Ic_snr'] = 96.8727372348
expected_results[2][0]['Ic_angle'] = 0.921151545297
expected_results[2][0]['Ic_argmax'] = 261041

expected_results[2][1]['Ip_snr'] = 97.6397283701
expected_results[2][1]['Ip_angle'] = -0.156578884423
expected_results[2][1]['Ip_argmax'] = 282
expected_results[2][1]['Ic_snr'] = 97.6397283701
expected_results[2][1]['Ic_angle'] = 1.41421744237
expected_results[2][1]['Ic_argmax'] = 282

expected_results[2][2]['Ip_snr'] = 100.0
expected_results[2][2]['Ip_angle'] = 2.41820532326e-18
expected_results[2][2]['Ip_argmax'] = 0
expected_results[2][2]['Ic_snr'] = 99.988543227
expected_results[2][2]['Ic_angle'] = 1.5377922043
expected_results[2][2]['Ic_argmax'] = 262143

expected_results[2][3]['Ip_snr'] = 97.3725007917
expected_results[2][3]['Ip_angle'] = -1.61086911817
expected_results[2][3]['Ip_argmax'] = 342
expected_results[2][3]['Ic_snr'] = 96.2035744912
expected_results[2][3]['Ic_angle'] = 0.442670138337
expected_results[2][3]['Ic_argmax'] = 379

expected_results[3][0]['Ip_snr'] = 96.4239530554
expected_results[3][0]['Ip_angle'] = 0.946889162447
expected_results[3][0]['Ip_argmax'] = 260697
expected_results[3][0]['Ic_snr'] = 94.4958639934
expected_results[3][0]['Ic_angle'] = 2.41579357775
expected_results[3][0]['Ic_argmax'] = 260693

expected_results[3][1]['Ip_snr'] = 99.4434573331
expected_results[3][1]['Ip_angle'] = 1.50330148916
expected_results[3][1]['Ip_argmax'] = 262086
expected_results[3][1]['Ic_snr'] = 99.4434573331
expected_results[3][1]['Ic_angle'] = 3.07409781595
expected_results[3][1]['Ic_argmax'] = 262086

expected_results[3][2]['Ip_snr'] = 97.3725007917
expected_results[3][2]['Ip_angle'] = 1.61086911817
expected_results[3][2]['Ip_argmax'] = 261802
expected_results[3][2]['Ic_snr'] = 97.3906866656
expected_results[3][2]['Ic_angle'] = -3.11216854627
expected_results[3][2]['Ic_argmax'] = 261802

expected_results[3][3]['Ip_snr'] = 100.0
expected_results[3][3]['Ip_angle'] = 9.03608368726e-19
expected_results[3][3]['Ip_argmax'] = 0
expected_results[3][3]['Ic_snr'] = 99.4335056063
expected_results[3][3]['Ic_angle'] = 2.02876392072
expected_results[3][3]['Ic_argmax'] = 36

def generate_detector_strain(template_params, h_plus, h_cross):
    polarization = 0

    if hasattr(template_params, 'polarization'):
        polarization = template_params.polarization

    f_plus = cos(polarization)
    f_cross = sin(polarization)

    return h_plus * f_plus + h_cross * f_cross

def make_padded_frequency_series(vec, filter_N=None, delta_f=None):
    """Convert vec (TimeSeries or FrequencySeries) to a FrequencySeries. If
    filter_N and/or delta_f are given, the output will take those values. If
    not told otherwise the code will attempt to pad a timeseries first such that
    the waveform will not wraparound. However, if delta_f is specified to be
    shorter than the waveform length then wraparound *will* be allowed.
    """
    if filter_N is None:
        power = ceil(log(len(vec), 2)) + 1
        N = 2 ** power
    else:
        N = filter_N
    n = N / 2 + 1

    if isinstance(vec, FrequencySeries):
        vectilde = FrequencySeries(zeros(n, dtype=complex_same_precision_as(vec)),
                                   delta_f=1.0, copy=False)
        if len(vectilde) < len(vec):
            cplen = len(vectilde)
        else:
            cplen = len(vec)
        vectilde[0:cplen] = vec[0:cplen]
        delta_f = vec.delta_f

    elif isinstance(vec, TimeSeries):
        # First determine if the timeseries is too short for the specified df
        # and increase if necessary
        curr_length = len(vec)
        new_length = int(nearest_larger_binary_number(curr_length))
        while new_length * vec.delta_t < 1./delta_f:
            new_length = new_length * 2
        vec.resize(new_length)
        # Then convert to frequencyseries
        v_tilde = vec.to_frequencyseries()
        # Then convert frequencyseries to required length and spacing by keeping
        # only every nth sample if delta_f needs increasing, and cutting at
        # Nyquist if the max frequency is too high.
        # NOTE: This assumes that the input and output data is using binary
        #       lengths.
        i_delta_f = v_tilde.get_delta_f()
        v_tilde = v_tilde.numpy()
        df_ratio = int(delta_f / i_delta_f)
        n_freq_len = int((n-1) * df_ratio +1)
        assert(n <= len(v_tilde))
        df_ratio = int(delta_f / i_delta_f)
        v_tilde = v_tilde[:n_freq_len:df_ratio]
        vectilde = FrequencySeries(v_tilde, delta_f=delta_f, dtype=complex128)

    return FrequencySeries(vectilde * DYN_RANGE_FAC, delta_f=delta_f,
                           dtype=complex128)

def get_waveform(wf_params, start_frequency, sample_rate, length,
                 filter_rate, sky_max_template=False):
    delta_f = filter_rate / float(length)
    if wf_params.approximant in fd_approximants():
        hp, hc = get_fd_waveform(wf_params, delta_f=delta_f,
                                 f_lower=start_frequency)

    elif wf_params.approximant in td_approximants():
        hp, hc = get_td_waveform(wf_params, delta_t=1./sample_rate,
                                 f_lower=start_frequency)

    if not sky_max_template:
        hvec = generate_detector_strain(wf_params, hp, hc)
        return make_padded_frequency_series(hvec, length, delta_f=delta_f)
    else:
        return make_padded_frequency_series(hp, length, delta_f=delta_f), \
            make_padded_frequency_series(hc, length, delta_f=delta_f)

class DummyClass(object):
    pass

class TestChisq(unittest.TestCase):
    def setUp(self, *args):
        # Where are my data files?
        if os.path.isfile('test/data/ZERO_DET_high_P.txt'):
            self.dataDir = 'test/data/'
        elif os.path.isfile('data/ZERO_DET_high_P.txt'):
            self.dataDir = 'data/'
        else:
            self.assertTrue(False, msg="Cannot find data files!")
        self.context = _context
        self.scheme = _scheme
        self.tolerance = 1e-6
        self.filter_t_length = 16
        self.low_freq_filter = 30.
        self.sample_rate = 16384
        self.filter_N = int(self.filter_t_length * self.sample_rate)
        self.filter_n = int(self.filter_N / 2 + 1)
        self.filter_delta_f = 1.0 / self.filter_t_length
        self.psd = psd.from_string('aLIGOZeroDetHighPowerGWINC',
                                   self.filter_n, self.filter_delta_f,
                                   self.low_freq_filter)
        self.psd *= DYN_RANGE_FAC*DYN_RANGE_FAC
        wps1 = DummyClass()
        wps1.mass1 = 123.7627
        wps1.mass2 = 72.55471
        wps1.inclination = 1.125029
        wps1.coa_phase = 2.906049
        wps1.approximant='EOBNRv2HM'
        self.wps1 = wps1

        wps2 = DummyClass()
        wps2.mass1 = 131.460647583
        wps2.mass2 = 69.0030059814
        wps2.inclination = 0.8432287
        wps2.coa_phase = 0.2
        wps2.approximant = 'SEOBNRv4_ROM'
        self.wps2 = wps2

        wps3 = copy.deepcopy(wps2)
        wps3.approximant = 'EOBNRv2HM_ROM'
        self.wps3 = wps3

        wps4 = copy.deepcopy(wps2)
        wps4.spin1x = 0.8
        wps4.spin2y = -0.9
        wps4.approximant = 'IMRPhenomPv2'
        self.wps4 = wps4

        self.wps_list = [wps1, wps2, wps3, wps4]

        self.sm_power_chisq = vetoes.SingleDetSkyMaxPowerChisq(num_bins='1')
        self.sm_power_chisq2 = vetoes.SingleDetSkyMaxPowerChisq(num_bins='10')
        self.power_chisq = vetoes.SingleDetPowerChisq(num_bins='1')
        self.power_chisq2 = vetoes.SingleDetPowerChisq(num_bins='10')


    def test_filtering(self):
        idx = self.idx
        jdx = self.jdx
        #w1 = self.wps_list[idx]
        #w2 = self.wps_list[jdx]
        #stilde = get_waveform(w1, self.low_freq_filter-1,
        #                      self.sample_rate, self.filter_N,
        #                      self.sample_rate)
        #try:
        #    stilde.save('data/skymaxtest_stilde_%d.hdf' % idx)
        #except:
        #    pass
        stilde = load_frequencyseries(self.dataDir + \
                                      'skymaxtest_stilde_%d.hdf' % idx)
        s_norm = sigmasq(stilde, psd=self.psd,
                         low_frequency_cutoff=self.low_freq_filter)
        stilde /= sqrt(float(s_norm))
        stilde *= 100
        #hplus, hcross = get_waveform(w2, self.low_freq_filter-1,
        #                             self.sample_rate, self.filter_N,
        #                             self.sample_rate, sky_max_template=True)
        #try:
        #    hplus.save('data/skymaxtest_hplus_%d.hdf' % jdx)
        #    hcross.save('data/skymaxtest_hcross_%d.hdf' % jdx)
        #except:
        #    pass
        hplus = load_frequencyseries(self.dataDir + \
                                     'skymaxtest_hplus_%d.hdf' % jdx)
        hcross = load_frequencyseries(self.dataDir + \
                                      'skymaxtest_hcross_%d.hdf' % jdx)
        hplus.f_lower = self.low_freq_filter
        hplus.params = random.randint(0,100000000000)
        hcross.f_lower = self.low_freq_filter
        hcross.params = random.randint(0,100000000000)
        hp_norm = sigmasq(hplus, psd=self.psd,
                          low_frequency_cutoff=self.low_freq_filter)
        hc_norm = sigmasq(hcross, psd=self.psd,
                          low_frequency_cutoff=self.low_freq_filter)
        hplus /= sqrt(float(hp_norm))
        hcross /= sqrt(float(hc_norm))
        hpc_corr = overlap_cplx(hplus, hcross, psd=self.psd,
                                low_frequency_cutoff=self.low_freq_filter,
                                normalized=False)
        hpc_corr_R = real(hpc_corr)
        I_plus, corr_plus, n_plus = matched_filter_core\
            (hplus, stilde, psd=self.psd,
             low_frequency_cutoff=self.low_freq_filter, h_norm=1.)
        # FIXME: Remove the deepcopies before merging with master
        I_plus = copy.deepcopy(I_plus)
        corr_plus = copy.deepcopy(corr_plus)
        I_cross, corr_cross, n_cross = matched_filter_core\
            (hcross, stilde, psd=self.psd,
             low_frequency_cutoff=self.low_freq_filter, h_norm=1.)
        I_cross = copy.deepcopy(I_cross)
        corr_cross = copy.deepcopy(corr_cross)
        I_plus = I_plus * n_plus
        I_cross = I_cross * n_cross
        IPM = abs(I_plus.data).argmax()
        ICM = abs(I_cross.data).argmax()
        self.assertAlmostEqual(abs(I_plus[IPM]),
                               expected_results[idx][jdx]['Ip_snr'])
        self.assertAlmostEqual(angle(I_plus[IPM]),
                               expected_results[idx][jdx]['Ip_angle'])
        self.assertEqual(IPM, expected_results[idx][jdx]['Ip_argmax'])
        self.assertAlmostEqual(abs(I_cross[ICM]),
                               expected_results[idx][jdx]['Ic_snr'])
        self.assertAlmostEqual(angle(I_cross[ICM]),
                               expected_results[idx][jdx]['Ic_angle'])
        self.assertEqual(ICM, expected_results[idx][jdx]['Ic_argmax'])

        #print "expected_results[{}][{}]['Ip_snr'] = {}" .format(idx,jdx,abs(I_plus[IPM]))
        #print "expected_results[{}][{}]['Ip_angle'] = {}".format(idx,jdx,angle(I_plus[IPM]))
        #print "expected_results[{}][{}]['Ip_argmax'] = {}".format(idx,jdx, IPM)
        #print "expected_results[{}][{}]['Ic_snr'] = {}" .format(idx,jdx,abs(I_cross[ICM]))
        #print "expected_results[{}][{}]['Ic_angle'] = {}".format(idx,jdx,angle(I_cross[ICM]))
        #print "expected_results[{}][{}]['Ic_argmax'] = {}".format(idx,jdx, ICM)

        det_stat_prec = compute_max_snr_over_sky_loc_stat\
            (I_plus, I_cross, hpc_corr_R, hpnorm=1., hcnorm=1.,
             thresh=0.1, analyse_slice=slice(0,len(I_plus.data)))
        det_stat_hom = compute_max_snr_over_sky_loc_stat_no_phase\
            (I_plus, I_cross, hpc_corr_R, hpnorm=1., hcnorm=1.,
             thresh=0.1, analyse_slice=slice(0,len(I_plus.data)))
        idx_max_prec = argmax(det_stat_prec.data)
        idx_max_hom = argmax(det_stat_hom.data)
        max_ds_prec = det_stat_prec[idx_max_prec]
        max_ds_hom = det_stat_hom[idx_max_hom]

        uvals_prec, _ = compute_u_val_for_sky_loc_stat\
            (I_plus.data, I_cross.data, hpc_corr_R, indices=[idx_max_prec],
             hpnorm=1., hcnorm=1.)
        uvals_hom, _ = compute_u_val_for_sky_loc_stat_no_phase\
            (I_plus.data, I_cross.data, hpc_corr_R, indices=[idx_max_hom],
             hpnorm=1., hcnorm=1.)

        ht = hplus * uvals_hom[0] + hcross
        ht_norm = sigmasq(ht, psd=self.psd,
                          low_frequency_cutoff=self.low_freq_filter)
        ht /= sqrt(float(ht_norm))
        ht.f_lower = self.low_freq_filter
        ht.params = random.randint(0,100000000000)
        I_t, corr_t, n_t = matched_filter_core\
            (ht, stilde, psd=self.psd,
             low_frequency_cutoff=self.low_freq_filter, h_norm=1.)
        I_t = I_t * n_t
        self.assertAlmostEqual(abs(real(I_t.data[idx_max_hom])), max_ds_hom)
        self.assertEqual(abs(real(I_t.data[idx_max_hom])),
                         max(abs(real(I_t.data))))
        chisq, _ = self.power_chisq.values\
            (corr_t, array([max_ds_hom]) / n_plus, n_t, self.psd,
             array([idx_max_hom]), ht)

        ht = hplus * uvals_prec[0] + hcross
        ht_norm = sigmasq(ht, psd=self.psd,
                          low_frequency_cutoff=self.low_freq_filter)
        ht /= sqrt(float(ht_norm))
        ht.f_lower = self.low_freq_filter
        ht.params = random.randint(0,100000000000)
        I_t, corr_t, n_t = matched_filter_core\
            (ht, stilde, psd=self.psd,
             low_frequency_cutoff=self.low_freq_filter, h_norm=1.)
        I_t = I_t * n_t

        chisq, _ = self.power_chisq.values\
            (corr_t, array([max_ds_prec]) / n_plus, n_t, self.psd,
             array([idx_max_prec]), ht)

        self.assertAlmostEqual(abs(I_t.data[idx_max_prec]), max_ds_prec)
        self.assertEqual(idx_max_prec, abs(I_t.data).argmax())
        self.assertTrue(chisq < 1E-4)



def test_maker(class_name, idx, jdx):
    class Test(class_name):
        idx = idx
        jdx = jdx

    Test.__name__ = "Test %s" % '_'.join([str(idx),str(jdx)])
    return Test


suite = unittest.TestSuite()
for idx in range(4):
    for jdx in range(4):
        curr_cls = test_maker(TestChisq, idx, jdx)
        suite.addTest(unittest.TestLoader().loadTestsFromTestCase(curr_cls))

if __name__ == '__main__':
    results = unittest.TextTestRunner(verbosity=2).run(suite)
    simple_exit(results)