import csv
import logging
import math
import os
import struct
import typing

import nptdms
import numpy as np
from scipy.io import wavfile

from dr_files.pb.headers_pb2 import MeasurementHeader

LINE_BUFFER = 1000
MAX_SHORT = 32767.0
ENGINEERING_UNITS_MAP = {
    "custom": 0,
    "V(default)": 1,
    "Pa": 2,
    "g": 3,
    "m/s^2": 4,
    "m/s": 5,
    "m": 6,
    "in/s^2": 7,
    "in/s": 8,
    "in": 9,
    "mil/s": 10,
    "mil": 11,
    "mm/s": 12,
    "mm": 13,
    "micron/s": 14,
    "micron": 15,
    "rad/s^2": 16,
    "rad/s": 17,
    "rad": 18,
    "strain": 19,
    "N": 20,
    "A": 21,
    "RPM": 22,
    "cycle/s^2": 23,
    "Hz": 24,
    "cycle": 25,
    "deg/s^2": 26,
    "deg/s": 27,
    "deg": 28,
}


def dr_to_int_stream(
    file_path: os.PathLike,
) -> (MeasurementHeader, typing.List[int]):
    """
    Reads a Mechbase file found in file_path, returning a tuple with the header and a stream of
    integers.
    """
    with open(file_path, "rb") as fin:
        # Read header
        header: MeasurementHeader = _proto_read(fin, MeasurementHeader)
        data = fin.read()
        sample_width = _detect_sample_width(data)
        return header, list(_read_int_data(data, sample_width))


def dr_to_raw_values(
    file_path: str,
) -> (MeasurementHeader, typing.List[np.array]):
    """
    Reads a Mechbase file found in file_path, returning a tuple with the header and a list of numpy
    arrays with the integer values of the signals.
    """
    header, int_stream = dr_to_int_stream(file_path)

    signals = list(map(lambda _: list(), range(len(header.channels))))
    for index, data_point in enumerate(int_stream):
        signals[index % len(signals)].append(data_point)

    # Convert signal and return
    return header, list(map(np.array, signals))


def dr_to_values(file_path: str) -> (MeasurementHeader, typing.List[np.array]):
    """
    Reads a Mechbase file found in file_path, returning a tuple with the header and a list of numpy
    arrays with the actual measurement values.
    """
    header, raw_signals = dr_to_raw_values(file_path)
    signals = list()
    converters = list(map(value_converter, header.channels))

    for index, s in enumerate(raw_signals):
        signals.append(
            np.array(list(map(converters[index % len(converters)], raw_signals[index])))
        )
    return header, signals


def dr_to_csv(file_path: os.PathLike, out_file_path: os.PathLike):
    """
    Reads a Mechbase file found in file_path, converts it to CSV and writes it at out_file_path.
    """
    header, signals = dr_to_values(file_path)
    h_row = ["time"]
    for i in range(len(header.channels)):
        h_row.append(f"channel {i + 1}")
    with open(out_file_path, "w+") as fout:
        w = csv.writer(fout)
        w.writerow(h_row)
        rows = []

        time = 0
        fs = 1 / header.sample_rate
        for signal_values in zip(*signals):
            row = [time]
            row.extend(signal_values)
            time += fs

            rows.append(list(map(lambda value: format(value, ".10f"), row)))
            if len(rows) < LINE_BUFFER:
                continue

            w.writerows(rows)
            rows = []

        w.writerows(rows)


def dr_to_wav(file_path: os.PathLike, out_file_path: os.PathLike):
    """
    Reads a Mechbase file found in file_path, returning a tuple with the header and a numpy array of the signal.
    """
    header, raw_signals = dr_to_raw_values(file_path)
    if len(header.channels) > 1:
        logging.warning(
            f'Converting dr file "{out_file_path}" to wav "{out_file_path}", but dr file has more than a channel'
        )
    wavfile.write(out_file_path, header.sample_rate, np.array(raw_signals[0], np.int16))


def dr_to_tdms(
    file_path: os.PathLike,
    out_file_path: os.PathLike,
    engineering_units_code: str = "",
):
    header, raw_signals = dr_to_raw_values(file_path)
    channels = header.channels

    signals = []
    for index, signal in enumerate(raw_signals):
        multiplier = channels[index].reference / MAX_SHORT
        signals.append(np.array(list(map(lambda v: v * multiplier, signal))))

    root_obj = nptdms.RootObject()
    group_obj = nptdms.GroupObject("MEASUREMENT")
    channel_objs = []
    engineering_unit_label = engineering_units_code
    engineering_units = 0
    if engineering_units_code in ENGINEERING_UNITS_MAP:
        engineering_units = ENGINEERING_UNITS_MAP[engineering_units_code]

    for index, signal in enumerate(signals):
        channel_obj = nptdms.ChannelObject(
            "MEASUREMENT",
            str(index),
            data=signal,
            properties={
                "db_reference": channels[index].db_reference,
                "engineering_unit_label": engineering_unit_label,
                "engineering_units": engineering_units,
                "offset": channels[index].offset,
                "pregain": channels[index].pregain,
                "sensitivity": channels[index].sensitivity,
                "wf_increment": 1.0 / header.sample_rate,
                "wf_samples": len(signal),
            },
        )
        channel_objs.append(channel_obj)
    with nptdms.TdmsWriter(out_file_path) as tdms_file:
        channel_objs.insert(0, group_obj)
        channel_objs.insert(0, root_obj)
        tdms_file.write_segment(channel_objs)


def value_converter(header: MeasurementHeader.ChannelHeader):
    """
    Returns a generator function, which converts a stream of integer values to actual measurement
    values using the given header.
    """
    multiplier = header.reference / MAX_SHORT

    def fun(value):
        offset = header.offset
        if not offset:
            offset = 0.0
        value = (value * multiplier + offset) * 1000.0 / header.sensitivity
        if header.db_reference > 0.0:
            if value == 0:
                value = 0.0000001
            value = 20 * math.log10(abs(value) / header.db_reference)
            if header.pregain > 0.0:
                value += header.pregain
        return value

    return fun


def float_value_converter(header: MeasurementHeader.ChannelHeader):
    """
    Returns a generator function, which converts a stream of integer values to actual measurement
    values using the given header.
    """
    multiplier = header.reference / MAX_SHORT

    def fun(value):
        offset = header.offset
        if header.db_reference > 0.0:
            if value == 0:
                value = 0.0000001
            if header.pregain > 0.0:
                value -= header.pregain
            value = pow(10, (value / 20)) * header.db_reference
        if not offset:
            offset = 0.0
        value = ((header.sensitivity * value) / 1000.0 - offset) / multiplier
        return value

    return fun


def _detect_sample_width(data: bytes) -> int:
    """
    Detects whether signal data uses 16-bit or 32-bit integers.

    Some acquisition devices write 32-bit integers instead of 16-bit. When read
    as int16 pairs, 32-bit values produce a pattern where every other sample is
    0 (for positive values) or -1 (for negative values, due to sign extension).
    """
    if len(data) < 8:
        return 2

    int16_values = np.frombuffer(data, dtype="<i2")
    if len(int16_values) < 4:
        return 2

    odd_values = int16_values[1::2]
    zeros_or_sign_ext = np.sum((odd_values == 0) | (odd_values == -1))
    if zeros_or_sign_ext / len(odd_values) > 0.95:
        return 4

    return 2


def _read_int_data(data: bytes, sample_width: int):
    """
    Reads signal data as a stream of integers with the given sample width.
    """
    dtype = "<i2" if sample_width == 2 else "<i4"
    usable = len(data) - (len(data) % sample_width)
    values = np.frombuffer(data[:usable], dtype=dtype)
    yield from values


def _read_dr_int_stream(fin):
    """
    Reads a file and converts it to a stream of bytes.
    """
    data = fin.read()
    sample_width = _detect_sample_width(data)
    yield from _read_int_data(data, sample_width)


def _proto_read(stream, clazz):
    bt = stream.read(2)
    proto_size = struct.unpack("<h", bt)[0]
    proto = clazz()
    msg_buf = stream.read(proto_size)
    proto.ParseFromString(msg_buf)
    return proto


def _proto_write(stream, proto):
    bt = proto.SerializeToString()
    stream.write(struct.pack("<h", len(bt)))
    stream.write(bt)


def write_int_data(stream, data):
    stream.write(int_to_bytes(data))


def int_to_bytes(number: int) -> bytes:
    return struct.pack("<h", number)


class DrWriter:
    def __init__(self, stream) -> None:
        self.stream = stream

    def write_header(self, header: dict) -> None:

        self._channels_count = len(header["channels"])
        measurement = MeasurementHeader()
        measurement.sample_rate = header["sample_rate"]
        for channel in header["channels"]:
            measurement.channels.append(
                MeasurementHeader.ChannelHeader(
                    reference=channel["reference"],
                    sensitivity=channel["sensitivity"],
                    db_reference=channel["db_reference"],
                    pregain=channel["pregain"],
                    offset=channel["pregain"],
                )
            )
        size = measurement.ByteSize()
        _proto_write(self.stream, measurement)

    def write_row(self, row) -> None:
        if len(row) != self._channels_count:
            raise ValueError("Length of row must match the length of channels")
        for item in row:
            write_int_data(self.stream, item)
