"""
A series of test functions for the core functions in EQcorrscan.
"""
import unittest
import numpy as np
import logging

from obspy import read, UTCDateTime

from eqcorrscan.utils.synth_seis import generate_synth_data
from eqcorrscan.utils.correlate import get_stream_xcorr
from eqcorrscan.core import lag_calc as lag_calc_module
from eqcorrscan.core.lag_calc import (
    _xcorr_interp, LagCalcError, _prepare_data, lag_calc, xcorr_pick_family,
    _concatenate_and_correlate)
from eqcorrscan.core.match_filter import Detection, Family, Party, Template
from eqcorrscan.helpers.mock_logger import MockLoggingHandler

np.random.seed(999)


class SyntheticTests(unittest.TestCase):
    """ Test lag-calc with synthetic data. """
    @classmethod
    def setUpClass(cls):
        samp_rate = 50
        cls.t_length = .75
        # Make some synthetic templates
        templates, data, seeds = generate_synth_data(
            nsta=5, ntemplates=5, nseeds=10, samp_rate=samp_rate,
            t_length=cls.t_length, max_amp=10, max_lag=15, phaseout="both",
            jitter=0, noise=False, same_phase=True)
        # Rename channels
        channel_mapper = {"SYN_Z": "HHZ", "SYN_H": "HHN"}
        for tr in data:
            tr.stats.channel = channel_mapper[tr.stats.channel]
        for template in templates:
            for tr in template:
                tr.stats.channel = channel_mapper[tr.stats.channel]
        cls.party = Party()
        t = 0
        data_start = data[0].stats.starttime
        for template, template_seeds in zip(templates, seeds):
            template_name = "template_{0}".format(t)
            detections = []
            for i, sample in enumerate(template_seeds["time"]):
                det = Detection(
                    template_name=template_name,
                    detect_time=data_start + (sample / samp_rate),
                    detect_val=template_seeds["SNR"][i] / len(data),
                    no_chans=len(data),
                    chans=[
                        (tr.stats.station, tr.stats.channel) for tr in data],
                    threshold=0.0, threshold_input=0.0, threshold_type="abs",
                    typeofdet="ccc")
                det._calculate_event(
                    template_st=template, estimate_origin=False)
                detections.append(det)
            # Make a fully formed Template
            _template = Template(
                name=template_name, st=template, lowcut=2.0, highcut=15.0,
                samp_rate=samp_rate, filt_order=4, process_length=86400,
                prepick=10. / samp_rate, event=None)
            family = Family(template=_template, detections=detections)
            cls.party += family
            t += 1
        cls.data = data

    def _prepare_data_checks(self, detect_stream_dict, shift_len, family):
        self.assertEqual(
            len(detect_stream_dict), len(family.detections))
        for detection_id, stream in detect_stream_dict.items():
            detection = [d for d in family if d.id == detection_id][0]
            self.assertEqual(len(stream), detection.no_chans)
            for tr in stream:
                self.assertAlmostEqual(
                    tr.stats.endtime - tr.stats.starttime,
                    (2 * shift_len) + self.t_length, 1)

    def test_prepare_data(self):
        shift_len = 0.2
        for family in self.party:
            detect_stream_dict = _prepare_data(
                family=family, detect_data=self.data, shift_len=shift_len)
            self._prepare_data_checks(detect_stream_dict=detect_stream_dict,
                                      family=family, shift_len=shift_len)

    def test_prepare_data_too_short(self):
        data = self.data.copy()
        data.trim(self.party[0][0].detect_time,
                  self.party[0][0].detect_time + 3)
        shift_len = 0.2
        detect_stream_dict = _prepare_data(
            family=self.party[0], detect_data=data, shift_len=shift_len)
        self.assertEqual(len(detect_stream_dict), 1)

    def test_prepare_data_duplicate_channels(self):
        data = self.data.copy()
        data += data[0].copy()
        with self.assertRaises(LagCalcError):
            _prepare_data(
                family=self.party[0], detect_data=data, shift_len=0.2)

    def test_prepare_data_masked(self):
        data = self.data.copy()
        data.cutout(self.party[0][0].detect_time,
                    self.party[0][0].detect_time + 3)
        data.merge()
        shift_len = 0.2
        detect_stream_dict = _prepare_data(
            family=self.party[0], detect_data=data, shift_len=shift_len)
        short_key = self.party[0][0].id
        for key, value in detect_stream_dict.items():
            detection = [d for d in self.party[0] if d.id == key][0]
            if key == short_key:
                self.assertNotEqual(len(value), detection.no_chans)
            else:
                self.assertEqual(len(value), detection.no_chans)

    def test_family_picking(self):
        catalog_dict = xcorr_pick_family(
            family=self.party[0], stream=self.data, shift_len=0.2, plot=False)
        for event in catalog_dict.values():
            self.assertEqual(len(event.picks), len(self.data))
            for pick in event.picks:
                self.assertTrue("cc_max=" in pick.comments[0].text)
                self.assertAlmostEqual(
                    float(pick.comments[0].text.split("=")[-1]), 1.0, 1)

    def test_family_picking_missing_data(self):
        """ Check that this all works okay when one channel has some gaps """
        gappy_data = self.data.copy()
        sr = gappy_data[0].stats.sampling_rate
        mask = np.zeros(len(gappy_data[0].data), dtype=bool)
        gap_start = self.party[0][5].detect_time - 5
        gap_length = 3600
        gap_end = gap_start + gap_length
        gap_start = int(sr * (gap_start - UTCDateTime(0)))
        gap_end = int(sr * (gap_end - UTCDateTime(0)))
        mask[gap_start: gap_end] = np.ones(int(gap_length * sr), dtype=bool)
        gappy_data[0].data = np.ma.masked_array(
            data=gappy_data[0].data, mask=mask)
        catalog_dict = xcorr_pick_family(
            family=self.party[0], stream=gappy_data, shift_len=0.2, plot=False)
        gap = gappy_data.split().get_gaps()
        for event in catalog_dict.values():
            if len(event.picks) != len(self.data):
                self.assertEqual(len(event.picks), len(self.data) - 1)
                # Check that the event happened to be in the gap
                self.assertTrue(gap[0][4] <= event.picks[0].time <= gap[0][5])
                # Check that there isn't a pick on the channel missing data
                self.assertNotIn('.'.join(gap[0][0:4]),
                                 [pick.waveform_id.get_seed_string()
                                  for pick in event.picks])
            for pick in event.picks:
                self.assertTrue("cc_max=" in pick.comments[0].text)
                self.assertAlmostEqual(
                    float(pick.comments[0].text.split("=")[-1]), 1.0, 1)

    def test_family_picking_with_interpolation(self):
        catalog_dict = xcorr_pick_family(
            family=self.party[0], stream=self.data, shift_len=0.2, plot=False,
            interpolate=True)
        for event in catalog_dict.values():
            for pick in event.picks:
                self.assertTrue("cc_max=" in pick.comments[0].text)
                self.assertAlmostEqual(
                    float(pick.comments[0].text.split("=")[-1]), 1.0, 1)

    def test_lag_calc_api(self):
        detections = [d for f in self.party for d in f]
        templates = [f.template.st for f in self.party]
        template_names = [f.template.name for f in self.party]
        output_cat = lag_calc(
            detections, self.data, template_names, templates,
            shift_len=0.2, min_cc=0.4, horizontal_chans=['E', 'N', '1', '2'],
            vertical_chans=['Z'], cores=1, interpolate=False,
            plot=False)
        self.assertEqual(len(output_cat), len(detections))
        for event in output_cat:
            self.assertEqual(len(event.picks), len(self.data))


class SimpleRealDataTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        st = read().detrend().filter("bandpass", freqmin=2, freqmax=20)
        cls.template = st.slice(starttime=st[0].stats.starttime + 4,
                                endtime=st[0].stats.starttime + 10).copy()
        cls.shift_len = 1.0
        cls.detect_streams = []
        for i in range(1, 8):
            cls.detect_streams.append(
                st.slice(starttime=st[0].stats.starttime + i - cls.shift_len,
                         endtime=st[0].stats.starttime + i + 6 + cls.shift_len
                         ).copy())

    def test_correlation_max_and_position(self):
        ccc, chans = _concatenate_and_correlate(
            streams=self.detect_streams, template=self.template, cores=1)
        samp_rate = self.template[0].stats.sampling_rate
        t_start, t_end = (
            self.template[0].stats.starttime, self.template[0].stats.endtime)
        for _ccc, detect_stream in zip(ccc, self.detect_streams):
            d_start, d_end = (detect_stream[0].stats.starttime,
                              detect_stream[0].stats.endtime)
            if d_start <= t_start and d_end >= t_end:
                for ccc_chan in _ccc:
                    self.assertEqual(ccc_chan.max(), 1.0)
                    self.assertEqual(ccc_chan.argmax(),
                                     samp_rate * (t_start - d_start))
            else:
                for ccc_chan in _ccc:
                    self.assertNotEqual(ccc_chan.max(), 1.0)

    def test_correlation_precision(self):
        """Compare to correlation function outputs"""
        ccc, chans = _concatenate_and_correlate(
            streams=self.detect_streams, template=self.template, cores=1)
        fftw_xcorr_func = get_stream_xcorr("fftw")
        for _ccc, detect_stream in zip(ccc, self.detect_streams):
            fftw_ccc, _, _ = fftw_xcorr_func(
                templates=[self.template], stream=detect_stream, stack=False)
            for chan_ccc, fftw_chan_ccc in zip(_ccc, fftw_ccc[0]):
                self.assertTrue(np.allclose(
                    chan_ccc, fftw_chan_ccc, atol=.00001))


class ShortTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        log = logging.getLogger(lag_calc_module.__name__)
        cls._log_handler = MockLoggingHandler(level='DEBUG')
        log.addHandler(cls._log_handler)
        cls.log_messages = cls._log_handler.messages

    def setUp(self):
        self._log_handler.reset()

    def test_error(self):
        with self.assertRaises(LagCalcError):
            raise LagCalcError('Generic error')
        err = LagCalcError('Generic error')
        self.assertEqual('Generic error', err.value)
        self.assertEqual('Generic error', err.__repr__())
        self.assertEqual('LagCalcError: Generic error', err.__str__())

    def test_bad_interp(self):
        ccc = np.array([-0.21483282, -0.59443731, 0.1898917, -0.67516038,
                        0.60129057, -0.71043723,  0.16709118, 0.96839009,
                        1.58283915, -0.3053663])

        _xcorr_interp(ccc, 0.1)
        self.assertEqual(len(self.log_messages['warning']), 1)
        self.assertTrue(
            'not give an accurate result' in self.log_messages['warning'][0])


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