import logging
import os
from struct import Struct
import astropy.units as u
import numpy as np
from peegy.definitions.channel_definitions import ChannelItem
from peegy.definitions.events import SingleEvent, Events
from peegy.tools.units.unit_tools import set_default_unit
from typing import List
log = logging.getLogger()
__author__ = 'jundurraga'
"These scripts are based on MNE-python and Neuroscan data structure http://paulbourke.net/dataformats/eeg/"


class CNTEventType1(object):
    def __init__(self,
                 stim_type,
                 key_board,
                 key_pad_accept,
                 offset):

        self.stim_type = stim_type
        self.key_board = key_board
        self.key_pad_accept = key_pad_accept
        self.offset = offset


class CNTEventType2(object):
    def __init__(self,
                 stim_type,
                 key_board,
                 key_pad_accept,
                 offset,
                 type,
                 code,
                 latency,
                 epoch_event,
                 accept2,
                 accuracy
                 ):
        self.stim_type = stim_type
        self.key_board = key_board
        self.key_pad_accept = key_pad_accept
        self.offset = offset
        self.type = type
        self.code = code
        self.latency = latency
        self.epoch_event = epoch_event
        self.accept2 = accept2
        self.accuracy = accuracy


class CNTEventType3(object):
    def __init__(self,
                 stim_type,
                 key_board,
                 key_pad_accept,
                 offset,
                 type,
                 code,
                 latency,
                 epoch_event,
                 accept2,
                 accuracy
                 ):
        self.stim_type = stim_type
        self.key_board = key_board
        self.key_pad_accept = key_pad_accept
        self.offset = offset
        self.type = type
        self.code = code
        self.latency = latency
        self.epoch_event = epoch_event
        self.accept2 = accept2
        self.accuracy = accuracy


class DeviceEventChannel:
    bdf_event_channel = 'Status'
    edf_event_channel = 'EDF Annotations'
    bdf_event_annotations = 'BDF Annotations'
    auto = 'auto'


def read_cnt_header(file_name: str):
    """Read the cnt header.
     http://paulbourke.net/dataformats/eeg/
    """
    data_offset = 900  # Size of the 'SETUP' header.
    with open(file_name, "rb", buffering=0) as f:
        f.seek(21)
        _subject_id = f.read(20).strip().decode('ascii', errors='ignore').replace('\x00', '')
        f.seek(121)
        subject_name = f.read(20).strip().decode('ascii', errors='ignore').replace('\x00', '').split()
        _last_name = subject_name[0] if len(subject_name) > 0 else ''
        _first_name = subject_name[-1] if len(subject_name) > 0 else ''
        f.seek(2, 1)
        _sex = f.read(1).strip().decode('ascii', errors='ignore').replace('\x00', '')
        _handedness = f.read(1).strip().decode('ascii', errors='ignore').replace('\x00', '')
        f.seek(205)
        _session_label = f.read(20).strip().decode('ascii', errors='ignore').replace('\x00', '')
        _session_date = f.read(10).strip().decode('ascii', errors='ignore').replace('\x00', '')
        _session_time = f.read(12).strip().decode('ascii', errors='ignore').replace('\x00', '')
        f.seek(353)
        n_events = np.fromfile(f, dtype="<u4", count=1).item()
        f.seek(370)
        n_channels = np.fromfile(f, dtype="<u2", count=1).item()
        f.seek(376)
        sampling_rate = np.fromfile(f, dtype="<u2", count=1).item() * u.Hz
        f.seek(438)
        lowpass_toggle = np.fromfile(f, "i1", count=1).item()
        highpass_toggle = np.fromfile(f, "i1", count=1).item()

        f.seek(864)
        n_samples = np.fromfile(f, dtype="<u4", count=1).item()
        _event_table_position = 886
        f.seek(_event_table_position)
        (event_offset,) = np.frombuffer(f.read(4), dtype="<i4")

        f.seek(869)
        low_cutoff = np.fromfile(f, dtype="f4", count=1).item()
        f.seek(2, 1)
        high_cutoff = np.fromfile(f, dtype="f4", count=1).item()
        n_bytes = 4

        f.seek(890)
        continuous_seconds = np.fromfile(f, dtype="<f4", count=1).item()
        events_size = os.path.getsize(file_name) - event_offset
        if events_size <= 0:  # no events
            data_size = n_samples * n_channels
        else:
            data_size = os.path.getsize(file_name) - (data_offset + 75 * n_channels) - events_size

        if n_samples == 0 or data_size // (n_samples * n_channels) not in [2, 4]:
            n_bytes = 2
            n_samples = data_size // (n_bytes * n_channels)
            # annotations = _read_annotations_cnt(input_fname, data_format="int16")
            # if len(annotations) and annotations.onset[-1] * sampling_rate > n_samples:
            #     n_bytes = 4
            #     n_samples = n_samples_header
        else:
            n_bytes = data_size // (n_samples * n_channels)

        # channel offset (blocks per channel)
        _channel_offset = np.fromfile(f, dtype="<i4", count=1).item()
        if _channel_offset > 1:
            _channel_offset = _channel_offset // n_bytes
        else:
            _channel_offset = 1

        if lowpass_toggle != 1:
            low_cutoff = None
        if highpass_toggle != 1:
            high_cutoff = None

        header = {'file_name': file_name,
                  'subject_id': _subject_id,
                  'first_name': _first_name,
                  'last_name': _last_name,
                  'sex': _sex,
                  'handedness': _handedness,
                  'session_label': _session_label,
                  'session_date': _session_date,
                  'n_channels': n_channels,
                  'lowpass_toggle': lowpass_toggle,
                  'highpass_toggle': highpass_toggle,
                  'low_cutoff': low_cutoff,
                  'high_cutoff': high_cutoff,
                  'file_size': os.path.getsize(file_name),
                  'start_date': _session_date,
                  'start_time': _session_time,
                  'bytes_in_header': data_offset,
                  'fs': sampling_rate,
                  'channel_offset': _channel_offset,
                  'n_bytes': n_bytes,
                  'n_samples': n_samples,
                  'has_events': events_size > 0,
                  'event_offset': event_offset,
                  'n_events': n_events,
                  'duration': n_samples / sampling_rate,
                  'continuous_seconds': continuous_seconds}
        (cals,
         baselines,
         sensitivities,
         channels,
         pos) = (list(), list(), list(), list(), list())

        ch = []
        for ch_idx in range(n_channels):
            f.seek(data_offset + 75 * ch_idx)
            ch_name = f.read(10).strip().decode('ascii', errors='ignore').replace('\x00', '')
            ch.append(ch_name)

            f.seek(data_offset + 75 * ch_idx + 19)
            xy = np.fromfile(f, dtype="f4", count=2)
            xy[1] *= -1  # invert y-axis
            pos.append(xy)
            f.seek(data_offset + 75 * ch_idx + 47)
            # baselines corrections
            baselines.append(np.fromfile(f, dtype="i2", count=1).item())

            f.seek(data_offset + 75 * ch_idx + 59)
            sensitivity = np.fromfile(f, dtype="f4", count=1).item()
            sensitivities.append(sensitivity)

            f.seek(data_offset + 75 * ch_idx + 71)
            cal = np.fromfile(f, dtype="f4", count=1).item()
            cals.append(cal * sensitivity / 204.8)  # calibration factor (micro Volts)
        for ch_name, (_pos_x, _pos_y) in zip(ch, pos):
            channels.append(ChannelItem(label=ch_name,
                                        idx=ch_idx,
                                        x=_pos_x,
                                        y=_pos_y,
                                        ))
        header['channels'] = np.array(channels)
        header['pos'] = np.array(pos)
        header['baselines'] = np.array(baselines)
        header['sensitivity'] = np.array(sensitivities)
        header['cals'] = np.array(cals)

    return header


def _get_event_parser(event_type):
    if event_type == 1:
        event_maker = CNTEventType1
        struct_pattern = "<HBcl"
    elif event_type == 2:
        event_maker = CNTEventType2
        struct_pattern = "<HBclhhfccc"
    elif event_type == 3:
        event_maker = CNTEventType3
        struct_pattern = "<HBclhhfccc"  # Same as event type 2
    else:
        raise ValueError(f"unknown CNT even type {event_type}")

    def parser(buffer):
        struct = Struct(struct_pattern)
        for chunk in struct.iter_unpack(buffer):
            yield event_maker(*chunk)

    return parser


def get_data(header,
             channels_idx: np.array = np.array([]),
             ini_time: u.quantity.Quantity = 0 * u.s,
             end_time: u.quantity.Quantity | None = None):
    fs_data = set_default_unit(header['fs'], u.Hz)
    ini_time = set_default_unit(ini_time, u.s)
    end_time = set_default_unit(end_time, u.s)

    n_channels = header['n_channels']
    channel_offset = header['channel_offset']
    baselines = header['baselines']
    cals = header['cals']
    n_bytes = header["n_bytes"]
    n_samples = header['n_samples']
    dtype = "<i4" if n_bytes == 4 else "<i2"
    chunk_size = channel_offset * n_channels
    if not channels_idx.size:
        channels_idx = np.arange(0, header['channels'].size).astype(int)
    # The data is divided into blocks of samples / channel.
    # channel_offset determines the amount of successive samples.
    # Offset is used to align the data because start can be in
    # the middle of these blocks.
    ini_sample = np.round(ini_time * fs_data).astype(int)
    if end_time is None:
        end_time = np.floor(header['duration'])
    else:
        end_time = np.minimum(end_time, header['duration'])
    end_sample = np.minimum(np.round(end_time * fs_data).astype(int), n_samples)
    data_left = (end_sample - ini_sample) * n_channels
    # Read up to 100 MB of data at a time, block_size is in data samples
    block_size = ((int(100e6) // n_bytes) // chunk_size) * chunk_size
    block_size = min(data_left, block_size)
    s_offset = ini_sample % channel_offset
    data = np.zeros((end_sample - ini_sample, n_channels))
    with open(header['file_name'], "rb", buffering=0) as f:
        f.seek(900 + n_channels * (75 + (ini_sample - s_offset) * n_bytes))
        for sample_start in np.arange(0, data_left, block_size) // n_channels:
            # Earlier comment says n_samples is unreliable, but I think it
            # is because it needed to be changed to unsigned int
            # See: PR #12393
            sample_stop = sample_start + min(
                (
                    n_samples,
                    block_size // n_channels,
                    data_left // n_channels - sample_start,
                )
            )
            n_samps = sample_stop - sample_start
            # In case channel offset and start time do not align perfectly,
            # extra sample sets are read here to cover the desired time
            # window. The whole (up to 100 MB) block is read at once and
            # then reshaped to (n_channels, n_samples).
            extra_samps = (
                chunk_size
                if (s_offset != 0 or n_samps % channel_offset != 0)
                else 0
            )
            if s_offset >= (channel_offset / 2):  # Extend at the end.
                extra_samps += chunk_size
            count = n_samps // channel_offset * chunk_size + extra_samps
            n_chunks = count // chunk_size
            samps = np.fromfile(f, dtype=dtype, count=count)
            samps = samps.reshape((n_chunks, n_channels, channel_offset), order="C")

            # Intermediate shaping to chunk sizes.
            block = np.zeros((channel_offset * n_chunks, n_channels))
            for set_idx, row in enumerate(samps):  # Final shape.
                block_slice = slice(
                    set_idx * channel_offset, (set_idx + 1) * channel_offset
                )
                block[block_slice, :n_channels] = row.T
            data[sample_start: sample_stop, channels_idx] = block[s_offset: n_samps + s_offset, channels_idx]
            print('reading {:}'.format((sample_stop/fs_data).to(u.s)))
    data = data[:, channels_idx]
    data -= baselines[channels_idx][None, :]
    data *= cals[channels_idx][None, :]
    data = data * u.uV
    # get events
    events = get_events(header=header,
                        ini_time=ini_time,
                        end_time=end_time)
    return data, events, channels_idx


def get_events(header,
               ini_time: u.quantity.Quantity = 0 * u.s,
               end_time: u.quantity.Quantity | None = None):
    # http://paulbourke.net/dataformats/eeg/
    events = None
    fs_data = set_default_unit(header['fs'], u.Hz)
    ini_time = set_default_unit(ini_time, u.s)
    end_time = set_default_unit(end_time, u.s)
    n_samples = header['n_samples']
    ini_sample = np.round(ini_time * fs_data).astype(int)
    if end_time is None:
        end_time = np.floor(header['duration'])
    else:
        end_time = np.minimum(end_time, header['duration'])
    end_sample = np.minimum(np.round(end_time * fs_data).astype(int), n_samples)
    if header['has_events']:
        teeg_parser = Struct("<Bll")
        with open(header['file_name'], "rb", buffering=0) as f:
            f.seek(header['event_offset'])
            (event_type, total_length, offset) = teeg_parser.unpack(f.read(teeg_parser.size))

        event_parser = _get_event_parser(event_type=event_type)

        with open(header['file_name'], "rb", buffering=0) as f:
            f.seek(header['event_offset'] + 9)  # the real table stats at +9
            buffer = f.read(total_length)

        orig_events = list(event_parser(buffer))
        events = _event_times(header=header,
                              original_events=orig_events,
                              ini_sample=ini_sample,
                              end_sample=end_sample)
    return events


def _event_times(header: dict | None = None,
                 original_events: List[object] | None = None,
                 ini_sample: int = 0,
                 end_sample: int | None = None) -> list:
    n_bytes = header['n_bytes']
    n_channels = header['n_channels']
    fs = header['fs']
    if end_sample is None:
        end_sample = np.inf
    events = []
    for _ev in original_events:
        offset = _ev.offset
        if isinstance(_ev, CNTEventType3):
            offset = offset * n_bytes * n_channels
        event_time = offset - 900 - (75 * n_channels)
        event_time = event_time // (n_channels * n_bytes)
        event_time = event_time - 1
        bad = False
        if "xd0" in str(_ev.key_pad_accept):
            bad = False
        if "xc0" in str(_ev.key_pad_accept):
            bad = True
        # Extract the 4-bit fields
        # Upper nibble (4 bits) currently not used
        # accept = (event.KeyPad_Accept[0] & 0xF0) >> 4
        # Lower nibble (4 bits) keypad button press
        keypad = _ev.key_pad_accept[0] & 0x0F
        description = None
        if str(keypad) != "0":
            description = f"KeyPad Response {keypad}"
        elif _ev.key_board != 0:
            description = f"Keyboard Response {_ev.key_board}"
        _event_time = (event_time - ini_sample) / fs
        _original_time = event_time / fs
        if (_event_time >= 0) and (_original_time <= end_sample / fs):
            events.append(SingleEvent(time_pos=_event_time,
                                      code=_ev.stim_type,
                                      bad_event=bad,
                                      description=description))
        if _original_time > end_sample / fs:
            break
    return Events(events)

    #     # Check to see if there are any button presses
    #     description = []
    #     for event in my_events:

    #
    #     description = np.array(description)
    #
    #     onset, duration, description = _update_bad_span_onset(
    #         accept_reject, onset / sfreq, duration, description
    #     )
    #     return Annotations(
    #         onset=onset, duration=duration, description=description, orig_time=None
    #     )
