"""
Utility functions for manipulating neuprint-python output.
"""
import sys
import inspect
import functools
import warnings
from textwrap import dedent
from collections.abc import Iterable, Iterator, Collection

import numpy as np
import pandas as pd


class NotNull:
    """Filter for existing properties.

    Translates to::

        WHERE neuron.{property} IS NOT NULL

    """

    pass


class IsNull:
    """Filter for missing properties.

    Translates to::

        WHERE neuron.{property} IS NULL

    """

    pass

#
# Import the notebook-aware version of tqdm if
# we appear to be running within a notebook context.
#
try:
    import ipykernel.iostream
    if isinstance(sys.stdout, ipykernel.iostream.OutStream):
        from tqdm.notebook import tqdm

        try:
            import ipywidgets
            ipywidgets
        except ImportError:
            msg = dedent("""\

                Progress bar will not work well in the notebook without ipywidgets.
                Run the following commands (for notebook and jupyterlab users):

                    conda install -c conda-forge ipywidgets
                    jupyter nbextension enable --py widgetsnbextension
                    jupyter labextension install @jupyter-widgets/jupyterlab-manager

                ...and then reload your jupyter session, and restart your kernel.
            """)
            warnings.warn(msg)
    else:
        from tqdm import tqdm

except ImportError:
    from tqdm import tqdm


class tqdm(tqdm):
    """
    Same as tqdm, but auto-disable the progress bar if there's only one item.
    """
    def __init__(self, iterable=None, *args, disable=None, **kwargs):
        if disable is None:
            disable = (iterable is not None
                       and hasattr(iterable, '__len__')
                       and len(iterable) <= 1)

        super().__init__(iterable, *args, disable=disable, **kwargs)


def trange(*args, **kwargs):
    return tqdm(range(*args), **kwargs)


def UMAP(*args, **kwargs):
    """
    UMAP is an optional dependency, so this wrapper emits
    a nicer error message if it's not available.
    """
    try:
        from umap import UMAP
    except ImportError as ex:
        msg = (
            "The 'umap' dimensionality reduction package is required for some "
            "plotting functionality, but it isn't currently installed.\n\n"
            "Please install it:\n\n"
            "  conda install -c conda-forge umap-learn\n\n"
        )
        raise RuntimeError(msg) from ex

    return UMAP(*args, **kwargs)


def make_iterable(x):
    """
    If ``x`` is already a list or array, return it unchanged.
    If ``x`` is Series, return its values.
    If ``x`` is ``None``, return an empty list ``[]``.
    Otherwise, wrap it in a list.
    """
    if x is None:
        return []

    if isinstance(x, np.ndarray):
        return x

    if isinstance(x, pd.Series):
        return x.values

    if isinstance(x, Collection) and not isinstance(x, str):
        return x
    else:
        return [x]


def make_args_iterable(argnames):
    """
    Returns a decorator.
    For the given argument names, the decorator converts the
    arguments into iterables via ``make_iterable()``.
    """
    def decorator(f):

        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            callargs = inspect.getcallargs(f, *args, **kwargs)
            for name in argnames:
                callargs[name] = make_iterable(callargs[name])
            return f(**callargs)

        wrapper.__signature__ = inspect.signature(f)
        return wrapper

    return decorator


@make_args_iterable(['properties'])
def merge_neuron_properties(neuron_df, conn_df, properties=['type', 'instance']):
    """
    Merge neuron properties to a connection table.

    Given a table of neuron properties and a connection table, append
    ``_pre`` and ``_post`` columns to the connection table for each of
    the given properties via the appropriate merge operations.

    Args:
        neuron_df:
            DataFrame with columns for 'bodyId' and any properties you want to merge

        conn_df:
            DataFrame with columns ``bodyId_pre`` and ``bodyId_post``

        properties:
            Column names from ``neuron_df`` to merge onto ``conn_df``.

    Returns:
        Updated ``conn_df`` with new columns.

    Example:

        .. code-block:: ipython

            In [1]: from neuprint import fetch_adjacencies, NeuronCriteria as NC, merge_neuron_properties
               ...: neuron_df, conn_df = fetch_adjacencies(rois='PB', min_roi_weight=120)
               ...: print(conn_df)
               bodyId_pre  bodyId_post roi  weight
            0   880875736   1631450739  PB     123
            1   880880259    849421763  PB     141
            2   910442723    849421763  PB     139
            3   910783961   5813070465  PB     184
            4   911129204    724280817  PB     127
            5   911134009    849421763  PB     125
            6   911565419   5813070465  PB     141
            7   911911004   1062526223  PB     125
            8   911919044    973566036  PB     122
            9  5813080838    974239375  PB     136

            In [2]: merge_neuron_properties(neuron_df, conn_df, 'type')
            Out[2]:
               bodyId_pre  bodyId_post roi  weight  type_pre    type_post
            0   880875736   1631450739  PB     123  Delta7_a  PEN_b(PEN2)
            1   880880259    849421763  PB     141  Delta7_a  PEN_b(PEN2)
            2   910442723    849421763  PB     139  Delta7_a  PEN_b(PEN2)
            3   910783961   5813070465  PB     184  Delta7_a  PEN_b(PEN2)
            4   911129204    724280817  PB     127  Delta7_a  PEN_b(PEN2)
            5   911134009    849421763  PB     125  Delta7_a  PEN_b(PEN2)
            6   911565419   5813070465  PB     141  Delta7_a  PEN_b(PEN2)
            7   911911004   1062526223  PB     125  Delta7_b  PEN_b(PEN2)
            8   911919044    973566036  PB     122  Delta7_a  PEN_b(PEN2)
            9  5813080838    974239375  PB     136       EPG          PEG
    """
    neuron_df = neuron_df[['bodyId', *properties]]

    newcols  = [f'{prop}_pre'  for prop in properties]
    newcols += [f'{prop}_post' for prop in properties]
    conn_df = conn_df.drop(columns=newcols, errors='ignore')

    conn_df = conn_df.merge(neuron_df, 'left', left_on='bodyId_pre', right_on='bodyId')
    del conn_df['bodyId']

    conn_df = conn_df.merge(neuron_df, 'left', left_on='bodyId_post', right_on='bodyId',
                            suffixes=['_pre', '_post'])
    del conn_df['bodyId']

    return conn_df


def connection_table_to_matrix(conn_df, group_cols='bodyId', weight_col='weight', sort_by=None, make_square=False):
    """
    Given a weighted connection table, produce a weighted adjacency matrix.

    Args:
        conn_df:
            A DataFrame with columns for pre- and post- identifiers
            (e.g. bodyId, type or instance), and a column for the
            weight of the connection.

        group_cols:
            Which two columns to use as the row index and column index
            of the returned matrix, respetively.
            Or give a single string (e.g. ``"body"``, in which case the
            two column names are chosen by appending the suffixes
            ``_pre`` and ``_post`` to your string.

            If a pair of pre/post values occurs more than once in the
            connection table, all of its weights will be summed in the
            output matrix.

        weight_col:
            Which column holds the connection weight, to be aggregated for each unique pre/post pair.

        sort_by:
            How to sort the rows and columns of the result.
            Can be two strings, e.g. ``("type_pre", "type_post")``,
            or a single string, e.g. ``"type"`` in which case the suffixes are assumed.

        make_square:
            If True, insert rows and columns to ensure that the same IDs exist in the rows and columns.
            Inserted entries will have value 0.0

    Returns:
        DataFrame, shape NxM, where N is the number of unique values in
        the 'pre' group column, and M is the number of unique values in
        the 'post' group column.

    Example:

        .. code-block:: ipython

            In [1]: from neuprint import fetch_simple_connections, NeuronCriteria as NC
               ...: kc_criteria = NC(type='KC.*')
               ...: conn_df = fetch_simple_connections(kc_criteria, kc_criteria)
            In [1]: conn_df.head()
            Out[1]:
               bodyId_pre  bodyId_post  weight type_pre type_post instance_pre instance_post                                       conn_roiInfo
            0  1224137495   5813032771      29      KCg       KCg          KCg    KCg(super)  {'MB(R)': {'pre': 26, 'post': 26}, 'gL(R)': {'...
            1  1172713521   5813067826      27      KCg       KCg   KCg(super)         KCg-d  {'MB(R)': {'pre': 26, 'post': 26}, 'PED(R)': {...
            2   517858947   5813032943      26   KCab-p    KCab-p       KCab-p        KCab-p  {'MB(R)': {'pre': 25, 'post': 25}, 'PED(R)': {...
            3   642680826   5812980940      25   KCab-p    KCab-p       KCab-p        KCab-p  {'MB(R)': {'pre': 25, 'post': 25}, 'PED(R)': {...
            4  5813067826   1172713521      24      KCg       KCg        KCg-d    KCg(super)  {'MB(R)': {'pre': 23, 'post': 23}, 'gL(R)': {'...

            In [2]: from neuprint.utils import connection_table_to_matrix
               ...: connection_table_to_matrix(conn_df, 'type')
            Out[2]:
            type_post   KC  KCa'b'  KCab-p  KCab-sc     KCg
            type_pre
            KC           3     139       6        5     365
            KCa'b'     154  102337     245      997    1977
            KCab-p       7     310   17899     3029     127
            KCab-sc      4    2591    3975   247038    3419
            KCg        380    1969      79     1526  250351
    """
    if isinstance(group_cols, str):
        group_cols = (f"{group_cols}_pre", f"{group_cols}_post")

    assert len(group_cols) == 2, \
        "Please provide two group_cols (e.g. 'bodyId_pre', 'bodyId_post')"

    assert group_cols[0] in conn_df, \
        f"Column missing: {group_cols[0]}"

    assert group_cols[1] in conn_df, \
        f"Column missing: {group_cols[1]}"

    assert weight_col in conn_df, \
        f"Column missing: {weight_col}"

    col_pre, col_post = group_cols
    dtype = conn_df[weight_col].dtype

    agg_weights_df = conn_df.groupby([col_pre, col_post], sort=False)[weight_col].sum().reset_index()
    matrix = agg_weights_df.pivot(col_pre, col_post, weight_col)
    matrix = matrix.fillna(0).astype(dtype)

    if sort_by:
        if isinstance(sort_by, str):
            sort_by = (f"{sort_by}_pre", f"{sort_by}_post")

        assert len(sort_by) == 2, \
            "Please provide two sort_by column names (e.g. 'type_pre', 'type_post')"

        pre_order = conn_df.sort_values(sort_by[0])[col_pre].unique()
        post_order = conn_df.sort_values(sort_by[1])[col_post].unique()
        matrix = matrix.reindex(index=pre_order, columns=post_order)
    else:
        # No sort: Keep the order as close to the input order as possible.
        pre_order = conn_df[col_pre].unique()
        post_order = conn_df[col_post].unique()
        matrix = matrix.reindex(index=pre_order, columns=post_order)

    if make_square:
        matrix, _ = matrix.align(matrix.T).fillna(0.0).astype(matrix.dtype)
        matrix = matrix.rename_axis('bodyId_pre', axis=0).rename_axis('bodyId_post', axis=1)
        matrix = matrix.loc[sorted(matrix.index), sorted(matrix.columns)]

    return matrix


def iter_batches(it, batch_size):
    """
    Iterator.

    Consume the given iterator/iterable in batches and
    yield each batch as a list of items.

    The last batch might be smaller than the others,
    if there aren't enough items to fill it.

    If the given iterator supports the __len__ method,
    the returned batch iterator will, too.
    """
    if hasattr(it, '__len__'):
        return _iter_batches_with_len(it, batch_size)
    else:
        return _iter_batches(it, batch_size)


class _iter_batches:
    def __init__(self, it, batch_size):
        self.base_iterator = it
        self.batch_size = batch_size


    def __iter__(self):
        return self._iter_batches(self.base_iterator, self.batch_size)


    def _iter_batches(self, it, batch_size):
        if isinstance(it, (pd.DataFrame, pd.Series)):
            for batch_start in range(0, len(it), batch_size):
                yield it.iloc[batch_start:batch_start+batch_size]
            return
        elif isinstance(it, (list, np.ndarray)):
            for batch_start in range(0, len(it), batch_size):
                yield it[batch_start:batch_start+batch_size]
            return
        else:
            if not isinstance(it, Iterator):
                assert isinstance(it, Iterable)
                it = iter(it)

            while True:
                batch = []
                try:
                    for _ in range(batch_size):
                        batch.append(next(it))
                except StopIteration:
                    return
                finally:
                    if batch:
                        yield batch


class _iter_batches_with_len(_iter_batches):
    def __len__(self):
        return int(np.ceil(len(self.base_iterator) / self.batch_size))


def compile_columns(client, core_columns=[]):
    """
    Compile list of columns from available :Neuron keys (excluding ROIs).

    Args:
        client:
            neu.Client to collect columns for.
        core_columns:
            List of core columns (optional). If provided, new columns will be
            added to the end of the list and non-existing columns will be
            dropped.

    Returns:
        columns:
            List of key names.
    """
    # Fetch existing keys. This call is cached.
    keys = client.fetch_neuron_keys()

    # Drop ROIs
    keys = [k for k in keys if k not in client.all_rois]

    # Drop missing columns from core_columns
    columns = [k for k in core_columns if k in keys]

    # Add new keys (sort to make deterministic)
    columns += [k for k in sorted(keys) if k not in columns]

    return columns
