"""
Class for reading/writing data from micromed (.trc).
Inspired by the Matlab code for EEGLAB from Rami K. Niazy.

Completed with matlab Guillaume BECQ code.

Author: Samuel Garcia
"""

import datetime
import struct
import io

import numpy as np

from .baserawio import (
    BaseRawIO,
    _signal_channel_dtype,
    _signal_stream_dtype,
    _spike_channel_dtype,
    _event_channel_dtype,
)

from neo.core import NeoReadWriteError


class StructFile(io.BufferedReader):
    def read_f(self, fmt, offset=None):
        if offset is not None:
            self.seek(offset)
        return struct.unpack(fmt, self.read(struct.calcsize(fmt)))


class MicromedRawIO(BaseRawIO):
    """
    Class for reading  data from micromed (.trc).

    Parameters
    ----------
    filename: str, default: None
        The *.trc file to be loaded
    """

    extensions = ["trc", "TRC"]
    rawmode = "one-file"

    def __init__(self, filename=""):
        BaseRawIO.__init__(self)
        self.filename = filename

    def _parse_header(self):
        with open(self.filename, "rb") as fid:
            f = StructFile(fid)

            # Name
            f.seek(64)
            surname = f.read(22).strip(b" ")
            firstname = f.read(20).strip(b" ")

            # Date
            day, month, year, hour, minute, sec = f.read_f("bbbbbb", offset=128)
            rec_datetime = datetime.datetime(year + 1900, month, day, hour, minute, sec)

            Data_Start_Offset, Num_Chan, Multiplexer, Rate_Min, Bytes = f.read_f("IHHHH", offset=138)

            # header version
            (header_version,) = f.read_f("b", offset=175)
            if header_version != 4:
                raise NotImplementedError(f"`header_version {header_version} is not implemented in neo yet")

            # area
            f.seek(176)
            zone_names = [
                "ORDER",
                "LABCOD",
                "NOTE",
                "FLAGS",
                "TRONCA",
                "IMPED_B",
                "IMPED_E",
                "MONTAGE",
                "COMPRESS",
                "AVERAGE",
                "HISTORY",
                "DVIDEO",
                "EVENT A",
                "EVENT B",
                "TRIGGER",
            ]
            zones = {}
            for zname in zone_names:
                zname2, pos, length = f.read_f("8sII")
                zones[zname] = zname2, pos, length
                if zname != zname2.decode("ascii").strip(" "):
                    raise NeoReadWriteError("expected the zone name to match")

            # raw signals memmap
            sig_dtype = "u" + str(Bytes)
            self._raw_signals = np.memmap(self.filename, dtype=sig_dtype, mode="r", offset=Data_Start_Offset).reshape(
                -1, Num_Chan
            )

            # Reading Code Info
            zname2, pos, length = zones["ORDER"]
            f.seek(pos)
            code = np.frombuffer(f.read(Num_Chan * 2), dtype="u2")

            units_code = {-1: "nV", 0: "uV", 1: "mV", 2: 1, 100: "percent", 101: "dimensionless", 102: "dimensionless"}
            signal_channels = []
            sig_grounds = []
            for c in range(Num_Chan):
                zname2, pos, length = zones["LABCOD"]
                f.seek(pos + code[c] * 128 + 2, 0)

                chan_name = f.read(6).strip(b"\x00").decode("ascii")
                ground = f.read(6).strip(b"\x00").decode("ascii")
                sig_grounds.append(ground)
                logical_min, logical_max, logical_ground, physical_min, physical_max = f.read_f("iiiii")
                (k,) = f.read_f("h")
                units = units_code.get(k, "uV")

                factor = float(physical_max - physical_min) / float(logical_max - logical_min + 1)
                gain = factor
                offset = -logical_ground * factor

                f.seek(8, 1)
                (sampling_rate,) = f.read_f("H")
                sampling_rate *= Rate_Min
                chan_id = str(c)
                stream_id = "0"
                signal_channels.append((chan_name, chan_id, sampling_rate, sig_dtype, units, gain, offset, stream_id))

            signal_channels = np.array(signal_channels, dtype=_signal_channel_dtype)

            signal_streams = np.array([("Signals", "0")], dtype=_signal_stream_dtype)

            if np.unique(signal_channels["sampling_rate"]).size != 1:
                raise NeoReadWriteError("The sampling rates must be the same across signal channels")
            self._sampling_rate = float(np.unique(signal_channels["sampling_rate"])[0])

            # Event channels
            event_channels = []
            event_channels.append(("Trigger", "", "event"))
            event_channels.append(("Note", "", "event"))
            event_channels.append(("Event A", "", "epoch"))
            event_channels.append(("Event B", "", "epoch"))
            event_channels = np.array(event_channels, dtype=_event_channel_dtype)

            # Read trigger and notes
            self._raw_events = []
            ev_dtypes = [
                ("TRIGGER", [("start", "u4"), ("label", "u2")]),
                ("NOTE", [("start", "u4"), ("label", "S40")]),
                ("EVENT A", [("label", "u4"), ("start", "u4"), ("stop", "u4")]),
                ("EVENT B", [("label", "u4"), ("start", "u4"), ("stop", "u4")]),
            ]
            for zname, ev_dtype in ev_dtypes:
                zname2, pos, length = zones[zname]
                dtype = np.dtype(ev_dtype)
                rawevent = np.memmap(self.filename, dtype=dtype, mode="r", offset=pos, shape=length // dtype.itemsize)

                keep = (
                    (rawevent["start"] >= rawevent["start"][0])
                    & (rawevent["start"] < self._raw_signals.shape[0])
                    & (rawevent["start"] != 0)
                )
                rawevent = rawevent[keep]
                self._raw_events.append(rawevent)

            # No spikes
            spike_channels = []
            spike_channels = np.array(spike_channels, dtype=_spike_channel_dtype)

            # fille into header dict
            self.header = {}
            self.header["nb_block"] = 1
            self.header["nb_segment"] = [1]
            self.header["signal_streams"] = signal_streams
            self.header["signal_channels"] = signal_channels
            self.header["spike_channels"] = spike_channels
            self.header["event_channels"] = event_channels

            # insert some annotation at some place
            self._generate_minimal_annotations()
            bl_annotations = self.raw_annotations["blocks"][0]
            seg_annotations = bl_annotations["segments"][0]

            for d in (bl_annotations, seg_annotations):
                d["rec_datetime"] = rec_datetime
                d["firstname"] = firstname
                d["surname"] = surname
                d["header_version"] = header_version

            sig_annotations = self.raw_annotations["blocks"][0]["segments"][0]["signals"][0]
            sig_annotations["__array_annotations__"]["ground"] = np.array(sig_grounds)

    def _source_name(self):
        return self.filename

    def _segment_t_start(self, block_index, seg_index):
        return 0.0

    def _segment_t_stop(self, block_index, seg_index):
        t_stop = self._raw_signals.shape[0] / self._sampling_rate
        return t_stop

    def _get_signal_size(self, block_index, seg_index, stream_index):
        if stream_index != 0:
            raise ValueError("`stream_index` must be 0")
        return self._raw_signals.shape[0]

    def _get_signal_t_start(self, block_index, seg_index, stream_index):
        if stream_index != 0:
            raise ValueError("`stream_index` must be 0")
        return 0.0

    def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes):
        if channel_indexes is None:
            channel_indexes = slice(channel_indexes)
        raw_signals = self._raw_signals[slice(i_start, i_stop), channel_indexes]
        return raw_signals

    def _spike_count(self, block_index, seg_index, unit_index):
        return 0

    def _event_count(self, block_index, seg_index, event_channel_index):
        n = self._raw_events[event_channel_index].size
        return n

    def _get_event_timestamps(self, block_index, seg_index, event_channel_index, t_start, t_stop):

        raw_event = self._raw_events[event_channel_index]

        if t_start is not None:
            keep = raw_event["start"] >= int(t_start * self._sampling_rate)
            raw_event = raw_event[keep]

        if t_stop is not None:
            keep = raw_event["start"] <= int(t_stop * self._sampling_rate)
            raw_event = raw_event[keep]

        timestamp = raw_event["start"]
        if event_channel_index < 2:
            durations = None
        else:
            durations = raw_event["stop"] - raw_event["start"]

        try:
            labels = raw_event["label"].astype("U")
        except UnicodeDecodeError:
            # sometimes the conversion do not work : here a simple fix
            labels = np.array([e.decode("cp1252") for e in raw_event["label"]], dtype="U")

        return timestamp, durations, labels

    def _rescale_event_timestamp(self, event_timestamps, dtype, event_channel_index):
        event_times = event_timestamps.astype(dtype) / self._sampling_rate
        return event_times

    def _rescale_epoch_duration(self, raw_duration, dtype, event_channel_index):
        durations = raw_duration.astype(dtype) / self._sampling_rate
        return durations
