"""
Functions for network matched-filter detection of seismic data.

Designed to cross-correlate templates generated by template_gen function
with data and output the detections.

:copyright:
    EQcorrscan developers.

:license:
    GNU Lesser General Public License, Version 3
    (https://www.gnu.org/copyleft/lesser.html)
"""
import os
import pickle
import logging
import numpy as np

from collections import defaultdict
from typing import List, Set
from timeit import default_timer

from concurrent.futures import ThreadPoolExecutor

from obspy import Stream, UTCDateTime

from eqcorrscan.core.match_filter.template import (
    Template, group_templates_by_seedid)
from eqcorrscan.core.match_filter.detection import Detection
from eqcorrscan.core.match_filter.party import Party
from eqcorrscan.core.match_filter.family import Family
from eqcorrscan.core.match_filter.helpers import _spike_test, _mad
from eqcorrscan.core.match_filter.matched_filter import MatchFilterError

from eqcorrscan.utils.correlate import (
    get_stream_xcorr, _stabalised_fmf, fftw_multi_normxcorr,
    _zero_invalid_correlation_sums, _set_inner_outer_threading)
from eqcorrscan.utils.pre_processing import (
    _check_daylong, _group_process)
from eqcorrscan.utils.findpeaks import multi_find_peaks
from eqcorrscan.utils.plotting import _match_filter_plot

Logger = logging.getLogger(__name__)


def _wildcard_fill(
    net: str, sta: str, loc: str, chan: str
) -> [str, str, str, str]:
    """
    Convert none to wildcards. Cope with seisan channel naming.

    .. rubric:: Example

    >>> _wildcard_fill(None, None, None, None)
    ('*', '*', '*', '*')
    >>> _wildcard_fill("NZ", "FOZ", "10", "HZ")
    ('NZ', 'FOZ', '10', 'H?Z')
    """
    if net in [None, '']:
        net = "*"
    if sta in [None, '']:
        sta = "*"
    if loc in [None, '']:
        loc = "*"
    if chan in [None, '']:
        chan = "*"
    # Cope with old seisan chans
    if len(chan) == 2:
        chan = f"{chan[0]}?{chan[-1]}"
    return net, sta, loc, chan


def _download_st(
    starttime: UTCDateTime,
    endtime: UTCDateTime,
    buff: float,
    min_gap: float,
    template_channel_ids: List[tuple],
    client,
    retries: int
) -> Stream:
    """
    Helper to download a stream from a client for a given start and end time.

    Applies `buff` to extend download to (heopfully) ensure all data are
    provided. Retries download up to `retries` times, and discards data
    with large gaps.

    :param starttime: Start time to download data from
    :param endtime: End time to download data to
    :param buff:
        Length to pad downloaded data by - some clients do not provide all
        data requested.
    :param min_gap: See core.match_filter.tribe.client_detect
    :param template_channel_ids:
    :param client:
        Client-like object with at least a .get_waveforms_bulk method.
    :param retries: See core.match_filter.tribe.client_detect

    :return: Stream as downloaded.
    """
    from obspy.clients.fdsn.header import FDSNException

    bulk_info = []
    for chan_id in template_channel_ids:
        bulk_info.append((
            chan_id[0], chan_id[1], chan_id[2], chan_id[3],
            starttime - buff, endtime + buff))

    for retry_attempt in range(retries):
        try:
            Logger.info(f"Downloading data between {starttime} and "
                        f"{endtime}")
            st = client.get_waveforms_bulk(bulk_info)
            Logger.info(
                "Downloaded data for {0} traces".format(len(st)))
            break
        except FDSNException as e:
            if "Split the request in smaller" in " ".join(e.args):
                Logger.warning(
                    "Datacentre does not support large requests: "
                    "splitting request into smaller chunks")
                st = Stream()
                for _bulk in bulk_info:
                    try:
                        st += client.get_waveforms_bulk([_bulk])
                    except Exception as e:
                        Logger.error("No data for {0}".format(_bulk))
                        Logger.error(e)
                        continue
                Logger.info("Downloaded data for {0} traces".format(
                    len(st)))
                break
        except Exception as e:
            Logger.error(e)
            continue
    else:
        raise MatchFilterError(
            "Could not download data after {0} attempts".format(
                retries))
    # Get gaps and remove traces as necessary
    if min_gap:
        gaps = st.get_gaps(min_gap=min_gap)
        if len(gaps) > 0:
            Logger.warning("Large gaps in downloaded data")
            st.merge()
            gappy_channels = list(
                set([(gap[0], gap[1], gap[2], gap[3])
                     for gap in gaps]))
            _st = Stream()
            for tr in st:
                tr_stats = (tr.stats.network, tr.stats.station,
                            tr.stats.location, tr.stats.channel)
                if tr_stats in gappy_channels:
                    Logger.warning(
                        "Removing gappy channel: {0}".format(tr))
                else:
                    _st += tr
            st = _st
            st.split()
    # Merge traces after gap checking
    st = st.merge()
    st.trim(starttime=starttime, endtime=endtime)

    st_ids = [tr.id for tr in st]
    # Remove traces that do not meet zero criteria
    st.traces = [tr for tr in st if _check_daylong(tr.data)]
    if len(st) < len(st_ids):
        lost_ids = " ".join([tr_id for tr_id in st_ids
                             if tr_id not in [tr.id for tr in st]])
        Logger.warning(
            f"Removed data for {lost_ids} due to more zero datapoints "
            f"than non-zero.")

    st_ids = [tr.id for tr in st]
    # Remove short traces
    st.traces = [
        tr for tr in st
        if tr.stats.endtime - tr.stats.starttime > 0.8 * (endtime - starttime)]
    if len(st) < len(st_ids):
        lost_ids = " ".join([tr_id for tr_id in st_ids
                             if tr_id not in [tr.id for tr in st]])
        Logger.warning(
            f"Removed data for {lost_ids} due to less than 80% of the "
            f"required length.")

    return st


def _pre_process(
    st: Stream,
    template_ids: set,
    pre_processed: bool,
    filt_order: int,
    highcut: float,
    lowcut: float,
    samp_rate: float,
    process_length: float,
    parallel: bool,
    cores: int,
    daylong: bool,
    ignore_length: bool,
    ignore_bad_data: bool,
    overlap: float, **kwargs
) -> Stream:
    """
    Basic matched-filter processing flow. Data are processing in-place.

    :param st: Stream to process
    :param template_ids:
        Iterable of seed ids in the template set. Only channels matching these
        seed ids will be retained.
    :param pre_processed: See core.match_filter.tribe.detect
    :param filt_order: See utils.pre_processing.multi_process
    :param highcut: See utils.pre_processing.multi_process
    :param lowcut: See utils.pre_processing.multi_process
    :param samp_rate: See utils.pre_processing.multi_process
    :param process_length: See utils.pre_processing.multi_process
    :param parallel: See utils.pre_processing.multi_process
    :param cores: See utils.pre_processing.multi_process
    :param daylong: See utils.pre_processing.multi_process
    :param ignore_length: See utils.pre_processing.multi_process
    :param overlap: See core.match_filter.tribe.detect
    :param ignore_bad_data: See utils.pre_processing.multi_process
    :param template_ids:

    :return: Processed stream
    """
    # Retain only channels that have matches in templates
    Logger.info(template_ids)
    st = Stream([tr for tr in st if tr.id in template_ids])
    Logger.info(f"Processing {(len(st))} channels")
    if len(st) == 0:
        raise IndexError(
            "No matching channels between stream and templates")
    tic = default_timer()
    _spike_test(st)
    toc = default_timer()
    Logger.info(f"Checking for spikes took {toc - tic:.4f} s")
    if not pre_processed:
        st_chunks = _group_process(
            filt_order=filt_order,
            highcut=highcut,
            lowcut=lowcut,
            samp_rate=samp_rate,
            process_length=process_length,
            parallel=parallel,
            cores=cores,
            stream=st,
            daylong=daylong,
            ignore_length=ignore_length,
            overlap=overlap,
            ignore_bad_data=ignore_bad_data)
    else:
        st_chunks = [st]
    Logger.info(f"Stream has been split into {len(st_chunks)} chunks")
    return st_chunks


def _group(
    sids: Set[str],
    templates: List[Template],
    group_size: int,
    groups: List[List[str]] = None
) -> List[List[Template]]:
    """
    Group templates either by seed id, or using pre-computed groups

    :param sids: Seed IDs available in stream
    :param templates: Templates to group
    :param group_size: Maximum group size
    :param groups: [Optional] List of List of template names in groups
    :return: Groups of templates.
    """
    Logger.info(f"Grouping for {sids}")
    if groups:
        Logger.info("Using pre-computed groups")
        t_dict = {t.name: t for t in templates}
        template_groups = []
        for grp in groups:
            template_group = [
                t_dict.get(t_name) for t_name in grp
                if t_name in t_dict.keys()]
            if len(template_group):
                template_groups.append(template_group)
        return template_groups
    template_groups = group_templates_by_seedid(
        templates=templates,
        st_seed_ids=sids,
        group_size=group_size)
    if len(template_groups) == 1 and len(template_groups[0]) == 0:
        Logger.error("No matching ids between stream and templates")
        raise IndexError("No matching ids between stream and templates")
    return template_groups


def _corr_and_peaks(
    templates: List[Template],
    template_names: List[str],
    stream: Stream,
    xcorr_func: str,
    concurrency: str,
    cores: int,
    i: int,
    export_cccsums: bool,
    parallel: bool,
    peak_cores: int,
    threshold: float,
    threshold_type: str,
    trig_int: float,
    sampling_rate: float,
    full_peaks: bool,
    plot: bool,
    plotdir: str,
    plot_format: str,
    prepped: bool = False,
    **kwargs
):
    """
    Compute cross-correlation between templates and a stream. Returns peaks in
     correlation function.

    :param templates: Templates to correlate
    :param template_names: Names of templates (ordered as templates)
    :param stream: Stream to correlate templates with
    :param xcorr_func: Cross-correlation function to use
    :param concurrency: Concurrency of cross-correlation function
    :param cores: Cores (threads) to use for cross-correlation
    :param i: Group-id (internal book-keeping)
    :param export_cccsums: Whether to export the raw cross-correlation sums
    :param parallel: Whether to compute peaks in parallel
    :param peak_cores: Number of cores (threads) to use for peak finding
    :param threshold: Threshold value (user-defined)
    :param threshold_type: Threshold type (e.g. MAD, ...)
    :param trig_int: Trigger interval in seconds
    :param sampling_rate: Sampling rate of data
    :param full_peaks: Whether to compute full peaks, or fast peaks.
    :param plot: Whether to plot correlation sums and peaks or not
    :param plotdir: Where to save plots if made
    :param plot_format: What format (extension) to use for plots.
    :param prepped:
        Whether data have already been prepared for correlation or not.
        If prepped, inputs change for a specific xcorr-function, see code.

    :return: Peaks, thresholds, number of channels, channels for each template
    """
    # Special cases for fmf and fftw to minimize reshaping time.
    Logger.info(
        f"Starting correlation run for template group {i}")
    tic = default_timer()
    if prepped and xcorr_func == "fmf":
        assert isinstance(templates, np.ndarray)
        assert isinstance(stream, np.ndarray)
        # These need to be passed from queues.
        pads = kwargs.get('pads')
        weights = kwargs.get('weights')
        chans = kwargs.get("chans")
        no_chans = kwargs.get("no_chans")
        # We do not care about removing the gain from our data, we copied it.
        multipliers = np.ones((len(stream), 1))
        step = 1  # We only implement single-step correlations
        if concurrency in ("multithread", "multiprocess"):
            arch = "cpu"
        else:
            arch = "gpu"
        cccsums = _stabalised_fmf(
            template_arr=templates, data_arr=stream, weights=weights,
            pads=pads, arch=arch, multipliers=multipliers, step=step)
    elif prepped and xcorr_func in ("fftw", None):
        assert isinstance(templates, dict)
        assert isinstance(stream, dict)
        pads = kwargs.pop('pads')
        seed_ids = kwargs.pop("seed_ids")
        num_cores_inner, num_cores_outer = _set_inner_outer_threading(
            kwargs.get('cores', None), kwargs.get("cores_outer", None),
            len(stream))

        cccsums, tr_chans = fftw_multi_normxcorr(
            template_array=templates, stream_array=stream,
            pad_array=pads, seed_ids=seed_ids, cores_inner=num_cores_inner,
            cores_outer=num_cores_outer, stack=True, **kwargs)
        n_templates = len(cccsums)
        # Post processing
        no_chans = np.sum(np.array(tr_chans).astype(int), axis=0)
        chans = [[] for _i in range(n_templates)]
        for seed_id, tr_chan in zip(seed_ids, tr_chans):
            for chan, state in zip(chans, tr_chan):
                if state:
                    chan.append(seed_id)
        cccsums = _zero_invalid_correlation_sums(cccsums, pads, chans)
        chans = [[(seed_id.split('.')[1], seed_id.split('.')[-1].split('_')[0])
                  for seed_id in _chans] for _chans in chans]
    else:
        # The default just uses stream xcorr funcs.
        multichannel_normxcorr = get_stream_xcorr(xcorr_func, concurrency)
        cccsums, no_chans, chans = multichannel_normxcorr(
            templates=templates, stream=stream, cores=cores, **kwargs
        )
    if len(cccsums[0]) == 0:
        raise MatchFilterError(
            f"Correlation has not run for group {i}, "
            f"zero length cccsum")
    toc = default_timer()
    Logger.info(
        f"Correlations for group {i} of {len(template_names)} "
        f"templates took {toc - tic:.4f} s")
    Logger.debug(
        f"The shape of the returned cccsums in group {i} "
        f"is: {cccsums.shape}")
    Logger.debug(
        f'This is from {len(templates)} templates correlated with '
        f'{len(stream)} channels of data in group {i}')

    # Handle saving correlation stats
    if export_cccsums:
        for i, cccsum in enumerate(cccsums):
            fname = (
                f"{template_names[i]}-{stream[0].stats.starttime}-"
                f"{stream[0].stats.endtime}_cccsum.npy")
            np.save(file=fname, arr=cccsum)
            Logger.info(
                f"Saved correlation statistic to {fname}")

    # Zero mean check
    if np.any(np.abs(cccsums.mean(axis=-1)) > 0.05):
        Logger.warning(
            'Mean of correlations is non-zero!  Check this!')
    if parallel:
        Logger.info(f"Finding peaks using {peak_cores} threads")
    else:
        Logger.info("Finding peaks in serial")
    # This is in the main process because transferring
    #  lots of large correlation sums in queues is very slow
    all_peaks, thresholds = _threshold(
        cccsums=cccsums, no_chans=no_chans,
        template_names=template_names, threshold=threshold,
        threshold_type=threshold_type,
        trig_int=int(trig_int * sampling_rate),
        parallel=parallel, full_peaks=full_peaks,
        peak_cores=peak_cores, plot=plot, stream=stream,
        plotdir=plotdir, plot_format=plot_format)
    return all_peaks, thresholds, no_chans, chans


def _threshold(
    cccsums: np.ndarray,
    no_chans: list,
    template_names: list,
    threshold: float,
    threshold_type: str,
    trig_int: int,  # converted to samples before getting to this func.
    parallel: bool,
    full_peaks: bool,
    peak_cores: int,
    plot: bool,
    stream: Stream,
    plotdir: str,
    plot_format: str,
):
    """
    Find peaks within correlation functions for given thresholds.

    :param cccsums: Numpy array of correlations [templates x samples]
    :param no_chans:
        Number of channels for each correlation (ordered as cccsums)
    :param template_names:
        Template names for each correlation (ordered as cccsums)
    :param threshold: Input threshold value
    :param threshold_type: Input threshold type (e.g. MAD, ...)
    :param trig_int: Trigger interval in SAMPLES.
    :param parallel: Whether to compute peaks in parallel
    :param full_peaks: Whether to compute full peaks or not
    :param peak_cores: Number of cores (threads) to use for peak finding.
    :param plot: Whether to plot the peak finding
    :param stream: Stream for plotting (not needed otherwise)
    :param plotdir: Directory to write plots to
    :param plot_format: Format to save plots in

    :return: (all peaks, used thresholds)
    """
    Logger.debug(f"Got cccsums shaped {cccsums.shape}")
    Logger.debug(f"From {len(template_names)} templates")

    tic = default_timer()
    if str(threshold_type) == str("absolute"):
        thresholds = [threshold for _ in range(len(cccsums))]
    elif str(threshold_type) == str('MAD'):
        median_cores = min([peak_cores, len(cccsums)])
        if cccsums.size < 2e7:  # parallelism not worth it
            median_cores = 1
        with ThreadPoolExecutor(max_workers=median_cores) as executor:
            # Because numpy releases GIL threading can use
            # multiple cores
            medians = executor.map(_mad, cccsums,
                                   chunksize=len(cccsums) // median_cores)
        thresholds = [threshold * median for median in medians]
    else:
        thresholds = [threshold * no_chans[i]
                      for i in range(len(cccsums))]
    toc = default_timer()
    Logger.info(f"Computing thresholds took {toc - tic: .4f} s")
    outtic = default_timer()
    all_peaks = multi_find_peaks(
        arr=cccsums, thresh=thresholds, parallel=parallel,
        trig_int=trig_int, full_peaks=full_peaks, cores=peak_cores)
    outtoc = default_timer()
    Logger.info(f"Finding peaks for group took {outtoc - outtic:.4f}s")

    # Plotting
    if plot and stream:
        for i, cccsum in enumerate(cccsums):
            _match_filter_plot(
                stream=stream, cccsum=cccsum,
                template_names=template_names,
                rawthresh=thresholds[i], plotdir=plotdir,
                plot_format=plot_format, i=i)
        else:
            Logger.error("Plotting enabled but not stream found to plot")

    return all_peaks, thresholds


def _detect(
    template_names: List[str],
    all_peaks: np.ndarray,
    starttime: UTCDateTime,
    delta: float,
    no_chans: List[int],
    chans: List[List[str]],
    thresholds: List[float]
) -> List[Detection]:
    """
    Convert peaks to Detection objects

    :param template_names: Lis of template names
    :param all_peaks: Array of peaks orders as template_names
    :param starttime: Starttime for peak index relative time
    :param delta: Sample interval to convert peaks from samples to time
    :param no_chans: Number of channels used (ordered as template_names)
    :param chans: Channels used (ordered as template_names)
    :param thresholds: Thresholds used (ordered as template_names)

    :return: List of detections.
    """
    tic = default_timer()
    detections = []
    for i, template_name in enumerate(template_names):
        if not all_peaks[i]:
            Logger.debug(f"Found 0 peaks for template {template_name}")
            continue
        Logger.debug(f"Found {len(all_peaks[i])} detections "
                     f"for template {template_name}")
        for peak in all_peaks[i]:
            detecttime = starttime + (peak[1] * delta)
            if peak[0] > no_chans[i]:
                Logger.error(f"Correlation sum {peak[0]} exceeds "
                             f"bounds ({no_chans[i]}")
            detection = Detection(
                template_name=template_name, detect_time=detecttime,
                no_chans=no_chans[i], detect_val=peak[0],
                threshold=thresholds[i], typeofdet='corr',
                chans=chans[i],
                threshold_type=None,
                # Update threshold_type and threshold outside of this func.
                threshold_input=None)
            detections.append(detection)
    toc = default_timer()
    Logger.info(f"Forming detections took {toc - tic:.4f} s")
    return detections


def _load_template(t_file: str) -> Template:
    """ Load a pickled template from a file """
    try:
        with open(t_file, "rb") as f:
            t = pickle.load(f)
    except Exception as e:
        Logger.warning(f"Could not read template from {t_file} due to {e}")
        return None
    assert isinstance(t, Template), "Loaded object is not a Template, aborting"
    return t


def _read_template_db(template_file_dict: dict) -> List[Template]:
    """
    Read templates from files on disk.

    :param template_file_dict: Template file names keyed by template name

    :returns: list of templates
    """
    with ThreadPoolExecutor() as executor:
        templates = executor.map(_load_template, template_file_dict.values())
    templates = [t for t in templates if t]
    Logger.info(f"Deserialized {len(templates)} templates")
    if len(templates) < len(template_file_dict):
        Logger.warning(f"Expected {len(template_file_dict)} templates, "
                       f"but found {len(templates)}")
    return templates


def _make_party(
    detections: List[Detection],
    threshold: float,
    threshold_type: str,
    templates: List[Template],
    chunk_start: UTCDateTime,
    chunk_id: int,
    save_progress: bool
) -> str:
    """
    Construct a Party from Detections.

    :param detections: List of detections
    :param threshold: Input threshold
    :param threshold_type: Input threshold type
    :param templates: Templates used in detections
    :param chunk_start: Starttime of party epoch
    :param chunk_id: Internal index for party epoch
    :param save_progress: Whether to save progress or not

    :return: The filename the party has been pickled to.
    """
    chunk_dir = os.path.join(
        ".parties", "{chunk_start.year}", "{chunk_start.julday:03d}")
    chunk_file_str = os.path.join(
        chunk_dir, "chunk_party_{chunk_start_str}_{chunk_id}_{pid}.pkl")
    # Process ID included in chunk file to avoid multiple processes writing
    # and reading and removing the same files.

    # Get the results out of the end!
    Logger.info(f"Made {len(detections)} detections")

    # post - add in threshold, threshold_type to all detections
    Logger.info("Adding threshold to detections")
    for detection in detections:
        detection.threshold_input = threshold
        detection.threshold_type = threshold_type

    # Select detections very quickly: detection order does not
    # change, make dict of keys: template-names and values:
    # list of indices and use indices to select
    Logger.info("Making dict of detections")
    detection_idx_dict = defaultdict(list)
    for n, detection in enumerate(detections):
        detection_idx_dict[detection.template_name].append(n)

    # Convert to Families and build party.
    Logger.info("Converting to party and making events")
    chunk_party = Party()

    # Make a dictionary of templates keyed by name - we could be passed a dict
    # of pickled templates
    if not isinstance(templates, dict):
        templates = {t.name: t for t in templates}

    for t_name, template in templates.items():
        family_detections = [
            detections[idx]
            for idx in detection_idx_dict[t_name]]
        # Make party sparse - only write out families with detections
        if len(family_detections):
            if not isinstance(template, Template):
                # Try and read this from disk
                with open(template, "rb") as f:
                    template = pickle.load(f)
            for d in family_detections:
                d._calculate_event(template=template)
            family = Family(
                template=template, detections=family_detections)
            chunk_party += family

    Logger.info("Pickling party")
    if not os.path.isdir(chunk_dir.format(chunk_start=chunk_start)):
        os.makedirs(chunk_dir.format(chunk_start=chunk_start))

    chunk_file = chunk_file_str.format(
        chunk_start_str=chunk_start.strftime("%Y-%m-%dT%H-%M-%S"),
        chunk_start=chunk_start,
        chunk_id=chunk_id, pid=os.getpid())
    with open(chunk_file, "wb") as _f:
        pickle.dump(chunk_party, _f)
    Logger.info("Completed party processing")

    if save_progress:
        Logger.info(f"Written chunk to {chunk_file}")
    return chunk_file


if __name__ == "__main__":
    import doctest

    doctest.testmod()
