"""Provides functions to parse datagrams in files generated by Simard SN90 and CS90 sonars.

This file is a modified version of a file of the same name in the pyEcholab repository, 
avaiable at https://github.com/noaa-afsc-mace/pyEcholab (see the echolab2/instruments/util 
directory). See that file for the licence details and author aknowledgements.

The classes defined here are used by code in the simrad_raw_file.py file.
"""

import struct
import numpy as np
from construct import Struct, Array, this, Container, PaddedString, GreedyString
from construct import If, IfThenElse, Timestamp, GreedyBytes
from construct import Int32ul, Int32sl, Int16ul, Float32l, Float64l, Int64ul, Bytes

__all__ = ['SimradSINParser', 'SimradVERParser', 'SimradPHYParser',
            'SimradPCOParser', 'SimradPINParser', 'SimradEOPParser',
            'SimradSENParser', 'SimradSECParser', 'SimradRAWParser']

def construct_to_dict(obj):
    """Recursively convert a construct data stucture into a dict."""
    if isinstance(obj, Container):
        result_dict = {}
        for key, value in obj.items():
            result_dict[key] = construct_to_dict(value)
        return result_dict
    elif isinstance(obj, list):
        return [construct_to_dict(item) for item in obj]
    else:
        return obj


class _SimradDatagramParser():
    """Base class for a Simrad datagram parser class."""

    def __init__(self, header_type, versions = []):
        self._id = header_type
        self._versions = versions

    def validate_data_header(self, data):

        if isinstance(data, dict):
            type_ = data['type'][:3]
            version   = int(data['type'][3])
        elif isinstance(data, str):
            type_ = data[:3]
            version   = int(data[3])
        else:
            raise TypeError('Expected a dict or str')

        if type_ != self._id:
            raise ValueError('Expected data of type %s, not %s' %(self._id, type_))

        if version not in self._versions:
            raise ValueError('No parser available for type %s version %d' %(self._id, version))

        return type_, version

    def from_string(self, raw_string, bytes_read):
        """Parse the dataframe from the raw_string."""

        header = raw_string[:4]
        header = header.decode()
        id_, version = self.validate_data_header(header)
        return self._unpack_contents(raw_string, bytes_read, version=version)

    def to_string(self, data={}):

        id_, version = self.validate_data_header(data)
        datagram_content_str = self._pack_contents(data, version=version)
        return self.finalize_datagram(datagram_content_str)

    def _unpack_contents(self, raw_string='', version=0):
        raise NotImplementedError

    def _pack_contents(self, data={}, version=0):
        raise NotImplementedError

    @classmethod
    def finalize_datagram(cls, datagram_content_str):
        datagram_size = len(datagram_content_str)
        final_fmt = '=l%dsl' % (datagram_size)
        return struct.pack(final_fmt, datagram_size, datagram_content_str, datagram_size)


class SimradUnknownParser(_SimradDatagramParser):
    """Parser for unknown datagram types.
    
    This parser only extracts the type and timestamp and
    returns the remainder of the data unparsed.
    """

    def __init__(self, dg_type):
        _SimradDatagramParser.__init__(self, dg_type, [0])

        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'data' / GreedyBytes()
            )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradSINParser(_SimradDatagramParser):
    """Parses SN90 system information datagrams"""
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "SIN", [0])

        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'transceiver_count' / Int32ul,
                'transceivers' / Array(this.transceiver_count, 
                                    Struct(
                                        'ip' / Int32ul,
                                        'port' / Int16ul,
                                        'name' / PaddedString(32, 'ascii')
                                    )
                                )
            )


    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data
    

class SimradVERParser(_SimradDatagramParser):
    """Parses SN90 version information datagram."""
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "VER", [0, 1])
        
        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'file_version' / PaddedString(32, 'ascii'),
                'software_version' / PaddedString(32, 'ascii'),
                'version_info' / PaddedString(64, 'ascii'),
                'product_name' / PaddedString(64, 'ascii')
            )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradPHYParser(_SimradDatagramParser):
    """Parses SN90 physical configuration datagrams"""
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "PHY", [0])

        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'platform_count' / Int32sl,
                'platforms' / Array(this.platform_count,
                    Struct(
                        'struct_size' / Int32sl,
                        'platform_type' / Int32sl,
                        'dimension' / If(this.platform_type == 0,
                        Struct(
                            'length' / Float32l,
                            'width' / Float32l,
                            'height' / Float32l
                        )
                        ),
                        'offset_from_centre' / If(this.platform_type == 0,
                            Struct(
                                'origin_offset_from_centre_x' / Float32l,
                                'origin_offset_from_centre_y' / Float32l,
                                'origin_offset_from_centre_z' / Float32l
                        )
                        ),
                        'name' / If(this.platform_type == 1, PaddedString(32, 'ascii')),
                        'parent_platform' / If(this.platform_type == 1, PaddedString(32, 'ascii')),
                        'rotation_x' / If(this.platform_type == 1, Float32l),
                        'rotation_y' / If(this.platform_type == 1, Float32l),
                        'rotation_z' / If(this.platform_type == 1, Float32l)
                    )
                )
            )
        
    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradPINParser(_SimradDatagramParser):
    """Parses SN90 ping information datagrams"""
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "PIN", [0, 1])
        
        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'ping_time' / Timestamp(Int64ul, 1e-7, 1600),
                'ping_number' / Int32sl,
                'latitude' / Float64l,
                'longitude' / Float64l,
                'speed'  / Float64l,
                'heading'  / Float64l,
                'heave' / Float64l,
                'roll' / Float64l,
                'pitch' / Float64l,
                'vessel_depth' / Float64l,
                'vessel_distance' / If(this.type == 'PIN1', Float64l),
                'transducer_offset_x' / Float64l,
                'transducer_offset_y' / Float64l,
                'transducer_offset_z' / Float64l,
                'relative_transducer_heading' / Float64l,
                'sound_velocity' / Float64l
            )

    def _unpack_contents(self, raw_string, bytes_read, version):

        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradEOPParser(_SimradDatagramParser):
    """Parses SN90 end of ping datagram."""
    
    def __init__(self):

        _SimradDatagramParser.__init__(self, "EOP", [0])

        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600)
            )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradSENParser(_SimradDatagramParser):
    """Parses SN90 sensor datagrams"""
    
    def __init__(self):

        _SimradDatagramParser.__init__(self, "SEN", [0])

        self.dg_def =\
            Struct(
                'type' / PaddedString(4, 'ascii'),
                'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
                'received_time' / Timestamp(Int64ul, 1e-7, 1600),
                'protocol' / PaddedString(32, 'ascii'),
                'port_name' / PaddedString(32, 'ascii'),
                'message_length' / Int32sl,
                'message' / IfThenElse(this.protocol == 'Nmea', 
                                        PaddedString(this.message_length, 'ascii'),
                                        Bytes(this.message_length))
            )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradPCOParser(_SimradDatagramParser):
    """Parses SN90 ping configuration datagrams"""
    
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "PCO", [0, 1])

        self.dg_def = Struct(
            'type' / PaddedString(4, 'ascii'),
            'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
            'ping_configuration' / Struct(
                'no_of_transceivers' / Int32sl,
                'transceiver_config' / Array(this.no_of_transceivers,
                    Struct(
                        'transceiver_config_size' / Int32sl,
                        'id' / Int32sl,
                        'transceiver_name_len' / Int16ul,
                        'transceiver_name' / PaddedString(this.transceiver_name_len*2, 'utf_16_le'),
                        'split_beam_percentage' /Int32sl,
                        'tx_config' / Struct(
                            'tx_configuration_size' / Int32sl,
                            'no_of_pings' / Int32sl,
                            'tx_ping_config' / Array(this.no_of_pings,
                                Struct(
                                    'tx_ping_config_size' / Int32sl,
                                    'id' / Int32sl,
                                    'ping_name_len' / Int16ul,
                                    'ping_name' / PaddedString(this.ping_name_len*2, 'utf_16_le'),
                                    'frequency' / Float32l,
                                    'pulse_duration' /Float32l,
                                    'pulse_form' / Int32sl,
                                    'pulse_sweep' / Float32l,
                                    'pulse_slope' / Float32l,
                                    'focus_x' / Int32sl,
                                    'focus_y' / Int32sl,
                                    'beam_width_x' / Float32l,
                                    'beam_width_y' / Float32l,
                                    'steering_x' / Float32l,
                                    'steering_y' / Float32l,
                                    'beam_delay' / Float32l,
                                    'tx_amplitude' / Float32l,
                                    'tx_voltage' / Float32l,
                                    'actual_beam_bandwidth_rx' / Float64l,
                                    'decimation' / Int32sl,
                                    'range' / Float64l,
                                    'steering_vector_hcs_x' / Float32l,
                                    'steering_vector_hcs_y' / Float32l,
                                    'steering_vector_hcs_z' / Float32l,
                                    'rotation_axis_vector_x' / Float32l,
                                    'rotation_axis_vector_y' / Float32l,
                                    'rotation_axis_vector_z' / Float32l,
                                    'tx_ping_weight_x_len' / Int32sl,
                                    'tx_ping_weight_x' / Array(this.tx_ping_weight_x_len, Float32l),
                                    'tx_ping_weight_y_len' / Int32sl,
                                    'tx_ping_weight_y' / Array(this.tx_ping_weight_y_len, Float32l),
                                    'performance_info' / Struct(
                                        'tx_ping_performance_info_size' / Int32sl,
                                        'tx_power' / Float32l,
                                        'source_level' / Float32l
                                    )
                                )
                            )
                        ),
                        'rx_config' / Struct(
                            'rx_configuration_size' / Int32sl,
                            'audio_beam_index' / Int32sl,
                            'no_of_fans' / Int32sl,
                            'fans' / Array(this.no_of_fans,
                                Struct(
                                    'fan_config_size' / Int32sl,
                                    'id' / Int32sl,
                                    'fan_name_len' / Int16ul,
                                    'fan_name' / PaddedString(this.fan_name_len*2, 'utf_16_le'),
                                    'sample_interval' / Float64l,
                                    'tx_ping_id' / Int32sl,
                                    'main_beam_rx_weight_x_len' / Int32sl,
                                    'main_beam_rx_weight_x' / Array(this.main_beam_rx_weight_x_len, Float32l),
                                    'main_beam_rx_weight_y_len' / Int32sl,
                                    'main_beam_rx_weight_y' / Array(this.main_beam_rx_weight_y_len, Float32l),
                                    'split_beam_rx_weight_x_len' / Int32sl,
                                    'split_beam_rx_weight_x' / Array(this.split_beam_rx_weight_x_len, Float32l),
                                    'split_beam_rx_weight_y_len' / Int32sl,
                                    'split_beam_rx_weight_y' / Array(this.split_beam_rx_weight_y_len, Float32l),
                                    'noise_filter' / Int32sl,
                                    'processing' / Struct(
                                        'fan_processing_size' / Int32sl,
                                        'tvg_a' / Float64l,
                                        'tvg_b' / Float64l,
                                        'tvg_c' / Float64l,
                                        'rcg' / Int32sl,
                                        'agc' / Int32sl,
                                        'amp_gain' / Int32sl
                                    ),
                                    'no_of_rx_beams' / Int32sl,
                                    'rx_beams' / Array(this.no_of_rx_beams, 
                                        Struct(
                                            'rx_beam_config_size' / Int32sl,
                                            'id' / Int32sl,
                                            'beam_name_len' / Int16ul,
                                            'beam_name' / PaddedString(this.beam_name_len*2, 'utf_16_le'),
                                            'beam_width_x' / Float32l,
                                            'beam_width_y' / Float32l,
                                            'steering_x' / Float32l,
                                            'steering_y' / Float32l,
                                            'beam_type' / Int32sl,
                                            'steering_vector_hcs_x' / Float32l,
                                            'steering_vector_hcs_y' / Float32l,
                                            'steering_vector_hcs_z' / Float32l,
                                            'processing_type' / Int32sl,
                                            'performance_info' / Struct(
                                                'rx_beam_performance_info_size' / Int32sl,
                                                'directivity_index' / Float32l,
                                                'gain' / Float32l,
                                                'gain_adjust' / Float32l,
                                                'sa_correction' / Float32l,
                                                'sa_correction_adjust' / Float32l,
                                                'equivalent_beam_angle' / Float32l,
                                                'absorption_coefficient' / Float32l,
                                                'angle_sensitivity_alongship' / Float32l,
                                                'angle_sensitivity_athwartship' / Float32l
                                            ),
                                            'rx_delay' / If (this._._._._._.type == 'PCO1', Int32sl)
                                        )
                                    ),
                                    'rx_delay' / If (this._._._._.type == 'PCO1', Int32sl) 
                                )
                            )
                        ),
                        'transmission_mode' / Int32sl
                    )
                ),
                'hint_since_last_len' / Int16ul,
                'hint_since_last' / PaddedString(this.hint_since_last_len*2, 'utf_16_le'),
                'hint_since_last_ping_len' / Int16ul,
                'hint_since_last_ping' / PaddedString(this.hint_since_last_ping_len*2, 'utf_16_le')
            )
        )
        
    def _unpack_contents(self, raw_string, bytes_read, version):

        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradSECParser(_SimradDatagramParser):
    """Parses SN90 sensor configuration datagrams"""
    
    def __init__(self):
        _SimradDatagramParser.__init__(self, "SEC", [0])

        self.dg_def = Struct(
            'type' / PaddedString(4, 'ascii'),
            'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
            'xml' / GreedyString('utf_8')
        )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)
        return data


class SimradRAWParser(_SimradDatagramParser):
    """Parse Simrad RAW datagrams."""

    def __init__(self):
        _SimradDatagramParser.__init__(self, "RAW", [2])

        self.dg_def = Struct(
            'type' / PaddedString(4, 'ascii'),
            'timestamp' / Timestamp(Int64ul, 1e-7, 1600),
            'ip_address' / Int32ul,
            'port' / Int16ul,
            'padding' / Int16ul,
            'message_length' / Int32sl,
            'length' / Int32ul,
            'id' / PaddedString(4, 'ascii'),
            'time' / Timestamp(Int64ul, 1e-7, 1600),
            'ping_number' / Int32ul,
            'datagram_number' / Int32ul,
            'sample_index' / Int32sl,
            'no_of_samples' / Int32ul,
            'heave' / Float32l,
            'roll' / Float32l,
            'pitch' / Float32l,
            'yaw' / Float32l,
            'beam_index_start' / Int16ul,
            'spare' / Bytes(2),
            'no_of_beams' / Int32ul,
            'data' / Array(this.no_of_samples * this.no_of_beams * 2, Float32l
            )
        )

    def _unpack_contents(self, raw_string, bytes_read, version):
        data = self.dg_def.parse(raw_string)
        data = construct_to_dict(data)

        real = np.array(data['data'][0::2])
        imag = np.array(data['data'][1::2])

        data['data'] = (real + 1.0j * imag).reshape(data['no_of_beams'],
                                                    data['no_of_samples'])

        return data
