"""Context methods dealing with run scanning and selection"""
import fnmatch
import re
import typing as ty
import warnings
from collections import defaultdict
import numpy as np
import pandas as pd
import pytz
import datetime
import strax
# use tqdm as loaded in utils (from tqdm.notebook when in a juypyter env)
tqdm = strax.utils.tqdm

export, __all__ = strax.exporter()


@strax.Context.add_method
def list_available(self, target, **kwargs):
    """Return sorted list of run_id's for which target is available
    """
    if len(kwargs):
        # noinspection PyMethodFirstArgAssignment
        self = self.new_context(**kwargs)

    if self.runs is None:
        self.scan_runs()

    keys = set(self.keys_for_runs(
        target,
        self.runs['name'].values))

    found = set()
    for sf in self.storage:
        remaining = keys - found
        is_found = sf.find_several(list(remaining), **self._find_options)
        found |= set([k for i, k in enumerate(remaining)
                      if is_found[i]])
    return list(sorted([x.run_id for x in found]))


@strax.Context.add_method
def keys_for_runs(self,
                  target: str,
                  run_ids: ty.Union[np.ndarray, list, tuple, str]
                  ) -> ty.List[strax.DataKey]:
    """
    Get the data-keys for a multitude of runs. If use_per_run_defaults
        is False which it preferably is (#246), getting many keys should
        be fast as we only only compute the lineage once.

    :param run_ids: Runs to get datakeys for
    :param target: datatype requested
    :return: list of datakeys of the target for the given runs.
    """
    run_ids = strax.to_str_tuple(run_ids)

    if self.context_config['use_per_run_defaults']:
        return [self.key_for(r, target) for r in run_ids]
    elif len(run_ids):
        # Get the lineage once, for the context specifies that the
        # defaults  may not change!
        p = self._get_plugins((target,), run_ids[0])[target]
        return [strax.DataKey(r, target, p.lineage) for r in run_ids]
    else:
        return []


@strax.Context.add_method
def scan_runs(self: strax.Context,
              check_available=tuple(),
              store_fields=tuple()):
    """Update and return self.runs with runs currently available
    in all storage frontends.
    :param check_available: Check whether these data types are available
    Availability of xxx is stored as a boolean in the xxx_available
    column.
    :param store_fields: Additional fields from run doc to include
    as rows in the dataframe.

    The context options scan_availability and store_run_fields list
    data types and run fields, respectively, that will always be scanned.
    """
    store_fields = tuple(set(
        list(strax.to_str_tuple(store_fields))
        + ['name', 'number', 'tags', 'mode',
           strax.RUN_DEFAULTS_KEY]
        + list(self.context_config['store_run_fields'])))
    check_available = tuple(set(
        list(strax.to_str_tuple(check_available))
        + list(self.context_config['check_available'])))

    for target in check_available:
        save_when = self.get_save_when(target)
        if save_when < strax.SaveWhen.ALWAYS:
            p = self._plugin_class_registry[target]
            self.log.warning(f'{p.__name__}-plugin is {save_when}. '
                             f'Therefore {target} is most likely not stored!')

    docs = None
    for sf in self.storage:
        _temp_docs = []
        for doc in sf._scan_runs(store_fields=store_fields):
            # If there is no number, make one from the name
            _is_superrun = doc.get('name', 'not_a_superrun_if_no_name').startswith('_')
            if 'number' not in doc:
                if 'name' not in doc:
                    raise ValueError(f"Invalid run doc {doc}, contains "
                                     f"neither name nor number.")
                if not _is_superrun:
                    # If there is no name, make one from the number
                    doc['number'] = int(doc['name'])
            if not _is_superrun:
                doc.setdefault('name', f"{doc['number']:06d}")

            # Convert tags/mode/source list to a,separated string if needed (mode and source can
            # be lists for superruns)
            doc.setdefault('mode', '')
            if type(doc['mode']) == list:
                doc['mode'] = ','.join(doc['mode'])

            doc.setdefault('source', '')
            if type(doc['source']) == list:
                doc['source'] = ','.join(doc['source'])

            doc['tags'] = ','.join([t['name'] if type(t) == dict else t
                                   for t in doc.get('tags', [])])

            # Set a default livetime if we have start and stop
            if ('start' in store_fields
                    and 'end' in store_fields
                    and 'livetime' in store_fields
                    and doc.get('start') is not None
                    and doc.get('end') is not None):
                doc.setdefault('livetime', doc['end'] - doc['start'])

            if _is_superrun:
                # In contrast to regular run-docs, 
                # superruns are timezone aware. So strip off timezone 
                # again:
                start = doc.get('start')
                if start:
                    doc['start'] = start.replace(tzinfo=None)
                    
                start = doc.get('end')
                if start:
                    doc['end'] = start.replace(tzinfo=None)
                    
                if type(doc.get('livetime')) == float:
                    # If we have a superrun livetime is stored as intger seconds
                    # as timedelta is not json serializable 
                    doc['livetime'] = datetime.timedelta(seconds=doc['livetime'])
            
            # Put the strax defaults stuff into a different cache
            if strax.RUN_DEFAULTS_KEY in doc:
                self._run_defaults_cache[doc['name']] = \
                    doc[strax.RUN_DEFAULTS_KEY]
                del doc[strax.RUN_DEFAULTS_KEY]

            doc = flatten_run_metadata(doc)

            _temp_docs.append(doc)

        if len(_temp_docs):
            new_docs = pd.DataFrame(_temp_docs)
        else:
            new_docs = pd.DataFrame([], columns=store_fields)

        if docs is None:
            docs = new_docs
        else:
            # Keep only new runs (not found by earlier frontends)
            docs = pd.concat([
                docs,
                new_docs[
                    ~np.in1d(new_docs['name'], docs['name'])]],
                sort=False)
            docs.reset_index(drop=True, inplace=True)

    # Rearrange columns
    if (not self.context_config['use_per_run_defaults']
        and strax.RUN_DEFAULTS_KEY in docs.columns):
        del docs[strax.RUN_DEFAULTS_KEY]
    docs = docs[['name'] + [x for x in docs.columns.tolist()
                            if x != 'name']]
    self.runs = docs

    for d in tqdm(check_available,
                  desc='Checking data availability'):
        self.runs[d + '_available'] = np.in1d(
            self.runs.name.values,
            self.list_available(d))

    return self.runs


@strax.Context.add_method
def select_runs(self, run_mode=None, run_id=None,
                include_tags=None, exclude_tags=None,
                available=tuple(),
                pattern_type='fnmatch', ignore_underscore=True):
    """Return pandas.DataFrame with basic info from runs
    that match selection criteria.
    :param run_mode: Pattern to match run modes (reader.ini.name)
    :param run_id: Pattern to match a run_id or run_ids
    :param available: str or tuple of strs of data types for which data
    must be available according to the runs DB.

    :param include_tags: String or list of strings of patterns
        for required tags
    :param exclude_tags: String / list of strings of patterns
        for forbidden tags.
        Exclusion criteria  have higher priority than inclusion criteria.
    :param pattern_type: Type of pattern matching to use.
        Defaults to 'fnmatch', which means you can use
        unix shell-style wildcards (`?`, `*`).
        The alternative is 're', which means you can use
        full python regular expressions.
    :param ignore_underscore: Ignore the underscore at the start of tags
        (indicating some degree of officialness or automation).

    Examples:
     - `run_selection(include_tags='blinded')`
        select all datasets with a blinded or _blinded tag.
     - `run_selection(include_tags='*blinded')`
        ... with blinded or _blinded, unblinded, blablinded, etc.
     - `run_selection(include_tags=['blinded', 'unblinded'])`
        ... with blinded OR unblinded, but not blablinded.
     - `run_selection(include_tags='blinded',
                      exclude_tags=['bad', 'messy'])`
       select blinded dsatasets that aren't bad or messy
    """
    if self.runs is None:
        self.scan_runs(check_available=strax.to_str_tuple(available))
    dsets = self.runs.copy()

    if pattern_type not in ('re', 'fnmatch'):
        raise ValueError("Pattern type must be 're' or 'fnmatch'")

    # Filter datasets by run mode and/or name
    for field_name, requested_value in (
            ('name', run_id),
            ('mode', run_mode)):

        if requested_value is None:
            continue
            
        requested_value = strax.to_str_tuple(requested_value)

        values = dsets[field_name].values
        mask = np.zeros(len(values), dtype=np.bool_)

        if pattern_type == 'fnmatch':
            for i, x in enumerate(values):
                mask[i] = np.any([fnmatch.fnmatch(x, rv) for rv in requested_value]) 
        elif pattern_type == 're':
            for i, x in enumerate(values):
                mask[i] = np.any([re.match(rv, x) for rv in requested_value])

        dsets = dsets[mask]

    dsets = _include_exclude_tags(dsets, include_tags, exclude_tags, pattern_type,
                                  ignore_underscore,
                                  )

    have_available = strax.to_str_tuple(available)
    for d in have_available:
        if not d + '_available' in dsets.columns:
            # Get extra availability info from the run db
            d_available = np.in1d(self.runs.name.values,
                                  self.list_available(d))
            # Save both in the context and for this selection using
            # available = ('data_type',)
            self.runs[d + '_available'] = d_available
            dsets[d + '_available'] = d_available
    for d in have_available:
        dsets = dsets[dsets[d + '_available']]

    return dsets


@strax.Context.add_method
def define_run(self: strax.Context,
               name: str,
               data: ty.Union[np.ndarray, pd.DataFrame, dict, list, tuple],
               from_run: ty.Union[str, None] = None):
    """
    Function for defining new superruns from a list of run_ids.

    Note:
        The function also allows to create a superrun from data
        (numpy.arrays/pandas.DataFframrs). However, this is currently
        not supported from the data loading side.

    :param name: Name/run_id of the superrun. Suoerrun names must start
        with an underscore.
    :param data: Data from which the superrun should be created. Can be
        either one of the following: a tuple/list of run_ids or a
        numpy.array/pandas.DataFrame containing some data.
    :param from_run: List of run_ids which were used to create the
        numpy.array/pandas.DataFrame passed in data.
    """
    if isinstance(data, (pd.DataFrame, np.ndarray)):
        if isinstance(data, np.ndarray):
            data = pd.DataFrame.from_records(data)
      
        # strax.endtime does not work with DataFrames due to numba
        if 'endtime' in data.columns:
            end = data['endtime']
        else:
            end = data['time'] + data['length'] * data['dt']
            
        # Array of events / regions of interest
        start, end = data['time'], end
        if from_run is not None:
            return self.define_run(
                name,
                {from_run: np.transpose([start, end])})
        elif 'run_id' not in data.columns:
            raise ValueError(
                "Must provide from_run or data with a run_id column "
                "to define a superrun")
        else:
            df = pd.DataFrame(dict(start=start, end=end,
                                   run_id=data['run_id']))
            return self.define_run(
                name,
                {run_id: rs[['start', 'end']].values.transpose()
                 for run_id, rs in df.groupby('run_id')})

    if isinstance(data, (list, tuple)):
        # list of runids
        data = strax.to_str_tuple(data)
        return self.define_run(
            name,
            {run_id: 'all' for run_id in data})

    if not isinstance(data, dict):
        raise ValueError(f"Can't define run from {type(data)}")

    # Find start and end time of the new run = earliest start time of other runs
    run_md = dict(start=datetime.datetime.max.replace(tzinfo=pytz.utc), 
                  end=datetime.datetime.min.replace(tzinfo=pytz.utc), 
                  livetime=0)
    keys = []
    starts = []
    tags = set()
    modes = set()
    sources = set()
    for _subrunid in data:
        doc = self.run_metadata(_subrunid, ['start', 'end', 'mode', 'tags', 'source'])
        doc.setdefault('tags', [{'name': ''},])
        doc.setdefault('mode', '')
        doc.setdefault('source', '')

        tags |= set([tag['name'] for tag in doc['tags']])

        modes |= set(strax.to_str_tuple(doc['mode']))
        sources |= set(strax.to_str_tuple(doc['source']))
        if len(sources) > 1:
            warnings.warn(f'You are defining a superrun with more than one source: "{sources}"')

        run_doc_start = doc['start'].replace(tzinfo=pytz.utc)
        run_doc_end = doc['end'].replace(tzinfo=pytz.utc)

        run_md['start'] = min(run_md['start'], run_doc_start)
        run_md['end'] = max(run_md['end'], run_doc_end)

        time_delta = run_doc_end - run_doc_start
        run_md['livetime'] += time_delta.total_seconds()
        keys.append(_subrunid)
        starts.append(run_doc_start)

    run_md['tags'] = tuple(tags)
    run_md['mode'] = tuple(modes)
    run_md['source'] = tuple(sources)

    # Make sure subruns are sorted in time
    sort_index = np.argsort(starts)
    data = {keys[i]: data[keys[i]] for i in sort_index}

    # Superrun names must start with an underscore
    if not name.startswith('_'):
        name = '_' + name
    # Dict mapping run_id: array of time ranges or all
    for sf in self.storage:
        if not sf.readonly and sf.can_define_runs:
            sf.define_run(name, sub_run_spec=data, **run_md)
            break
    else:
        raise RuntimeError("No storage frontend registered that allows"
                           " run definition")


@strax.Context.add_method
def available_for_run(self: strax.Context,
                      run_id: str,
                      include_targets: ty.Union[None, list, tuple, str] = None,
                      exclude_targets: ty.Union[None, list, tuple, str] = None,
                      pattern_type: str = 'fnmatch') -> pd.DataFrame:
    """
    For a given single run, check all the targets if they are stored.
        Excludes the target if never stored anyway.
    :param run_id: requested run
    :param include_targets: targets to include e.g. raw_records,
        raw_records* or *_nv. If multiple targets (e.g. a list) is
        provided, the target should match any of the arguments!
    :param exclude_targets: targets to exclude e.g. raw_records,
        raw_records* or *_nv. If multiple targets (e.g. a list) is
        provided, the target should match none of the arguments!
    :param pattern_type: either 'fnmatch' (Unix filename pattern
        matching) or 're' (Regular expression operations).
    :return: Table of available data per target
    """
    if not isinstance(run_id, str):
        raise ValueError(f'Only single run_id is allowed (str),'
                         f' got {run_id} ({type(run_id)})')

    if exclude_targets is None:
        exclude_targets = []
    if include_targets is None:
        include_targets = []

    is_stored = defaultdict(list)
    for target in self._plugin_class_registry.keys():
        # Skip targets that are not stored
        save_when = self.get_save_when(target)
        if save_when == strax.SaveWhen.NEVER:
            continue

        # Should we include this target or exclude it?
        include_t = []
        exclude_t = False

        for excl in strax.to_str_tuple(exclude_targets):
            # Simple logic, if we match the excluded target, we should
            # should not continue
            if _tag_match(target, excl, pattern_type, False):
                exclude_t = True
                break

        # We can match any of the "incl" targets, keep a list and check
        # of any of the "incl" matches the target.
        for incl in strax.to_str_tuple(include_targets):
            include_t.append(_tag_match(target, incl, pattern_type, False))

        # Convert to simple bool. If no include_targets is specified,
        # all are fine, otherwise check at least one is matching.
        include_t = True if not len(include_t) else any(include_t)

        if include_t and not exclude_t:
            is_stored['target'].append(target)
            is_stored['is_stored'].append(self.is_stored(run_id, target))
    return pd.DataFrame(is_stored)


def _include_exclude_tags(dsets,
                          include_tags,
                          exclude_tags,
                          pattern_type,
                          ignore_underscore,
                          ):
    if include_tags is not None:
        dsets = dsets[_tags_match(dsets,
                                  include_tags,
                                  pattern_type,
                                  ignore_underscore)]

    if exclude_tags is not None:
        dsets = dsets[True ^ _tags_match(dsets,
                                         exclude_tags,
                                         pattern_type,
                                         ignore_underscore)]
    return dsets


def _tags_match(dsets, patterns, pattern_type, ignore_underscore):
    result = np.zeros(len(dsets), dtype=np.bool_)

    if isinstance(patterns, str):
        patterns = [patterns]

    for i, tags in enumerate(dsets.tags):
        result[i] = any([any([_tag_match(tag, pattern,
                                         pattern_type,
                                         ignore_underscore)
                              for tag in tags.split(',')
                              for pattern in patterns])])

    return result


def _tag_match(tag, pattern, pattern_type, ignore_underscore):
    if ignore_underscore and tag.startswith('_') and not pattern.startswith('_'):
        tag = tag[1:]
    if pattern_type == 'fnmatch':
        return fnmatch.fnmatch(tag, pattern)
    elif pattern_type == 're':
        return bool(re.match(pattern, tag))
    raise NotImplementedError


@export
def flatten_run_metadata(md):
    # Flatten the tags field. Note this sets it to an empty string
    # if it does not exist.
    return strax.flatten_dict(
        md,
        separator='.',
        keep=[strax.RUN_DEFAULTS_KEY, 'sub_run_spec', 'tags'])
