#
# ------------------------------------------------------------
# Copyright (c) All rights reserved
# SiLab, Institute of Physics, University of Bonn
# ------------------------------------------------------------
#

'''
    Online data analysis functions
'''

import ctypes
import logging
import multiprocessing
import time
import queue

import numpy as np
from numba import njit

logger = logging.getLogger('OnlineAnalysis')


# Error codes
TRIGGER_NUMBER_ERROR = 0x00000001  # Trigger number has not increased by one
NO_TRIGGER_WORD_ERROR = 0x00000002  # Event has no trigger word associated
TRIGGER_TIMESTAMP_OVERFLOW = 0x00000004  # Indicating the overflow of the trigger timestamp
TRIGGER_NUMBER_OVERFLOW = 0x00000008  # Indicating the overflow of the trigger number
DATA_ERROR = 0x00000010  # Indicating any occurrence of data errors in the Momosa26 protocol (e.g., invalid column/row, invalid data length, data loss)
TIMESTAMP_OVERFLOW = 0x00000020  # Indicating the overflow of the Mimosa26 timestamp
FRAME_ID_OVERFLOW = 0x00000040  # Indicating the overflow of the Mimosa26 frame ID
OVERFLOW_FLAG = 0x00000080  # Indicating the occurrence of the overflow flag for a particular Mimosa26 row


# Mimosa26 raw data
@njit
def is_mimosa_data(word):  # Check for Mimosa data word
    return (0xff000000 & word) == 0x20000000


@njit
def get_plane_number(word):  # There are 6 planes in the stream, starting from 1; return plane number
    return (word >> 20) & 0xf


# Frame header
@njit
def is_frame_header(word):  # Check if frame header high word (frame start flag is set by R/0)
    return (0x00010000 & word) == 0x00010000


@njit
def is_data_loss(word):  # Indicates data loss
    return (0x00020000 & word) == 0x00020000


@njit
def get_m26_timestamp_low(word):  # Timestamp of Mimosa26 data from frame header low (generated by R/0)
    return 0x0000ffff & word


@njit
def get_m26_timestamp_high(word):  # Timestamp of Mimosa26 data from frame header high (generated by R/0)
    return (0x0000ffff & word) << 16


@njit
def is_frame_header0(word):  # Check if frame header0 word
    return (0x0000ffff & word) == 0x00005555


@njit
def is_frame_header1(word, plane):  # Check if frame header1 word for the actual plane
    return (0x0000ffff & word) == (0x00005550 | plane)


# Frame counter
@njit
def get_frame_id_low(word):  # Get the frame id from the frame id low word
    return 0x0000ffff & word


@njit
def get_frame_id_high(word):  # Get the frame id from the frame id high word
    return (0x0000ffff & word) << 16


# Data length
@njit
def get_frame_length(word):  # Get length of Mimosa26 frame
    return (0x0000ffff & word)


# Status / line word
@njit
def get_n_words(word):  # Return the number of data words for the actual row
    return 0x0000000f & word


@njit
def get_row(word):  # Extract row from Mimosa26 hit word
    return (0x00007ff0 & word) >> 4


@njit
def has_overflow(word):
    return (0x00008000 & word) != 0


# State word
@njit
def get_n_hits(word):  # Returns the number of hits given by actual column word
    return 0x00000003 & word


@njit
def get_column(word):  # Extract column from Mimosa26 hit word
    return (0x00001ffc & word) >> 2


# Frame trailer
@njit
def is_frame_trailer0(word):  # Check if frame trailer0 word
    return (0x0000ffff & word) == 0xaa50


@njit
def is_frame_trailer1(word, plane):  # Check if frame trailer1 word for the actual plane
    return (0x0000ffff & word) == (0xaa50 | plane)


# Trigger words
@njit
def is_trigger_word(word):  # Check if TLU word (trigger)
    return (0x80000000 & word) == 0x80000000


@njit
def histogram(raw_data, occ_hist, m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids):
    ''' Raw data to 2D occupancy histogram '''

    # Loop over the raw data words
    for raw_data_word in raw_data:
        if is_mimosa_data(raw_data_word):  # Check if word is from Mimosa26.
            # Check to which plane the data belongs
            plane_id = get_plane_number(raw_data_word) - 1  # The actual_plane if the actual word belongs to (0 to 5)
            # In the following, interpretation of the raw data words of the actual plane
            # Check for data loss bit set by the M26 RX FSM
            if is_data_loss(raw_data_word):
                # Setting the data loss flag to true.
                # The data loss bit is set by the M26 RX FSM.
                # The bit is set only once after each data loss, i.e.,
                # the first data word after the lost data words.
                m26_data_loss[plane_id] = True
            if is_frame_header(raw_data_word):  # New frame for actual plane, M26 timestamp (LSB), frame header0
                # Get Mimosa26 timestamp from raw data word (LSB)
                last_m26_timestamps[plane_id] = m26_timestamps[plane_id]
                m26_timestamps[plane_id] = (m26_timestamps[plane_id] & 0x7fffffffffff0000) | get_m26_timestamp_low(raw_data_word)
                m26_word_index[plane_id] = 0
                # Reset parameters after header
                m26_frame_length[plane_id] = 0
                m26_n_words[plane_id] = 0
                m26_data_loss[plane_id] = False
                m26_frame_status[plane_id] = 0
            elif m26_data_loss[plane_id] is True:  # Trash data
                # Nothing to do, do not trust data
                continue
            else:  # Interpreting M26 raw data
                m26_word_index[plane_id] += 1
                if m26_word_index[plane_id] == 1:  # Mimosa26 timestamp, M26 timestamp (MSB), frame header1
                    # Check for 32bit timestamp overflow
                    if m26_timestamps[plane_id] >= 0 and get_m26_timestamp_high(raw_data_word) < (m26_timestamps[plane_id] & 0x00000000ffff0000):
                        m26_frame_status[plane_id] |= TIMESTAMP_OVERFLOW
                        m26_timestamps[plane_id] = np.int64(2**32) + m26_timestamps[plane_id]
                    # Get Mimosa26 timestamp from raw data word (MSB)
                    m26_timestamps[plane_id] = get_m26_timestamp_high(raw_data_word) | (m26_timestamps[plane_id] & 0x7fffffff0000ffff)
                elif m26_word_index[plane_id] == 2:  # Mimosa26 frame ID
                    # Get Mimosa26 frame ID from raw data word (LSB)
                    m26_frame_ids[plane_id] = (m26_frame_ids[plane_id] & 0x7fffffffffff0000) | get_frame_id_low(raw_data_word)
                elif m26_word_index[plane_id] == 3:  # Mimosa26 frame ID
                    # Check for 32bit frame ID overflow
                    if m26_frame_ids[plane_id] >= 0 and get_frame_id_high(raw_data_word) < (m26_frame_ids[plane_id] & 0x00000000ffff0000):
                        m26_frame_status[plane_id] |= FRAME_ID_OVERFLOW
                        m26_frame_ids[plane_id] = np.int64(2**32) + m26_frame_ids[plane_id]
                    # Get Mimosa26 frame ID from raw data word (MSB)
                    m26_frame_ids[plane_id] = get_frame_id_high(raw_data_word) | (m26_frame_ids[plane_id] & 0x7fffffff0000ffff)
                elif m26_word_index[plane_id] == 4:  # Mimosa26 frame length
                    m26_frame_length[plane_id] = get_frame_length(raw_data_word)
                    if m26_frame_length[plane_id] > 570:  # Defined in the Mimosa26 protocol, no more than 570 "useful" data words
                        m26_data_loss[plane_id] = True
                        continue
                elif m26_word_index[plane_id] == 5:  # Mimosa26 frame length, a second time
                    if m26_frame_length[plane_id] != get_frame_length(raw_data_word):  # DO0 & DO1 should always have the same data length
                        m26_data_loss[plane_id] = True
                        continue
                    else:
                        m26_frame_length[plane_id] += get_frame_length(raw_data_word)
                elif m26_word_index[plane_id] == 5 + m26_frame_length[plane_id] + 1:  # Frame trailer0
                    if not is_frame_trailer0(raw_data_word):
                        m26_data_loss[plane_id] = True
                        continue
                elif m26_word_index[plane_id] == 5 + m26_frame_length[plane_id] + 2:  # Frame trailer1
                    if not is_frame_trailer1(raw_data_word, plane=plane_id + 1):
                        m26_data_loss[plane_id] = True
                        continue
                    else:
                        last_completed_m26_frame_ids[plane_id] = m26_frame_ids[plane_id]
                elif m26_word_index[plane_id] > 5 + m26_frame_length[plane_id] + 2:  # Ignore any occurrence of additional raw data words
                    m26_data_loss[plane_id] = True
                    continue
                else:  # Column / Row words (actual data word with hits)
                    if m26_n_words[plane_id] == 0:  # First word contains the row info and the number of data words for this row
                        if m26_word_index[plane_id] == 5 + m26_frame_length[plane_id]:  # Always even amount of words or this fill word is used
                            # Ignore this fill word
                            continue
                        else:
                            m26_n_words[plane_id] = get_n_words(raw_data_word)
                            m26_rows[plane_id] = get_row(raw_data_word)  # Get row from data word
                            if m26_rows[plane_id] >= 576:  # Row overflow
                                m26_data_loss[plane_id] = True
                                continue
                        if has_overflow(raw_data_word):
                            m26_frame_status[plane_id] |= OVERFLOW_FLAG  # set overflow bit
                        else:
                            m26_frame_status[plane_id] & ~OVERFLOW_FLAG  # unset overflow bit
                    else:
                        m26_n_words[plane_id] = m26_n_words[plane_id] - 1  # Count down the words
                        n_hits = get_n_hits(raw_data_word)
                        column = get_column(raw_data_word)  # Get column from data word
                        if column >= 1152:  # Column overflow
                            m26_data_loss[plane_id] = True
                            continue
                        for k in range(n_hits + 1):
                            if column + k >= 1152:
                                m26_data_loss[plane_id] = True
                                break
                            # Fill occupancy hist
                            occ_hist[column + k, m26_rows[plane_id], plane_id] += 1
        elif is_trigger_word(raw_data_word):  # Raw data word is TLU/trigger word
            pass  # Not needed here
        else:  # Raw data contains unknown word, neither M26 nor TLU word
            for tmp_plane_index in range(6):
                m26_data_loss[tmp_plane_index] = True

    return m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids


class OccupancyHistogramming(object):
    ''' Fast histogramming of raw data to a 2D hit histogramm

        No event building and seperate process for speed up
    '''
    _queue_timeout = 0.01  # max blocking time to delete object [s]
    _raw_data_queue = multiprocessing.Queue()
    stop = multiprocessing.Event()
    lock = multiprocessing.Lock()

    def __init__(self):
        # Create shared memory 32 bit unsigned int numpy array
        shared_array_base = multiprocessing.Array(ctypes.c_uint, 1152 * 576 * 6)
        shared_array = np.ctypeslib.as_array(shared_array_base.get_obj())
        self.occ_hist = shared_array.reshape(1152, 576, 6)

        self.p = multiprocessing.Process(target=self.worker,
                                         args=(self._raw_data_queue, shared_array_base,
                                               self.lock, self.stop, ))
        self.p.start()

    def add(self, raw_data):
        ''' Add raw data to be histogrammed '''
        self._raw_data_queue.put(raw_data)

    def reset(self, wait=True, timeout=0.5):
        ''' Reset histogram '''
        if not wait:
            if self._raw_data_queue.qsize() != 0:
                logger.warning('Resetting histogram while adding data')
        else:
            n_time = 0
            while self._raw_data_queue.qsize() != 0:
                time.sleep(0.01)
                n_time += 1
                if n_time * 0.01 > timeout:
                    logger.warning('Resetting histogram while adding data')
                    break
        with self.lock:
            # No overwrite with a new zero array due to shared memory
            for col in range(1152):
                for row in range(576):
                    for plane in range(6):
                        self.occ_hist[col, row, plane] = 0

    def get(self, wait=True, timeout=0.5, reset=True):
        ''' Get the result histogram '''
        if not wait:
            if self._raw_data_queue.qsize() != 0:
                logger.warning('Getting histogram while adding data')
        else:
            n_time = 0
            while self._raw_data_queue.qsize() != 0:
                time.sleep(0.01)
                n_time += 1
                if n_time * 0.01 > timeout:
                    logger.warning('Getting histogram while adding data')
                    break
        with self.lock:
            if reset:
                occ_hist = self.occ_hist.copy()
                # No overwrite with a new zero array due to shared memory
                for col in range(1152):
                    for row in range(576):
                        for plane in range(6):
                            self.occ_hist[col, row, plane] = 0
                return occ_hist
            else:
                return self.occ_hist

    def worker(self, raw_data_queue, shared_array_base, lock, stop):
        ''' Histogramming in seperate process '''
        occ_hist = np.ctypeslib.as_array(shared_array_base.get_obj()).reshape(1152, 576, 6)
        # Raw data interpreter
        # Per frame variables
        m26_frame_ids = np.zeros(shape=(6, ), dtype=np.int64)  # The Mimosa26 frame ID of the actual frame
        m26_frame_length = np.zeros(shape=(6, ), dtype=np.uint32)  # The number of "useful" data words for the actual frame
        m26_data_loss = np.ones((6, ), dtype=np.bool_)  # The data loss status for the actual frame
        m26_word_index = np.zeros(shape=(6, ), dtype=np.uint32)  # The word index per device of the actual frame
        m26_timestamps = np.zeros(shape=(6, ), dtype=np.int64)  # The timestamp for each plane (in units of 40 MHz)
        last_m26_timestamps = np.zeros(shape=(6, ), dtype=np.int64)
        m26_n_words = np.zeros(shape=(6, ), dtype=np.uint32)  # The number of words containing column / row info
        m26_rows = np.zeros(shape=(6, ), dtype=np.uint32)  # The actual readout row (rolling shutter)
        m26_frame_status = np.zeros(shape=(6, ), dtype=np.uint32)  # The status flags for the actual frames
        last_completed_m26_frame_ids = -1 * np.ones(shape=6, dtype=np.int64)  # The status if the frame is complete for the actual frame
        while not stop.is_set():
            try:
                raw_data = raw_data_queue.get(timeout=self._queue_timeout)
                with lock:
                    m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids = histogram(
                        raw_data, occ_hist, m26_frame_ids, m26_frame_length, m26_data_loss, m26_word_index, m26_timestamps, last_m26_timestamps, m26_n_words, m26_rows, m26_frame_status, last_completed_m26_frame_ids)
            except queue.Empty:
                continue

    def close(self):
        ''' Close process and wait till done. Likely needed to give access to pytable file handle.'''
        logger.info('Stopping process %d', self.p.pid)
        self._raw_data_queue.close()
        self._raw_data_queue.join_thread()  # Needed otherwise IOError: [Errno 232] The pipe is being closed
        self.stop.set()
        self.p.join()
        del self.p  # explicit delete required to free memory
        self.p = None

    def __del__(self):
        if self.p and self.p.is_alive():
            logger.warning('Process still running. Was close() called?')
            self.close()

    # def __del__(self):
    #     self._raw_data_queue.close()
    #     self._raw_data_queue.join_thread()  # Needed otherwise IOError: [Errno 232] The pipe is being closed
    #     self.stop.set()
    #     self.p.join()


if __name__ == "__main__":
    pass
