import matplotlib.pylab as plt
import importlib
import pandas as pd
import numpy as np
from typing import Callable, Iterator
import inspect
from types import ModuleType


class Required:
    def __repr__(self):
        return '_REQUIRED_'


_REQUIRED_ = Required()

callable_filt_for = {
    'callable': callable,
    'class': inspect.isclass,
    'function': inspect.isfunction,
    'function_or_class': lambda x: inspect.isclass(x) or inspect.isfunction(x)
}

DFLT_CALLABLE_FILT = 'callable'  # other choices: inspect.isfunction, inspect.isclass


def get_callable_filt(callable_filt):
    """
    Returns a boolean function that will be used to filter objects.
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :return:
    """
    if isinstance(callable_filt, str):
        callable_filt = callable_filt_for.get(callable_filt, None)
        if callable_filt is None:
            raise ValueError(f"No such callable_filt: {callable_filt}")
    return callable_filt


def module_classes(module):
    return filter(inspect.isclass, module.__dict__.values())


def callables_of_module(module, callable_filt=DFLT_CALLABLE_FILT):
    """

    :param module:
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :return:
    """
    callable_filt = get_callable_filt(callable_filt)
    if isinstance(module, str):
        module = importlib.import_module(module)
    return (x[1] for x in inspect.getmembers(module, predicate=callable_filt))


def callables_of_class(cls, callable_filt='function'):
    """

    :param cls:
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :return:
    """
    callable_filt = get_callable_filt(callable_filt)
    return (x[1] for x in inspect.getmembers(cls, predicate=callable_filt))


def get_callables_from(callables, callable_filt=DFLT_CALLABLE_FILT):
    """
    Get a generator of callable objects from the input x.
    :param callables: A module (or module dot-path string), class, callable, or list thereof
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :return: A generator of callable objects
    """
    callable_filt = get_callable_filt(callable_filt)
    if callable(callables):
        if not inspect.isclass(callables):
            yield from [callables]
        else:
            yield from get_callables_from(callables_of_class(callables), callable_filt=callable_filt)
    else:
        if isinstance(callables, str) or isinstance(callables, ModuleType):
            yield from get_callables_from(callables_of_module(callables), callable_filt=callable_filt)
        elif inspect.isclass(callables):
            yield callables
            yield from get_callables_from(callables_of_class(callables), callable_filt=callable_filt)
        else:
            try:
                for c in callables:
                    if not hasattr(c, '__qualname__'):
                        continue  # hack to avoid a case I don't understand
                    yield from get_callables_from(c, callable_filt=callable_filt)
            except:
                raise TypeError(f"Don't know how to handle the input: {c}\n"
                                "Wasn't a callable, string, module, or iterable of such.")


class SignatureExtractor:
    def __init__(self, normalize_var_names=True):
        def param_mint(param):
            if normalize_var_names:
                if param.kind == inspect.Parameter.VAR_KEYWORD:
                    d = {'name': '**kwargs'}
                elif param.kind == inspect.Parameter.VAR_POSITIONAL:
                    d = {'name': '*args'}
                else:
                    d = {'name': param.name}
            else:
                d = {'name': param.name}

            if param.default == inspect._empty:
                d['default'] = _REQUIRED_
            else:
                d['default'] = param.default

            return d

        self.param_mint = param_mint

    def __call__(self, obj):
        return [self.param_mint(p) for p in inspect.signature(obj).parameters.values()]


# Note: A function would make more sense, but chose to make a class to demo how classes are handled as callables
extract_name_and_default = SignatureExtractor(normalize_var_names=True)


def name_arg_default_dict_of_callables(callables: Iterator[Callable]) -> dict:
    """
    Get an {callable_name: {arg_name: arg_default, ...}, ...} dict from a collection of callables.
    See also: name_arg_default_dict_of_callables and arg_default_dict_of_module_classes
    :param callables: Iterable of callables
    :return: A dict
    """
    d = dict()
    for obj in callables:
        try:
            d[obj.__qualname__] = {x['name']: x['default'] for x in extract_name_and_default(obj)}
        except Exception as e:
            pass  # TODO: Give choice to warn instead of ignore
    return d


def arg_default_dict_of_callables(callables, callable_filt=DFLT_CALLABLE_FILT) -> dict:
    """
    Get an {callable_name: {arg_name: arg_default, ...}, ...} dict from a collection of callables taken from
    a specification of callables.
    :param callables = get_callables_from(callables, callable_filt=callable_filt)
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :return: A dict
    """
    callable_filt = get_callable_filt(callable_filt)
    return name_arg_default_dict_of_callables(get_callables_from(callables, callable_filt=callable_filt))


def non_null_counts(df: pd.DataFrame, null_val=np.nan):
    if null_val is np.nan:
        non_null_lidx = ~df.isna()
    else:
        non_null_lidx = df != null_val
    row_null_zero_count = non_null_lidx.sum(axis=1)
    col_null_zero_count = non_null_lidx.sum(axis=0)
    return row_null_zero_count, col_null_zero_count


def _df_of_callable_arg_default_dict(callable_arg_default_dict, null_fill='') -> pd.DataFrame:
    """
    Get a dataframe from a callable_arg_default_dict
    :param module:
    :param null_fill:
    :return:
    """
    d = pd.DataFrame.from_dict(callable_arg_default_dict)
    row_null_zero_count, col_null_zero_count = non_null_counts(d, null_val=np.nan)
    row_argsort = np.argsort(row_null_zero_count)[::-1]
    col_argsort = np.argsort(col_null_zero_count)[::-1]
    return d.iloc[row_argsort, col_argsort].fillna(null_fill).T


def callables_signatures_df(
        callables, callable_filt=DFLT_CALLABLE_FILT, null_fill='') -> pd.DataFrame:
    """
    Get a dataframe representing the signatures of the input callables.
    :param callables = get_callables_from(callables, callable_filt=callable_filt)
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :param null_fill: What to fill the empty cells with
    :return: A dataframe
    """
    callable_filt = get_callable_filt(callable_filt)
    d = arg_default_dict_of_callables(callables, callable_filt=callable_filt)
    return _df_of_callable_arg_default_dict(d, null_fill=null_fill)


def heatmap(X, y=None, col_labels=None, figsize=None, cmap=None, return_gcf=False, ax=None,
            xlabel_top=True, ylabel_left=True, xlabel_bottom=True, ylabel_right=True, **kwargs):
    """
    Make a heatmap of a matrix or pandas.DataFrame, but let the function figure stuff out.
    """
    import pandas as pd
    import numpy as np
    n_items, n_cols = X.shape
    if col_labels is not None:
        if col_labels is not False:
            assert len(col_labels) == n_cols, \
                "col_labels length should be the same as the number of columns in the matrix"
    elif isinstance(X, pd.DataFrame):
        col_labels = list(X.columns)

    if figsize is None:
        x_size, y_size = X.shape
        if x_size >= y_size:
            figsize = (6, min(18, 6 * x_size / y_size))
        else:
            figsize = (min(18, 6 * y_size / x_size), 6)

    if cmap is None:
        if X.min(axis=0).min(axis=0) < 0:
            cmap = 'RdBu_r'
        else:
            cmap = 'hot_r'

    kwargs['cmap'] = cmap
    kwargs = dict(kwargs, interpolation='nearest', aspect='auto')

    if figsize is not False:
        plt.figure(figsize=figsize)

    if ax is None:
        plt.imshow(X, **kwargs)
    else:
        ax.imshow(X, **kwargs)
    plt.grid(None)

    if y is not None:
        y = np.array(y)
        assert all(sorted(y) == y), "This will only work if your row_labels are sorted"

        unik_ys, unik_ys_idx = np.unique(y, return_index=True)
        for u, i in zip(unik_ys, unik_ys_idx):
            plt.hlines(i - 0.5, 0 - 0.5, n_cols - 0.5, colors='b', linestyles='dotted', alpha=0.5)
        plt.hlines(n_items - 0.5, 0 - 0.5, n_cols - 0.5, colors='b', linestyles='dotted', alpha=0.5)
        plt.yticks(unik_ys_idx + np.diff(np.hstack((unik_ys_idx, n_items))) / 2, unik_ys)
    elif isinstance(X, pd.DataFrame):
        y_tick_labels = list(X.index)
        plt.yticks(list(range(len(y_tick_labels))), y_tick_labels);

    if col_labels is not None:
        plt.xticks(list(range(len(col_labels))), col_labels)
    else:
        plt.xticks([])

    plt.gca().xaxis.set_tick_params(labeltop=xlabel_top, labelbottom=xlabel_bottom)
    plt.gca().yaxis.set_tick_params(labelleft=ylabel_left, labelright=ylabel_right)

    if return_gcf:
        return plt.gcf()


def heatmap_of_signatures(callables, callable_filt=DFLT_CALLABLE_FILT, figsize=None, cmap='gray_r'):
    """
    Visualize a matrix containing the all functions of the module, and their arguments.
    :param callables: A module (or module dot-path string), class, callable, or list thereof
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :param figsize: Size of the figure (will try to figure one out if not specified - pun noticed)
    :return:
    """
    callable_filt = get_callable_filt(callable_filt)
    callables = get_callables_from(callables, callable_filt=callable_filt)
    core_comp_df = callables_signatures_df(callables, callable_filt=callable_filt)
    t = (core_comp_df != '').astype(int) + (core_comp_df == _REQUIRED_).astype(int)
    heatmap(t, figsize=figsize, cmap=cmap)
    plt.xticks(rotation=90)
    plt.grid(False)
    plt.show()


def plot_nonnull_counts_of_signatures(
        callables, callable_filt=DFLT_CALLABLE_FILT, n_top_items=50, figsize=None, hspace=0.5):
    """
    For all callables of the input module, will plot (as a bar graph):
        The argument count of each callable
        The callable count of each argument (i.e., in how many callables are the given argument used)
    :param callables: A module (or module dot-path string), class, callable, or list thereof
    :param callable_filt: A boolean function or the strings 'callable', 'class', 'function', or 'function_or_class'
    :param n_top_items: Number of (top) items to plot
    :param figsize: Size of the figure (will try to figure one out if not specified - pun noticed)
    :param hspace: The space between both bar plots.
    :return:
    """
    callable_filt = get_callable_filt(callable_filt)
    callables = get_callables_from(callables, callable_filt=callable_filt)
    core_comp_df = callables_signatures_df(callables, callable_filt=callable_filt)
    t, tt = non_null_counts(core_comp_df, null_val='')
    plt.figure(figsize=figsize)
    plt.subplot(2, 1, 1)
    t.iloc[:n_top_items].plot(kind='bar')
    plt.subplot(2, 1, 2)
    tt.iloc[:n_top_items].plot(kind='bar')
    plt.subplots_adjust(hspace=hspace)
    plt.show()


if __name__ == '__main__':
    import argh

    argh.dispatch_commands([callables_signatures_df,
                            heatmap_of_signatures,
                            plot_nonnull_counts_of_signatures])