# -*- mode: python -*-
""" Front-end scripts for song detector """
import argparse
import logging
import os
import pickle
from pathlib import Path

import arf
import h5py as h5
import numpy as np
from arf import DataTypes, create_dataset, set_uuid, timestamp_to_datetime
from dateutil import parser as date_parser

from quicksong import __version__
from quicksong.core import IntervalFinder

script_name = "quicksong"

_example_config = dict(
    spectrogram={
        "window_ms": 20.0,
        "shift_ms": 10.0,
        "frequency_range": [0.7, 20.0],
    },
    features=["total_power", "wiener_entropy"],
    target_label="x",
    smoothing={"kernel_widths": [11, 21, 51]},
    classifier={"gamma": 2, "C": 1},
    testing={"test_size": 0.4, "random_state": 45},
    intervals={"max_gap_ms": 100, "min_interval_ms": 500},
    datasets=[],
)

log = logging.getLogger("")


def valid_datetime(date_string):
    """
    Convert string to datetime object or raise ArgumentTypeError
    """
    try:
        return date_parser.parse(date_string)
    except ValueError as err:
        msg = f"'{date_string}' is not a valid date/time"
        raise argparse.ArgumentTypeError(msg) from err


def setup_log(log, debug=False):
    """Set up logging for a module.

    log: generate by calling e.g. `log = logging.getLogger("dlab.extracellular")`
    """
    ch = logging.StreamHandler()
    formatter = logging.Formatter("%(message)s")
    loglevel = logging.DEBUG if debug else logging.INFO
    log.setLevel(loglevel)
    ch.setLevel(loglevel)
    ch.setFormatter(formatter)
    log.addHandler(ch)


def all_same(seq):
    """If all elements of seq are equal, returns the value, otherwise None"""
    it = iter(seq)
    first = next(it)
    for e in it:
        if e != first:
            return None
    return first


def make_config(args):
    import yaml

    cfg = _example_config.copy()
    for lblfile in args.label_files:
        cfg["datasets"].append(lblfile.stem)
    with open(args.config, "w") as fp:
        yaml.dump(cfg, fp)
        log.info("created new configuration file '%s'", args.config)
    if args.train:
        args.dataset_dir = lblfile.parent
        args.model = None
        train_classifier(args)


def make_extractor(cfg, sampling_rate_khz):
    from quicksong import features
    from quicksong.core import FeatureExtractor, make_hanning_kernels
    from quicksong.streaming import STFT, OverlapSaveConvolver

    speccr = STFT(
        sampling_rate_khz,
        cfg["spectrogram"]["window_ms"],
        cfg["spectrogram"]["shift_ms"],
        cfg["spectrogram"]["frequency_range"],
    )
    feats = tuple(getattr(features, f) for f in cfg["features"])
    conv = OverlapSaveConvolver(make_hanning_kernels(cfg["smoothing"]["kernel_widths"]))
    return FeatureExtractor(speccr, feats, conv)


def train_classifier(args):
    import ewave
    import yaml
    from arfx import lblio

    from quicksong import core, model

    with open(args.config) as fp:
        cfg = yaml.safe_load(fp)
    log.debug(cfg)
    label_ids = {cfg["target_label"]: 0}
    log.info("Training song classifier:")
    log.info("- version: %s", __version__)

    X_all = []
    y_all = []
    for dset in cfg["datasets"]:
        log.info("- dataset: %s", dset)
        path = args.dataset_dir / dset
        with ewave.open(path.with_suffix(".wav"), "r") as fp:
            Fs = fp.sampling_rate / 1000.0
            signal = ewave.rescale(fp.read(), "f")
            log.info(" - recording: %d samples @ %.2f kHz", signal.size, Fs)
        extractor = make_extractor(cfg, Fs)
        feats = extractor.process_all(signal)
        log.info("   - calculated features %s", feats.shape)

        with open(path.with_suffix(".lbl")) as fp:
            out = []
            for lbl in lblio.read(fp):
                n = lbl["name"]
                try:
                    id = label_ids[n]
                except KeyError:
                    id = len(label_ids)
                    label_ids[n] = id
                out.append((id, lbl["start"] * 1000, lbl["stop"] * 1000))
            labels = np.array(
                out, dtype=[("id", "i4"), ("start", "f8"), ("stop", "f8")]
            )
            log.info(" - labels: %d segments", labels.size)
            X, y = core.make_training_dataset(
                feats, labels, cfg["spectrogram"]["shift_ms"]
            )
            log.info("  - labeled frames: %d", y.size)
            X_all.append(X)
            y_all.append(y)
    X = np.vstack(X_all)
    y = np.concatenate(y_all)
    assert X.shape[0] == y.size
    log.info("- training classifier...")
    classifier, score = model.train_classifier(X, y, cfg["classifier"], cfg["testing"])
    log.info("  - model score: %.3f", score)
    if args.model is None:
        args.model = args.config.with_suffix(".pkl")
    with open(args.model, "wb") as fp:
        pickle.dump(
            {"config": cfg, "classifier": classifier, "version": __version__}, fp
        )
        log.info("- wrote model to '%s'", args.model)


class ArfWriter:
    """Utility class that writes to sequentially numbered entries"""

    def __init__(self, arfp, template="song_{index:04}"):
        self.arfp = arfp
        self.entry_count = len(arfp.keys())
        self.template = template

    @property
    def filename(self):
        return self.arfp.filename

    def create_entry(self, timestamp, **attributes):
        import arf

        next_entry_name = self.template.format(index=self.entry_count)
        self.entry_count += 1
        return arf.create_entry(self.arfp, next_entry_name, timestamp, **attributes)


def extract_songs(args):
    os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
    log.info("Extracting songs:")
    log.info("- version: %s", __version__)
    with open(args.model, "rb") as fp:
        log.info("- loading classifier from '%s'", args.model)
        model = pickle.load(fp)
    # TODO check model version
    classifier = model["classifier"]
    dt = model["config"]["spectrogram"]["shift_ms"]
    finder = IntervalFinder(max_gap=args.max_gap, min_duration=args.min_duration, dt=dt)
    # TODO: more flexibility with output format
    with arf.open_file(args.output, mode="a") as ofp:
        writer = ArfWriter(ofp)
        log.info("- saving songs to '%s'", args.output)
        for fname in args.input:
            with h5.File(fname, "r") as ifp:
                log.info("- processing recordings in '%s'", ifp.filename)
                last_sample = None
                buffer = []
                intervals = []
                for entry in ifp.values():
                    if not isinstance(entry, h5.Group):
                        log.info("  ✗ %s: not an entry, skipping", entry.name)
                        continue
                    entry_time = timestamp_to_datetime(entry.attrs["timestamp"])
                    if args.ignore_before and entry_time < args.ignore_before:
                        log.info("  ✗ %s: %s is before %s, skipping", entry.name, entry_time, args.ignore_before)
                        continue
                    if args.ignore_after and entry_time > args.ignore_after:
                        log.info("  ✗ %s: %s is after %s, skipping", entry.name, entry_time, args.ignore_after)
                        continue
                    try:
                        dset = entry[args.dataset_name]
                    except KeyError:
                        log.info(
                            "  ✗ %s: does not have dataset `%s`, skipping",
                            entry.name,
                            args.dataset_name,
                        )
                        continue
                    sampling_rate = dset.attrs["sampling_rate"] / 1000.0
                    entry_start = entry.attrs["jack_frame"]
                    log.debug(
                        " - %s: start=%s, end=%s, sampling_rate=%.2f kHz",
                        entry.name,
                        entry_start,
                        entry_start + dset.size,
                        sampling_rate,
                    )
                    # is this a continuation of the previous entry?
                    if last_sample is not None and last_sample > entry_start:
                        log.debug("   - continues previous entry")
                    else:
                        # process the intervals
                        extract_intervals(
                            writer,
                            buffer,
                            intervals,
                            args.pad_before,
                            args.pad_after,
                        )
                        extractor = make_extractor(model["config"], sampling_rate)
                        finder.reset()
                        buffer = []
                        intervals = []
                    for block in extractor.process(dset):
                        pred = classifier.predict(block)
                        intervals.extend(
                            (start * dt, stop * dt)
                            for start, stop in finder.process(pred)
                        )
                    # jack_frame is a np.uint32 so it should overflow to allow
                    # comparison with next entry
                    last_sample = entry_start + dset.size
                    buffer.append(dset)


def extract_intervals(writer, dsets, intervals, pad_before, pad_after):
    import posixpath
    from datetime import timedelta

    from quicksong.core import pad_intervals

    fields_to_drop = (
        "entry_creator",
        "timestamp",
        "jack_frame",
        "jack_sampling_rate",
        "jack_usec",
        "trial_off",
    )

    if len(dsets) == 0:
        return
    dset_name = ", ".join(dset.parent.name for dset in dsets)
    if len(intervals) == 0:
        log.info("  ✗ %s: no song detected", dset_name)
        return
    sampling_rate = all_same(dset.attrs["sampling_rate"] for dset in dsets) / 1000.0
    assert sampling_rate is not None, "Entries don't have the same sampling rate!"
    parent_entry = dsets[0].parent
    entry_attrs = {
        k: v for k, v in parent_entry.attrs.items() if k not in fields_to_drop
    }
    entry_attrs.update(
        source_file=parent_entry.file.filename,
        source_uuid=b",".join(dset.parent.attrs["uuid"] for dset in dsets),
        entry_creator=f"org.meliza.quicksong/{script_name} {__version__}",
    )
    entry_attrs.pop("uuid")
    dset_attrs = dict(dsets[0].attrs)
    dset_attrs["source_uuid"] = b",".join(dset.attrs["uuid"] for dset in dsets)
    dset_attrs["datatype"] = DataTypes.ACOUSTIC
    dset_attrs.pop("uuid")
    signal = np.concatenate(dsets)
    log.info("  ✓ %s: ", dset_name)
    for start, end in pad_intervals(intervals, pad_before, pad_after):
        start_sample = max(0, int(start * sampling_rate))
        end_sample = min(signal.size, int(end * sampling_rate))
        song_timestamp = timestamp_to_datetime(
            parent_entry.attrs["timestamp"]
        ) + timedelta(seconds=start_sample / sampling_rate / 1000)
        new_entry = writer.create_entry(song_timestamp, **entry_attrs)
        new_dset = create_dataset(
            new_entry,
            posixpath.basename(dsets[0].name),
            signal[start_sample:end_sample],
            compression="gzip",
            **dset_attrs,
        )
        set_uuid(new_dset)
        log.info(
            "    - samples %d–%d (%.1f ms) -> %s%s",   # noqa: RUF001
            start_sample,
            end_sample,
            (end_sample - start_sample) / sampling_rate,
            writer.filename,
            new_dset.name,
        )


def main(argv=None):
    p = argparse.ArgumentParser(
        description="train the song detector on labeled recordings"
    )
    p.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"%(prog)s {__version__})",
    )
    p.add_argument("--debug", help="show verbose log messages", action="store_true")
    sub = p.add_subparsers(title="subcommands")

    pp = sub.add_parser("init", help="generate a starting configuration file")
    pp.set_defaults(func=make_config)
    pp.add_argument(
        "--train",
        action="store_true",
        help="run training with the intialized config file",
    )
    pp.add_argument(
        "config", type=Path, help="path where the configuration file should be saved"
    )
    pp.add_argument(
        "label_files",
        type=Path,
        nargs="*",
        help="label files to add to config file as training dataset",
    )

    pp = sub.add_parser("train", help="train the classifier")
    pp.set_defaults(func=train_classifier)
    pp.add_argument(
        "--dataset-dir",
        "-d",
        type=Path,
        default=Path("."),
        help="directory where training data files are stored",
    )
    pp.add_argument("config", type=Path, help="path of the configuration file (will be generated with default values if it doesn't exist)")
    pp.add_argument(
        "model",
        type=Path,
        nargs="?",
        help="path where the model should be saved (default is `config.pkl`)",
    )

    pp = sub.add_parser(
        "extract", help="use a trained model to extract songs from ARF files"
    )
    pp.set_defaults(func=extract_songs)
    pp.add_argument(
        "--dataset-name",
        default="pcm_000",
        help="the name of the dataset to process in each entry (default `%(default)s`)",
    )
    pp.add_argument(
        "--max-gap",
        type=float,
        default=100,
        help="merge segments separated by gaps less than this value (default %(default).1f ms)",
    )
    pp.add_argument(
        "--min-duration",
        type=float,
        default=500,
        help="segments must be at least this long (default %(default).1f ms)",
    )
    pp.add_argument(
        "--pad-before",
        type=float,
        default=1000,
        help="the amount of time before song onset to extract (default %(default).1f ms)",
    )
    pp.add_argument(
        "--pad-after",
        type=float,
        default=1000,
        help="the amount of time after song end to extract (default %(default).1f ms)",
    )
    pp.add_argument("--ignore-before", type=valid_datetime, help="ignore source entries before this time")
    pp.add_argument("--ignore-after", type=valid_datetime, help="ignore source entries after this time")
    pp.add_argument("model", help="trained classifier model (pkl file saved by `train`")
    pp.add_argument("input", help="the ARF file to process", nargs="+")
    pp.add_argument("output", help="output where the song segments should be saved")

    args = p.parse_args(argv)

    setup_log(log, args.debug)
    if not hasattr(args, "func"):
        p.print_usage()
        return 0

    args.func(args)


if __name__ == "__main__":
    main()
