# -*- coding: utf-8 -*-
# -*- mode: python -*-
""" Utilities for extracellular experiments """
import os
import re
import json
import logging

import quickspikes as qs
import nbank
import numpy as np
import h5py as h5
from dlab import pprox
from dlab.util import memodict

log = logging.getLogger("dlab")


### general


def entry_time(entry):
    from arf import timestamp_to_float

    return timestamp_to_float(entry.attrs["timestamp"])


def iter_entries(data_file):
    """Iterate through the entries in an arf file in order of time"""
    return enumerate(sorted(data_file.values(), key=entry_time))


### present-audio


def audiolog_to_trials(trials, data_file, sync_dset="channel37", sync_thresh=1.0):
    """Parses a logfile from Margot's present_audio scripts for experiment structure

    trials: the "presentation" field in the json output of present_audio.py
    data_file: open handle to the hdf5 file generated by open-ephys during the recording
    sync_dset: the name of the dataset containing the synchronization signal
    sync_thresh: the threshold for detecting the sync signal
    trials: number of trials per stimulus
    """
    from arf import timestamp_to_float

    # Each element in this structure corresponds to a trial. In some cases the
    # data are stored as a dict/map, but the keys are just strings of the trial
    # number. The indices correspond to the entries in the arf file.
    n_trials = len(trials)
    expt_start = None
    sample_count = 0
    det = qs.detector(sync_thresh, 10)
    for i in range(n_trials):
        pproc = {"events": [], "index": i}
        pproc.update(trials[str(i)])
        entry_name = "/rec_%d" % i
        entry = data_file[entry_name]

        # get time relative to first trial
        time = entry_time(entry)
        if expt_start is None:
            expt_start = time
        pproc["offset"] = time - expt_start
        # find the sync signal - we expect one and only one click
        dset = entry[sync_dset]
        pproc["recording"] = {
            "entry": entry_name,
            "start": int(sample_count),
            "stop": int(sample_count + dset.size),
            "sampling_rate": dset.attrs["sampling_rate"],
        }
        sample_count += dset.size
        sync = dset[:].astype("d")
        det.scale_thresh(sync.mean(), sync.std())
        clicks = det(sync)
        if len(clicks) != 1:
            log.error("%s: expected 1 click, detected %d", dset.name, len(clicks))
        else:
            pproc["stim_on"] = clicks[0] / dset.attrs["sampling_rate"]
        yield pproc


def audiolog_to_pprox_script(argv=None):
    """CLI to generate a pprox from present_audio log"""
    import sys
    import argparse
    import json
    from dlab.util import setup_log, json_serializable

    __version__ = "0.1.0"

    p = argparse.ArgumentParser(
        description="generate pprox from trial structure in present_audio logfile"
    )
    p.add_argument(
        "-v", "--version", action="version", version="%(prog)s " + __version__
    )
    p.add_argument("--debug", help="show verbose log messages", action="store_true")
    p.add_argument(
        "--output",
        "-o",
        type=argparse.FileType("w", encoding="utf-8"),
        default=sys.stdout,
        help="name of output file. If absent, outputs to stdout",
    )
    p.add_argument(
        "--sync",
        default="channel37",
        help="name of channel with synchronization signal",
    )
    p.add_argument(
        "--sync-thresh",
        default="30.0",
        type=float,
        help="threshold (z-score) for detecting synchronization clicks",
    )
    p.add_argument("logfile", help="log file generated by present_audio.py")
    p.add_argument("recording", help="neurobank id or URL for the ARF recording")
    args = p.parse_args(argv)
    setup_log(log, args.debug)

    resource_url = nbank.full_url(args.recording)
    datafile = nbank.get(args.recording, local_only=True)
    if datafile is None:
        p.error(
            "unable to locate resource %s - is it deposited in neurobank?"
            % args.recording
        )
    log.info("loading recording from resource %s", resource_url)

    with h5.File(datafile, "r") as afp:
        with open(args.logfile, "rt") as lfp:
            expt_log = json.load(lfp)
            trials = pprox.from_trials(
                audiolog_to_trials(
                    expt_log.pop("presentation"), afp, args.sync, args.sync_thresh
                ),
                recording=resource_url,
                processed_by=["{} {}".format(p.prog, __version__)],
                **expt_log
            )
    json.dump(trials, args.output, default=json_serializable)
    if args.output != sys.stdout:
        log.info("wrote trial data to '%s'", args.output.name)


############### open-ephys-audio specific functions

def find_stim_dset(entry):
    """Returns the first dataset that matches 'Network_Events.*_TEXT'"""
    rex = re.compile(r"Network_Events-.*?TEXT")
    for name in entry:
        if rex.match(name) is not None:
            log.debug("  - stim log dataset: %s", name)
            return entry[name]


def parse_stim_id(path):
    """Extracts the stimulus id from the path"""
    return os.path.splitext(os.path.basename(path))[0]


@memodict
def stim_duration(stim_name):
    """
    Returns the duration of a stimulus (in s). This can only really be done by
    downloading the stimulus from the registry, because the start/stop times are
    not reliable. We try to speed this up by memoizing the function and caching
    the downloaded files.

    """
    import wave
    from dlab.core import fetch_resource

    neurobank_registry = nbank.default_registry()
    target = fetch_resource(neurobank_registry, stim_name)
    with open(target, "rb") as fp:
        reader = wave.open(fp)
        return 1.0 * reader.getnframes() / reader.getframerate()


def entry_metadata(entry):
    """Extracts metadata from an entry in an oeaudio-present experiment ARF file.

    Metadata are passed to open-ephys through the network events socket as a
    json-encoded dictionary. There should be one and only one metadata message
    per entry, so only the first is returned.

    """
    re_metadata = re.compile(r"metadata: (\{.*\})")
    stims = find_stim_dset(entry)
    for row in stims:
        message = row["message"].decode("utf-8")
        m = re_metadata.match(message)
        try:
            metadata = json.loads(m.group(1))
        except (AttributeError, json.JSONDecodeError):
            pass
        else:
            metadata.update(name=entry.name, sampling_rate=stims.attrs["sampling_rate"])
            return metadata
