"""ASoCv3Parser

Use parser to parse the data from the Hardware.
The boards sends a raw bitstream and it needs to be parsed into events.
The parser class is a tool to parse the raw bitstream of 8-bit chunks
into 16 bit words then extracting the 12-bit data, the header info,
channels and window number data.

This data is then returned in a dictionary.
The raw data is preserved.

This module was previously located under the naludaq.daq.workers
subpackage, and was relocated to improve the organization of the
naludaq package.
"""

import logging

import numpy as np

from naludaq.parsers.parser import Parser

LOGGER = logging.getLogger(__name__)


class ASoCv3Parser(Parser):
    def __init__(self, params):
        super().__init__(params)

    def parse_digital_data_old(self, in_data) -> dict:
        """Parse the raw data from the board.
        Assuming that the data packets are constant length, we can extract the data
        in place with matrix operations, speeding up the parsing.

        Parses the data spit out by the digitial readout portion of the firmware

        Channel data is interleaved, e.g
        pkt0,4,8,.... - ch0,
        pkt1,5,9,.... - ch1,
        etc

        TODO: To speed this up more, we do not want two matrices (bitshift & raw),
        we would perferably only want one.

        Args:
            in_data (bytearray): Raw data from the board
        Returns:
            Parsed event as a dict.
        Raises:
            BadDataException if no data is found or if the data contains errors.

        Additional info:
        The word is 16-bit and the data is 12-bit.
        The data is in the form:
        ZZZX XXXX XXXX XXXZ
        Where Z are unused bits.
        0x00 is not a valid byte, ever

        Example format for ASoC:
        first three bits are the "identifier", last bit is always 1
        4 x event headers:
        '001' 'prev_final_window' '1'
        '001' 'trigger_time_ns(23 dt 12)' '1'
        '001' 'trigger_time_ns(11 dt 0 )' '1'
        '001' 'eventNumber' '1'
        1x window header:
        '010' 'chan' '00' 'window' '1'
        64x data packets
        '100' 'data' '1'
        1x end packet (x"FACE")

        """
        return self._parse_event(in_data)
    
    def _parse_event(self, in_data) -> dict:
        raw_data = np.frombuffer(in_data["rawdata"], dtype=">H")
        num_headers = self.params.get("headers", 3)
        wlc_stride: "int | None" = self._get_wlc_stride(in_data["rawdata"])
        is_wlc_event: bool = wlc_stride is not None
        # The raw_data is in the form:
        # ZZZX XXXX XXXX XXXZ
        # Where Z are unused bits.

        # Bitshifted data will make it in the form:
        # 0000 XXXX XXXX XXXX
        raw_event_matrix = self._reshape_data(raw_data, is_wlc_event, wlc_stride)
        
        packet_size = self.num_wind_header + self.params["samples"]
        # Pull channel info from all packets
        channels = (raw_event_matrix[0, :] & self.params["chanmask"]) >> self.params[
            "chanshift"
        ]
        available_channels = np.unique(channels[: self.params["channels"]])
        num_windows = raw_event_matrix.shape[1] // len(available_channels)
        curr_event = self._get_event_template(num_windows)

        # Parse headers
        self._parse_event_headers(curr_event, self._bitshift_data(raw_data[:num_headers]))
        first_window_header_idx = num_headers + 1 if is_wlc_event else num_headers
        curr_event["start_window"] = np.uint16(
            (raw_data[first_window_header_idx] & self.params["windmask"])
            >> self.last_bits
        )

        # Parse event data
        for idx, chan in enumerate(available_channels):
            curr_event["data"][chan] = (
            (
            raw_event_matrix[1:packet_size, idx :: len(available_channels)].flatten(order="F") 
            ) & self.data_bitmask
            ) >> self.last_bits
            curr_event["window_labels"][chan] = (
                (
                    raw_event_matrix[0, idx :: len(available_channels)]
                    & self.params["windmask"]
                )
                >> self.last_bits
            )

        if is_wlc_event:
            wlc_headers = (
                raw_data[num_headers:-self.num_footers:wlc_stride]
                & self.params.get("wlcmask", 510)
            ) >> self.last_bits
            curr_event["wlc_headers"] = wlc_headers

        return curr_event


    def _get_event_template(self, num_windows: int) -> dict:
        """Returns a template for the event dictionary."""
        data = np.full((self.params["channels"], num_windows * self.params["samples"]), fill_value=np.nan, dtype=np.float32)
        window_labels = np.full((self.params["channels"], num_windows), fill_value=np.nan, dtype=np.float16)
        event_template = {
            "window_labels": window_labels,
            "data": data,
        }
        return event_template
    
    def _reshape_data(self, data: np.ndarray, is_wlc_event=False, wlc_stride=None) -> np.ndarray:
        """
        Reshapes the data into an np.array of shape (packet_size, num_packets).
        Will ignore headers, footers, and wlc_headers.
        """
        headers = self.params.get("headers", 3)
        data_end = len(data) - self.num_footers
        packet_size = self.num_wind_header + self.params["samples"]
        if is_wlc_event:
            num_channels = (wlc_stride - 1) // packet_size
            num_windows = (len(data) - headers - self.num_footers + 1) // wlc_stride
            num_packets = num_windows * num_channels
            data_by_wlc_blocks = np.reshape(
                data[headers:-1],
                [num_windows, wlc_stride],
            )
            reshaped_data = np.reshape(
                data_by_wlc_blocks[:, 1:].flatten(),
                [packet_size, num_packets],
                order="F",
            )
        else:
            event_size = len(data) - self.num_footers - headers + 1
            num_packets = event_size // packet_size
            reshaped_data = np.reshape(
                data[headers : data_end + 1],
                [packet_size, num_packets],
                order="F",
            )
        
        return reshaped_data

    def _get_wlc_stride(self, raw_data: bytearray) -> "int | None":
        """Checks the length between the first and second occurrence of 0x9E."""
        possible_strides = [(self.params["samples"] + self.num_wind_header) * (chan + 1) + 1 for chan in range(self.params["channels"])]
        try:
            start = raw_data[::2].index(0x9E)
            for possible_stride in possible_strides:
                if raw_data[2 * (start + possible_stride)] == 0x9E:
                    stride = possible_stride
                    break
            else:
                stride = None
        except ValueError:
            stride = None
        return stride

    def _add_xaxis_to_event(self, event):
        """Adds an x-axis to the data.

        Based on the amount of channels and samples it will add a timeaxis
        per channel to the event dict.

        Returns:
            numpy array with sample numbers for each channel.
        """
        if "wlc_headers" in event.keys():
            samples = self.samples
            channels = self.channels
            tim = (samples * np.repeat(event["wlc_headers"], samples)) + (
                np.ones((len(event["wlc_headers"]), samples), dtype=int)
                * np.arange(0, samples, dtype=int)
            ).flatten()
            times = np.tile(tim, (channels, 1))
        else:
            times = super()._add_xaxis_to_event(event)
        return times