#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
The psd test suite.
"""

from obspy import Trace, Stream, UTCDateTime
from obspy.core.util.base import NamedTemporaryFile
from obspy.signal.spectral_estimation import PPSD, psd, welch_window, \
    welch_taper
import numpy as np
import os
import unittest
import warnings


class PsdTestCase(unittest.TestCase):
    """
    Test cases for psd.
    """
    def setUp(self):
        # directory where the test files are located
        self.path = os.path.join(os.path.dirname(__file__), 'data')

    def test_obspy_psd_vs_pitsa(self):
        """
        Test to compare results of PITSA's psd routine to the
        :func:`matplotlib.mlab.psd` routine wrapped in
        :func:`obspy.signal.spectral_estimation.psd`.
        The test works on 8192 samples long gaussian noise with a standard
        deviation of 0.1 generated with PITSA, sampling rate for processing in
        PITSA was 100.0 Hz, length of nfft 512 samples. The overlap in PITSA
        cannot be controlled directly, instead only the number of overlapping
        segments can be specified.  Therefore the test works with zero overlap
        to have full control over the data segments used in the psd.
        It seems that PITSA has one frequency entry more, i.e. the psd is one
        point longer. I dont know were this can come from, for now this last
        sample in the psd is ignored.
        """
        SAMPLING_RATE = 100.0
        NFFT = 512
        NOVERLAP = 0
        file_noise = os.path.join(self.path, "pitsa_noise.npy")
        fn_psd_pitsa = "pitsa_noise_psd_samprate_100_nfft_512_noverlap_0.npy"
        file_psd_pitsa = os.path.join(self.path, fn_psd_pitsa)

        noise = np.load(file_noise)
        # in principle to mimic PITSA's results detrend should be specified as
        # some linear detrending (e.g. from matplotlib.mlab.detrend_linear)
        psd_obspy, _ = psd(noise, NFFT=NFFT, Fs=SAMPLING_RATE,
                           window=welch_taper, noverlap=NOVERLAP)
        psd_pitsa = np.load(file_psd_pitsa)

        # mlab's psd routine returns Nyquist frequency as last entry, PITSA
        # seems to omit it and returns a psd one frequency sample shorter.
        psd_obspy = psd_obspy[:-1]

        # test results. first couple of frequencies match not as exactly as all
        # the rest, test them separately with a little more allowance..
        np.testing.assert_array_almost_equal(psd_obspy[:3], psd_pitsa[:3],
                                             decimal=4)
        np.testing.assert_array_almost_equal(psd_obspy[1:5], psd_pitsa[1:5],
                                             decimal=5)
        np.testing.assert_array_almost_equal(psd_obspy[5:], psd_pitsa[5:],
                                             decimal=6)

    def test_welch_window_vs_pitsa(self):
        """
        Test that the helper function to generate the welch window delivers the
        same results as PITSA's routine.
        Testing both even and odd values for length of window.
        Not testing strange cases like length <5, though.
        """
        file_welch_even = os.path.join(self.path, "pitsa_welch_window_512.npy")
        file_welch_odd = os.path.join(self.path, "pitsa_welch_window_513.npy")

        for file, N in zip((file_welch_even, file_welch_odd), (512, 513)):
            window_pitsa = np.load(file)
            window_obspy = welch_window(N)
            np.testing.assert_array_almost_equal(window_pitsa, window_obspy)

    def test_PPSD(self):
        """
        Test PPSD routine with some real data. Data was downsampled to 100Hz
        so the ppsd is a bit distorted which does not matter for the purpose
        of testing.
        """
        # load test file
        file_data = os.path.join(
            self.path, 'BW.KW1._.EHZ.D.2011.090_downsampled.asc.gz')
        file_histogram = os.path.join(
            self.path,
            'BW.KW1._.EHZ.D.2011.090_downsampled__ppsd_hist_stack.npy')
        file_binning = os.path.join(
            self.path, 'BW.KW1._.EHZ.D.2011.090_downsampled__ppsd_mixed.npz')
        # parameters for the test
        data = np.loadtxt(file_data)
        stats = {'_format': 'MSEED',
                 'calib': 1.0,
                 'channel': 'EHZ',
                 'delta': 0.01,
                 'endtime': UTCDateTime(2011, 3, 31, 2, 36, 0, 180000),
                 'location': '',
                 'mseed': {'dataquality': 'D', 'record_length': 512,
                           'encoding': 'STEIM2', 'byteorder': '>'},
                 'network': 'BW',
                 'npts': 936001,
                 'sampling_rate': 100.0,
                 'starttime': UTCDateTime(2011, 3, 31, 0, 0, 0, 180000),
                 'station': 'KW1'}
        tr = Trace(data, stats)
        st = Stream([tr])
        paz = {'gain': 60077000.0,
               'poles': [(-0.037004 + 0.037016j), (-0.037004 - 0.037016j),
                         (-251.33 + 0j), (-131.04 - 467.29j),
                         (-131.04 + 467.29j)],
               'sensitivity': 2516778400.0,
               'zeros': [0j, 0j]}
        ppsd = PPSD(tr.stats, paz, db_bins=(-200, -50, 0.5))
        ppsd.add(st)
        # read results and compare
        result_hist = np.load(file_histogram)
        self.assertEqual(len(ppsd.times), 4)
        self.assertEqual(ppsd.nfft, 65536)
        self.assertEqual(ppsd.nlap, 49152)
        np.testing.assert_array_equal(ppsd.hist_stack, result_hist)
        # add the same data a second time (which should do nothing at all) and
        # test again - but it will raise UserWarnings, which we omit for now
        with warnings.catch_warnings(record=True):
            warnings.simplefilter('ignore', UserWarning)
            ppsd.add(st)
            np.testing.assert_array_equal(ppsd.hist_stack, result_hist)
        # test the binning arrays
        binning = np.load(file_binning)
        np.testing.assert_array_equal(ppsd.spec_bins, binning['spec_bins'])
        np.testing.assert_array_equal(ppsd.period_bins, binning['period_bins'])

        # test saving and loading of the PPSD (using a temporary file)
        with NamedTemporaryFile() as tf:
            filename = tf.name
            # test saving and loading an uncompressed file
            ppsd.save(filename, compress=False)
            ppsd_loaded = PPSD.load(filename)
            self.assertEqual(len(ppsd_loaded.times), 4)
            self.assertEqual(ppsd_loaded.nfft, 65536)
            self.assertEqual(ppsd_loaded.nlap, 49152)
            np.testing.assert_array_equal(ppsd_loaded.hist_stack, result_hist)
            np.testing.assert_array_equal(ppsd_loaded.spec_bins,
                                          binning['spec_bins'])
            np.testing.assert_array_equal(ppsd_loaded.period_bins,
                                          binning['period_bins'])
            # test saving and loading a compressed file
            ppsd.save(filename, compress=True)
            ppsd_loaded = PPSD.load(filename)
            self.assertEqual(len(ppsd_loaded.times), 4)
            self.assertEqual(ppsd_loaded.nfft, 65536)
            self.assertEqual(ppsd_loaded.nlap, 49152)
            np.testing.assert_array_equal(ppsd_loaded.hist_stack, result_hist)
            np.testing.assert_array_equal(ppsd_loaded.spec_bins,
                                          binning['spec_bins'])
            np.testing.assert_array_equal(ppsd_loaded.period_bins,
                                          binning['period_bins'])


def suite():
    return unittest.makeSuite(PsdTestCase, 'test')


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