#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
The gse2.core test suite.
"""
import copy
import os
import unittest
import warnings

import numpy as np

from obspy import Stream, Trace, UTCDateTime, read
from obspy.core.util import NamedTemporaryFile
from obspy.io.gse2.libgse2 import ChksumError


class CoreTestCase(unittest.TestCase):
    """
    Test cases for libgse2 core interface
    """
    def setUp(self):
        # directory where the test files are located
        self.path = os.path.dirname(__file__)

    def test_read_via_obspy(self):
        """
        Read files via L{obspy.Trace}
        """
        gse2file = os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')
        testdata = [12, -10, 16, 33, 9, 26, 16, 7, 17, 6, 1, 3, -2]
        # read
        st = read(gse2file, verify_checksum=True)
        st.verify()
        tr = st[0]
        self.assertEqual(tr.stats['station'], 'RJOB')
        self.assertEqual(tr.stats.npts, 12000)
        self.assertEqual(tr.stats['sampling_rate'], 200)
        self.assertEqual(tr.stats.get('channel'), 'Z')
        self.assertAlmostEqual(tr.stats.get('calib'), 9.49e-02)
        self.assertEqual(tr.stats.gse2.get('vang'), -1.0)
        self.assertEqual(tr.stats.gse2.get('hang'), -1.0)
        self.assertEqual(tr.stats.gse2.get('calper'), 1.0)
        self.assertEqual(tr.stats.gse2.get('instype'), '')
        self.assertAlmostEqual(tr.stats.starttime.timestamp,
                               1125455629.850, 6)
        self.assertEqual(tr.data[0:13].tolist(), testdata)

    def test_read_head_via_obspy(self):
        """
        Read header of files via L{obspy.Trace}
        """
        gse2file = os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')
        # read
        st = read(gse2file, headonly=True)
        tr = st[0]
        self.assertEqual(tr.stats['station'], 'RJOB')
        self.assertEqual(tr.stats.npts, 12000)
        self.assertEqual(tr.stats['sampling_rate'], 200)
        self.assertEqual(tr.stats.get('channel'), 'Z')
        self.assertAlmostEqual(tr.stats.get('calib'), 9.49e-02)
        self.assertEqual(tr.stats.gse2.get('vang'), -1.0)
        self.assertEqual(tr.stats.gse2.get('calper'), 1.0)
        self.assertAlmostEqual(tr.stats.starttime.timestamp,
                               1125455629.850, 6)
        self.assertEqual(str(tr.data), '[]')

    def test_read_and_write_via_obspy(self):
        """
        Read and Write files via L{obspy.Trace}
        """
        gse2file = os.path.join(self.path, 'data', 'loc_RNON20040609200559.z')
        # read trace
        st1 = read(gse2file)
        st1.verify()
        tr1 = st1[0]
        # write comparison trace
        with NamedTemporaryFile() as tf:
            tempfile = tf.name
            st2 = Stream()
            st2.traces.append(Trace())
            tr2 = st2[0]
            tr2.data = copy.deepcopy(tr1.data)
            tr2.stats = copy.deepcopy(tr1.stats)
            # raises "UserWarning: Bad value in GSE2 header field"
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', UserWarning)
                st2.write(tempfile, format='GSE2')
            # read comparison trace
            st3 = read(tempfile)
        st3.verify()
        tr3 = st3[0]
        # check if equal
        self.assertEqual(tr3.stats['station'], tr1.stats['station'])
        self.assertEqual(tr3.stats.npts, tr1.stats.npts)
        self.assertEqual(tr3.stats['sampling_rate'],
                         tr1.stats['sampling_rate'])
        self.assertEqual(tr3.stats.get('channel'),
                         tr1.stats.get('channel'))
        self.assertEqual(tr3.stats.get('starttime'),
                         tr1.stats.get('starttime'))
        self.assertEqual(tr3.stats.get('calib'),
                         tr1.stats.get('calib'))
        self.assertEqual(tr3.stats.gse2.get('vang'),
                         tr1.stats.gse2.get('vang'))
        self.assertEqual(tr3.stats.gse2.get('calper'),
                         tr1.stats.gse2.get('calper'))
        np.testing.assert_equal(tr3.data, tr1.data)

    def test_read_and_write_streams_via_obspy(self):
        """
        Read and Write files containing multiple GSE2 parts via L{obspy.Trace}
        """
        files = [os.path.join(self.path, 'data', 'loc_RNON20040609200559.z'),
                 os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')]
        testdata = [12, -10, 16, 33, 9, 26, 16, 7, 17, 6, 1, 3, -2]
        # write test file containing multiple GSE2 parts
        with NamedTemporaryFile() as tf:
            for filename in files:
                with open(filename, 'rb') as f1:
                    # raises "UserWarning: Bad value in GSE2 header field"
                    with warnings.catch_warnings():
                        warnings.simplefilter('ignore', UserWarning)
                        tf.write(f1.read())
            tf.flush()
            st1 = read(tf.name)
        st1.verify()
        self.assertEqual(len(st1), 2)
        tr11 = st1[0]
        tr12 = st1[1]
        self.assertEqual(tr11.stats['station'], 'RNON')
        self.assertEqual(tr12.stats['station'], 'RJOB')
        self.assertEqual(tr12.data[0:13].tolist(), testdata)
        # write and read
        with NamedTemporaryFile() as tf:
            tmpfile = tf.name
            # raises "UserWarning: Bad value in GSE2 header field"
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', UserWarning)
                st1.write(tmpfile, format='GSE2')
            st2 = read(tmpfile)
        st2.verify()
        self.assertEqual(len(st2), 2)
        tr21 = st1[0]
        tr22 = st1[1]
        self.assertEqual(tr21.stats['station'], 'RNON')
        self.assertEqual(tr22.stats['station'], 'RJOB')
        self.assertEqual(tr22.data[0:13].tolist(), testdata)
        np.testing.assert_equal(tr21.data, tr11.data)
        np.testing.assert_equal(tr22.data, tr12.data)

    def test_write_integers_via_obspy(self):
        """
        Write file test via L{obspy.Trace}.
        """
        npts = 1000
        # data cloud of integers - float won't work!
        np.random.seed(815)  # make test reproducible
        data = np.random.randint(-1000, 1000, npts).astype(np.int32)
        stats = {'network': 'BW', 'station': 'TEST', 'location': '',
                 'channel': 'EHE', 'npts': npts, 'sampling_rate': 200.0}
        start = UTCDateTime(2000, 1, 1)
        stats['starttime'] = start
        tr = Trace(data=data, header=stats)
        st = Stream([tr])
        st.verify()
        # write
        with NamedTemporaryFile() as tf:
            tempfile = tf.name
            st.write(tempfile, format="GSE2")
            # read again
            stream = read(tempfile)
        stream.verify()
        np.testing.assert_equal(data, stream[0].data)
        # test default attributes
        self.assertEqual('CM6', stream[0].stats.gse2.datatype)
        self.assertEqual(-1, stream[0].stats.gse2.vang)
        self.assertEqual(1.0, stream[0].stats.gse2.calper)
        self.assertEqual(1.0, stream[0].stats.calib)

    def test_tab_complete_stats(self):
        """
        Read files via L{obspy.Trace}
        """
        gse2file = os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')
        # read
        tr = read(gse2file)[0]
        self.assertIn('station', dir(tr.stats))
        self.assertIn('npts', dir(tr.stats))
        self.assertIn('sampling_rate', dir(tr.stats))
        self.assertEqual(tr.stats['station'], 'RJOB')
        self.assertEqual(tr.stats.npts, 12000)
        self.assertEqual(tr.stats['sampling_rate'], 200)

    def test_write_wrong_format(self):
        """
        Write floating point encoded data
        """
        np.random.seed(815)
        st = Stream([Trace(data=np.random.randn(1000))])
        with NamedTemporaryFile() as tf:
            tmpfile = tf.name
            self.assertRaises(Exception, st.write, tmpfile, format="GSE2")

    def test_read_with_wrong_checksum(self):
        """
        Test if additional kwarg verify_chksum can be given
        """
        # read original file
        gse2file = os.path.join(self.path, 'data',
                                'loc_RJOB20050831023349.z.wrong_chksum')
        # should not fail
        read(gse2file, verify_chksum=False)
        # should fail
        self.assertRaises(ChksumError, read, gse2file, verify_chksum=True)

    def test_read_with_wrong_parameters(self):
        """
        Test if additional kwargs can be given
        """
        # read original file
        gse2file = os.path.join(self.path, 'data',
                                'loc_RJOB20050831023349.z.wrong_chksum')
        # add wrong starttime flag of mseed, should also not fail
        read(gse2file, verify_chksum=False, starttime=None)

    def test_read_gse1_via_obspy(self):
        """
        Read files via L{obspy.Trace}
        """
        gse1file = os.path.join(self.path, 'data', 'loc_STAU20031119011659.z')
        testdata = [-818, -814, -798, -844, -806, -818, -800, -790, -818, -780]
        # read
        st = read(gse1file, verify_checksum=True)
        st.verify()
        tr = st[0]
        self.assertEqual(tr.stats['station'], 'LE0083')
        self.assertEqual(tr.stats.npts, 3000)
        self.assertAlmostEqual(tr.stats['sampling_rate'], 124.9999924)
        self.assertEqual(tr.stats.get('channel'), '  Z')
        self.assertAlmostEqual(tr.stats.get('calib'), 16.0000001)
        self.assertEqual(str(tr.stats.starttime),
                         '2003-11-19T01:16:59.990000Z')
        self.assertEqual(tr.data[0:10].tolist(), testdata)

    def test_read_gse1_head_via_obspy(self):
        """
        Read header via L{obspy.Trace}
        """
        gse1file = os.path.join(self.path, 'data', 'loc_STAU20031119011659.z')
        # read
        st = read(gse1file, headonly=True)
        tr = st[0]
        self.assertEqual(tr.stats['station'], 'LE0083')
        self.assertEqual(tr.stats.npts, 3000)
        self.assertAlmostEqual(tr.stats['sampling_rate'], 124.9999924)
        self.assertEqual(tr.stats.get('channel'), '  Z')
        self.assertAlmostEqual(tr.stats.get('calib'), 16.0000001)
        self.assertEqual(str(tr.stats.starttime),
                         '2003-11-19T01:16:59.990000Z')

    def test_read_intv_gse1_via_obspy(self):
        """
        Read file via L{obspy.Trace}
        """
        gse1file = os.path.join(self.path, 'data',
                                'GRF_031102_0225.GSE.wrong_chksum')
        data1 = [-334, -302, -291, -286, -266, -252, -240, -214]
        data2 = [-468, -480, -458, -481, -481, -435, -432, -389]

        # verify checksum fails
        self.assertRaises(ChksumError, read, gse1file, verify_chksum=True)
        # reading header only
        st = read(gse1file, headonly=True)
        self.assertEqual(len(st), 2)
        # reading without checksum verification
        st = read(gse1file, verify_chksum=False)
        st.verify()
        # first trace
        self.assertEqual(len(st), 2)
        self.assertEqual(st[0].stats['station'], 'GRA1')
        self.assertEqual(st[0].stats.npts, 6000)
        self.assertAlmostEqual(st[0].stats['sampling_rate'], 19.9999997)
        self.assertEqual(st[0].stats.get('channel'), ' BZ')
        self.assertAlmostEqual(st[0].stats.get('calib'), 0.9900001)
        self.assertEqual(st[0].stats.starttime,
                         UTCDateTime('2003-11-02T02:25:00.000000Z'))
        # second trace
        self.assertEqual(len(st), 2)
        self.assertEqual(st[1].stats['station'], 'GRA1')
        self.assertEqual(st[1].stats.npts, 6000)
        self.assertAlmostEqual(st[1].stats['sampling_rate'], 19.9999997)
        self.assertEqual(st[1].stats.get('channel'), ' BN')
        self.assertAlmostEqual(st[1].stats.get('calib'), 0.9200001)
        self.assertEqual(st[1].stats.starttime,
                         UTCDateTime('2003-11-02T02:25:00.000000Z'))
        # check first 8 samples
        self.assertEqual(st[0].data[0:8].tolist(), data1)
        # check last 8 samples
        self.assertEqual(st[1].data[-8:].tolist(), data2)

    def test_read_dos(self):
        """
        Read file with dos newlines / encoding, that is
        Line Feed (LF) and Carriage Return (CR)
        see #355
        """
        filedos = os.path.join(self.path, 'data',
                               'loc_RJOB20050831023349_first100_dos.z')
        fileunix = os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')
        st = read(filedos, verify_chksum=True)
        st2 = read(fileunix, verify_chksum=True)
        np.testing.assert_equal(st[0].data, st2[0].data[:100])
        self.assertEqual(st[0].stats['station'], 'RJOB')

    def test_read_apply_calib(self):
        """
        Tests apply_calib parameter in read method.
        """
        gse2file = os.path.join(self.path, 'data', 'loc_RJOB20050831023349.z')
        testdata = [12, -10, 16, 33, 9, 26, 16, 7, 17, 6, 1, 3, -2]
        # read w/ apply_calib = False
        st = read(gse2file, apply_calib=False)
        tr = st[0]
        self.assertEqual(tr.data[0:13].tolist(), testdata)
        # read w/ apply_calib = True
        st = read(gse2file, apply_calib=True)
        tr = st[0]
        testdata = [n * tr.stats.calib for n in testdata]
        self.assertEqual(tr.data[0:13].tolist(), testdata)

    def test_write_and_read_correct_network(self):
        """
        Tests that writing and reading the STA2 line works (otherwise the
        network code of the data is missing), even if some details like e.g.
        latitude are not present.
        """
        tr = Trace(np.arange(5, dtype=np.int32))
        tr.stats.network = "BW"
        with NamedTemporaryFile() as tf:
            tmpfile = tf.name
            tr.write(tmpfile, format='GSE2')
            tr = read(tmpfile)[0]
        self.assertEqual(tr.stats.network, "BW")

    def test_read_gse2_int_datatype(self):
        """
        Test reading of GSE2 files with data type INT.
        """
        gse2file = os.path.join(self.path, 'data', 'boa___00_07a.gse')
        testdata = [-4, -4, 1, 3, 2, -3, -6, -4, 2, 5]
        # read
        st = read(gse2file, verify_checksum=True)
        st.verify()
        tr = st[0]
        self.assertEqual(tr.stats['station'], 'BBOA')
        self.assertEqual(tr.stats.npts, 6784)
        self.assertAlmostEqual(tr.stats['sampling_rate'], 50.0)
        self.assertEqual(tr.stats.get('channel'), 'CPZ')
        self.assertAlmostEqual(tr.stats.get('calib'), 0.313)
        self.assertEqual(str(tr.stats.starttime),
                         '1990-04-07T00:07:33.000000Z')
        self.assertEqual(tr.data[0:10].tolist(), testdata)


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


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