# Authors: Adam Li <adam2392@gmail.com>
#
# License: BSD (3-clause)
import os
import os.path as op
from collections import OrderedDict
from datetime import datetime, timezone

import numpy as np

from ..base import BaseRaw
from ..constants import FIFF
from ..meas_info import create_info
from ..utils import _mult_cal_one
from ...annotations import Annotations
from ...utils import logger, verbose, fill_doc, warn


@fill_doc
def read_raw_persyst(fname, preload=False, verbose=None):
    """Reader for a Persyst (.lay/.dat) recording.

    Parameters
    ----------
    fname : str
        Path to the Persyst header (.lay) file.
    %(preload)s
    %(verbose)s

    Returns
    -------
    raw : instance of RawPersyst
        A Raw object containing Persyst data.

    See Also
    --------
    mne.io.Raw : Documentation of attribute and methods.
    """
    return RawPersyst(fname, preload, verbose)


@fill_doc
class RawPersyst(BaseRaw):
    """Raw object from a Persyst file.

    Parameters
    ----------
    fname : str
        Path to the Persyst header (.lay) file.
    %(preload)s
    %(verbose)s

    See Also
    --------
    mne.io.Raw : Documentation of attribute and methods.
    """

    @verbose
    def __init__(self, fname, preload=False, verbose=None):
        logger.info('Loading %s' % fname)

        if not fname.endswith('.lay'):
            fname = fname + '.lay'
        curr_path, lay_fname = op.dirname(fname), op.basename(fname)
        if not op.exists(fname):
            raise FileNotFoundError(f'The path you specified, '
                                    f'"{lay_fname}",does not exist.')

        # sections and subsections currently unused
        keys, data, sections = _read_lay_contents(fname)

        # these are the section headers in the Persyst file layout
        # Note: We do not make use of "SampleTimes" yet
        fileinfo_dict = OrderedDict()
        channelmap_dict = OrderedDict()
        patient_dict = OrderedDict()
        comments_dict = OrderedDict()

        # loop through each line in the lay file
        for key, val, section in zip(keys, data, sections):
            if key == '':
                continue

            # make sure key are lowercase for everything, but electrodes
            if key is not None and section != 'channelmap':
                key = key.lower()

            # FileInfo
            if section == 'fileinfo':
                # extract the .dat file name
                if key == 'file':
                    dat_fname = val
                    dat_path = op.dirname(dat_fname)
                    dat_fpath = op.join(curr_path, op.basename(dat_fname))

                    # determine if .dat file exists where it should
                    error_msg = f'The data path you specified ' \
                                f'does not exist for the lay path, {lay_fname}'
                    if op.isabs(dat_path) and not op.exists(dat_fname):
                        raise FileNotFoundError(error_msg)
                    if not op.exists(dat_fpath):
                        raise FileNotFoundError(error_msg)
                fileinfo_dict[key] = val
            # ChannelMap
            elif section == 'channelmap':
                # channel map has <channel_name>=<number> for <key>=<val>
                channelmap_dict[key] = val
            # Patient (All optional)
            elif section == 'patient':
                patient_dict[key] = val
            elif section == 'comments':
                comments_dict[key] = val

        # get numerical metadata
        # datatype is either 7 for 32 bit, or 0 for 16 bit
        datatype = fileinfo_dict.get('datatype')
        cal = float(fileinfo_dict.get('calibration'))
        n_chs = int(fileinfo_dict.get('waveformcount'))

        # Store subject information from lay file in mne format
        # Note: Persyst also records "Physician", "Technician",
        #       "Medications", "History", and "Comments1" and "Comments2"
        #       and this information is currently discarded
        subject_info = _get_subjectinfo(patient_dict)

        # set measurement date
        testdate = patient_dict.get('testdate')
        if testdate is not None:
            # TODO: Persyst may change its internal date schemas
            #  without notice
            # These are the 3 "so far" possible datatime storage
            # formats in Persyst .lay
            if '/' in testdate:
                testdate = datetime.strptime(testdate, '%m/%d/%Y')
            elif '-' in testdate:
                testdate = datetime.strptime(testdate, '%d-%m-%Y')
            elif '.' in testdate:
                testdate = datetime.strptime(testdate, '%Y.%m.%d')

            if not isinstance(testdate, datetime):
                warn('Cannot read in the measurement date due '
                     'to incompatible format. Please set manually '
                     'for %s ' % lay_fname)
                meas_date = None
            else:
                testtime = datetime.strptime(patient_dict.get('testtime'),
                                             '%H:%M:%S')
                meas_date = datetime(
                    year=testdate.year, month=testdate.month,
                    day=testdate.day, hour=testtime.hour,
                    minute=testtime.minute, second=testtime.second,
                    tzinfo=timezone.utc)

        # Create mne structure
        ch_names = list(channelmap_dict.keys())
        if n_chs != len(ch_names):
            raise RuntimeError('Channels in lay file do not '
                               'match the number of channels '
                               'in the .dat file.')  # noqa
        # get rid of the "-Ref" in channel names
        ch_names = [ch.upper().split('-REF')[0] for ch in ch_names]

        # get the sampling rate and default channel types to EEG
        sfreq = fileinfo_dict.get('samplingrate')
        ch_types = 'eeg'
        info = create_info(ch_names, sfreq, ch_types=ch_types)
        info.update(subject_info=subject_info)
        for idx in range(n_chs):
            # calibration brings to uV then 1e-6 brings to V
            info['chs'][idx]['cal'] = cal * 1.0e-6
        info['meas_date'] = meas_date

        # determine number of samples in file
        # Note: We do not use the lay file to do this
        # because clips in time may be generated by Persyst that
        # DO NOT modify the "SampleTimes" section
        with open(dat_fpath, 'rb') as f:
            # determine the precision
            if int(datatype) == 7:
                # 32 bit
                dtype = np.dtype('i4')
            elif int(datatype) == 0:
                # 16 bit
                dtype = np.dtype('i2')
            else:
                raise RuntimeError(f'Unknown format: {datatype}')

            # allow offset to occur
            f.seek(0, os.SEEK_END)
            n_samples = f.tell()
            n_samples = n_samples // (dtype.itemsize * n_chs)

            logger.debug(f'Loaded {n_samples} samples '
                         f'for {n_chs} channels.')

        raw_extras = {
            'dtype': dtype,
            'n_chs': n_chs,
            'n_samples': n_samples
        }
        # create Raw object
        super(RawPersyst, self).__init__(
            info, preload, filenames=[dat_fpath],
            last_samps=[n_samples - 1],
            raw_extras=[raw_extras], verbose=verbose)

        # set annotations based on the comments read in
        num_comments = len(comments_dict)
        onset = np.zeros(num_comments, float)
        duration = np.zeros(num_comments, float)
        description = [''] * num_comments
        for t_idx, (_description, (_onset, _duration)) in \
                enumerate(comments_dict.items()):
            # extract the onset, duration, description to
            # create an Annotations object
            onset[t_idx] = _onset
            duration[t_idx] = _duration
            description[t_idx] = _description
        annot = Annotations(onset, duration, description)
        self.set_annotations(annot)

    def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
        """Read a segment of data from a file.

        The Persyst software  records raw data in either 16 or 32 bit
        binary files. In addition, it stores the calibration to convert
        data to uV in the lay file.
        """
        dtype = self._raw_extras[fi]['dtype']
        n_chs = self._raw_extras[fi]['n_chs']
        dat_fname = self._filenames[fi]

        # compute samples count based on start and stop
        time_length_samps = stop - start

        # read data from .dat file into array of correct size, then calibrate
        # records = recnum rows x inf columns
        count = time_length_samps * n_chs

        # seek the dat file
        with open(dat_fname, 'rb') as dat_file_ID:
            # allow offset to occur
            dat_file_ID.seek(n_chs * dtype.itemsize * start, 1)

            # read in the actual record starting at possibly offset
            record = np.fromfile(dat_file_ID, dtype=dtype,
                                 count=count)

        # chs * rows
        # cast as float32; more than enough precision
        record = np.reshape(record, (n_chs, -1), 'F').astype(np.float32)

        # calibrate to convert to V and handle mult
        _mult_cal_one(data, record, idx, cals, mult)


def _get_subjectinfo(patient_dict):
    # attempt to parse out the birthdate, but if it doesn't
    # meet spec, then it will set to None
    birthdate = patient_dict.get('birthdate')
    if '/' in birthdate:
        try:
            birthdate = datetime.strptime(birthdate, '%m/%d/%y')
        except ValueError:
            birthdate = None
            print('Unable to process birthdate of %s ' % birthdate)
    elif '-' in birthdate:
        try:
            birthdate = datetime.strptime(birthdate, '%d-%m-%y')
        except ValueError:
            birthdate = None
            print('Unable to process birthdate of %s ' % birthdate)

    subject_info = {
        'first_name': patient_dict.get('first'),
        'middle_name': patient_dict.get('middle'),
        'last_name': patient_dict.get('last'),
        'sex': patient_dict.get('sex'),
        'hand': patient_dict.get('hand'),
        'his_id': patient_dict.get('id'),
        'birthday': birthdate,
    }

    # Recode sex values
    sex_dict = dict(
        m=FIFF.FIFFV_SUBJ_SEX_MALE,
        male=FIFF.FIFFV_SUBJ_SEX_MALE,
        f=FIFF.FIFFV_SUBJ_SEX_FEMALE,
        female=FIFF.FIFFV_SUBJ_SEX_FEMALE,
    )
    subject_info['sex'] = sex_dict.get(subject_info['sex'],
                                       FIFF.FIFFV_SUBJ_SEX_UNKNOWN)

    # Recode hand values
    hand_dict = dict(
        r=FIFF.FIFFV_SUBJ_HAND_RIGHT,
        right=FIFF.FIFFV_SUBJ_HAND_RIGHT,
        l=FIFF.FIFFV_SUBJ_HAND_LEFT,
        left=FIFF.FIFFV_SUBJ_HAND_LEFT,
        a=FIFF.FIFFV_SUBJ_HAND_AMBI,
        ambidextrous=FIFF.FIFFV_SUBJ_HAND_AMBI,
        ambi=FIFF.FIFFV_SUBJ_HAND_AMBI,
    )
    # no handedness is set when unknown
    try:
        subject_info['hand'] = hand_dict[subject_info['hand']]
    except KeyError:
        subject_info.pop('hand')

    return subject_info


def _read_lay_contents(fname):
    """Lay file are laid out like a INI file."""
    # keep track of sections, keys and data
    sections = []
    keys, data = [], []

    # initialize all section to empty str
    section = ''
    with open(fname, 'r') as fin:
        for line in fin:
            # break a line into a status, key and value
            status, key, val = _process_lay_line(line, section)

            # handle keys and values if they are
            # Section, Subsections, or Line items
            if status == 1:  # Section was found
                section = val.lower()
                continue

            # keep track of all sections, subsections,
            # keys and the data of the file
            sections.append(section)
            data.append(val)
            keys.append(key)

    return keys, data, sections


def _process_lay_line(line, section):
    """Process a line read from the Lay (INI) file.

    Each line in the .lay file will be processed
    into a structured ``status``, ``key`` and ``value``.

    Parameters
    ----------
    line : str
        The actual line in the Lay file.
    section : str
        The section in the Lay file.

    Returns
    -------
    status : int
        Returns the following integers based on status.
        -1  => unknown string found
        0   => empty line found
        1   => section found
        2   => key-value pair found
    key : str
        The string before the ``'='`` character. If section is "Comments",
        then returns the text comment description.
    value : str
        The string from the line after the ``'='`` character. If section is
        "Comments", then returns the onset and duration as a tuple.
    """
    key = ''  # default; only return value possibly not set
    line = line.strip()  # remove leading and trailing spaces
    end_idx = len(line) - 1  # get the last index of the line

    # empty sequence evaluates to false
    if not line:
        status = 0
        key = ''
        value = ''
        return status, key, value
    # section found
    elif (line[0] == '[') and (line[end_idx] == ']') \
            and (end_idx + 1 >= 3):
        status = 1
        value = line[1:end_idx].lower()
    # key found
    else:
        # handle Comments section differently from all other sections
        # TODO: utilize state and var_type in code.
        #  Currently not used
        if section == 'comments':
            # Persyst Comments output 5 variables "," separated
            time_sec, duration, state, var_type, text = line.split(',')
            status = 2
            key = text
            value = (time_sec, duration)
        # all other sections
        else:
            if '=' not in line:
                raise RuntimeError('The line %s does not conform '
                                   'to the standards. Please check the '
                                   '.lay file.' % line)  # noqa
            pos = line.index('=')
            status = 2

            # the line now is composed of a
            # <key>=<value>
            key = line[0:pos]
            key.strip()
            value = line[pos + 1:end_idx + 1]
            value.strip()
    return status, key, value
