"""
Seismometer logfile reader.
"""

import argparse
import csv
import io
import os
import struct
import traceback
import zlib


class LogReader(object):
    """
    Seismometer logfile reader.
    """

    def __init__(self):
        """
        Constructor.
        """
        self.data = []
        self.num_channels = 0

    def clear(self):
        """
        Initializes the reader.
        """
        del self.data[:]
        self.num_channels = 0

    def unwrap(self):
        """
        Resolves the counter overflows.
        """
        self.sort()
        for c in range(1, self.num_channels+1):
            last = None
            correction = 0
            for i in range(len(self.data)):
                if i > 0:
                    dx = self.data[i][c] - last
                    if dx < -127:
                        correction += 256
                    elif dx > 127:
                        correction -= 256
                last = self.data[i][c]
                self.data[i][c] += correction

    def sort(self):
        """
        Sorts the samples by time.
        """
        self.data.sort(key=lambda item: item[0])

    def load_file(self, file_name):
        """
        Loads the data from a logfile.

        Args:
            file_name (str): path to a logfile
        """
        with open(file_name, 'rb') as handle:
            b = io.BytesIO(zlib.decompress(handle.read()))
            start_time, end_time, self.num_channels = struct.unpack('ddB', b.read(17))
            num_samples, = struct.unpack('i', b.read(4))
            # num_samples /= 2
            sample_interval = (end_time - start_time) / (num_samples - 1)
            t = start_time
            timestamps = []
            channel_data = []
            for c in range(self.num_channels):
                channel_data.append([])
            for sample in range(num_samples):
                try:
                    for c in range(self.num_channels):
                        v, = struct.unpack('B', b.read(1))
                        channel_data[c].append(v)
                    timestamps.append(t)
                    t += sample_interval
                except:
                    pass
        for i in range(len(timestamps)):
            record = [timestamps[i]]
            for c in range(self.num_channels):
                record.append(channel_data[c][i])
            self.data.append(record)

    def load_dir(self, directory):
        """
        Loads data from all files in a directory.

        Args:
            directory (str): directory with logfiles
        """
        for file_name in os.listdir(directory):
            if not file_name.lower().endswith('.log'):
                continue
            try:
                self.load_file(os.path.join(directory, file_name))
            except:
                traceback.print_exc()
        self.sort()

    def save_to_csv(self, file_name):
        """
        Saves data as CSV.

        Args:
            file_name (str): CSV file name (output)
        """
        self.sort()
        if not self.data:
            return
        with open(file_name, 'w') as handle:
            writer = csv.writer(handle)
            if self.data:
                writer.writerow(['TIMESTAMP'] + ['CHANNEL%d' % c for c in range(self.num_channels)])
                writer.writerows(self.data)

    def save_to_wav(self, file_name, channel=0, rate=128000):
        """
        Saves data as a WAV file.

        Args:
            file_name (str): a WAV file name (output)
        """
        import scipy.io.wavfile

        self.sort()
        if not self.data:
            return
        with open(file_name, 'w') as handle:
            scipy.io.wavfile.write(handle, data=self.data[:][channel], rate=rate)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--files', nargs='*', help='load data from those files')
    parser.add_argument('-d', '--directories', nargs='*', help='load data from those directories')
    parser.add_argument('-csv', '--save_to_csv', help='file name to save as CSV')
    args = parser.parse_args()
    logreader = LogReader()
    if args.directories:
        for directory in args.directories:
            logreader.load_dir(directory)
    if args.files:
        for file_name in args.files:
            logreader.load_file(file_name)
    logreader.unwrap()
    if args.save_to_csv:
        logreader.save_to_csv(args.save_to_csv)
