import numpy as np

import streamlit as st
import collections
import functools
import inspect
import textwrap
from streamlit.report_thread import get_report_ctx
from streamlit.server.server import Server


def cache_on_button_press(label, **cache_kwargs):
    """Function decorator to memoize function executions.
    Parameters
    ----------
    label : str
        The label for the button to display prior to running the cached funnction.
    cache_kwargs : Dict[Any, Any]
        Additional parameters (such as show_spinner) to pass into the underlying @st.cache decorator.
    Example
    -------
    This show how you could write a username/password tester:
    >>> @cache_on_button_press('Authenticate')
    ... def authenticate(username, password):
    ...     return username == "buddha" and password == "s4msara"
    ...
    ... username = st.text_input('username')
    ... password = st.text_input('password')
    ...
    ... if authenticate(username, password):
    ...     st.success('Logged in.')
    ... else:
    ...     st.error('Incorrect username or password')
    """
    internal_cache_kwargs = dict(cache_kwargs)
    internal_cache_kwargs["allow_output_mutation"] = True
    internal_cache_kwargs["show_spinner"] = False

    def function_decorator(func):
        @functools.wraps(func)
        def wrapped_func(*args, **kwargs):
            @st.cache(**internal_cache_kwargs)
            def get_cache_entry(func, args, kwargs):
                class ButtonCacheEntry:
                    def __init__(self):
                        self.evaluated = False
                        self.return_value = None

                    def evaluate(self):
                        self.evaluated = True
                        self.return_value = func(*args, **kwargs)

                return ButtonCacheEntry()

            cache_entry = get_cache_entry(func, args, kwargs)
            if not cache_entry.evaluated:
                if st.sidebar.button(label):
                    cache_entry.evaluate()
                else:
                    raise st.ScriptRunner.StopException
            return cache_entry.return_value

        return wrapped_func

    return function_decorator


"""Hack to add per-session state to Streamlit.
Usage
-----
>>> import SessionState
>>>
>>> session_state = SessionState.get(user_name='', favorite_color='black')
>>> session_state.user_name
''
>>> session_state.user_name = 'Mary'
>>> session_state.favorite_color
'black'
Since you set user_name above, next time your script runs this will be the
result:
>>> session_state = get(user_name='', favorite_color='black')
>>> session_state.user_name
'Mary'
"""


class SessionState(object):
    def __init__(self, **kwargs):
        """A new SessionState object.
        Parameters
        ----------
        **kwargs : any
            Default values for the session state.
        Example
        -------
        >>> session_state = SessionState(user_name='', favorite_color='black')
        >>> session_state.user_name = 'Mary'
        ''
        >>> session_state.favorite_color
        'black'
        """
        for key, val in kwargs.items():
            setattr(self, key, val)


def get(**kwargs):
    """Gets a SessionState object for the current session.
    Creates a new object if necessary.
    Parameters
    ----------
    **kwargs : any
        Default values you want to add to the session state, if we're creating a
        new one.
    Example
    -------
    >>> session_state = get(user_name='', favorite_color='black')
    >>> session_state.user_name
    ''
    >>> session_state.user_name = 'Mary'
    >>> session_state.favorite_color
    'black'
    Since you set user_name above, next time your script runs this will be the
    result:
    >>> session_state = get(user_name='', favorite_color='black')
    >>> session_state.user_name
    'Mary'
    """
    # Hack to get the session object from Streamlit.

    ctx = get_report_ctx()

    this_session = None

    session_info = Server.get_current()._get_session_info(ctx.session_id)
    this_session = session_info.session

    if this_session is None:
        raise RuntimeError(
            "Oh noes. Couldn't get your Streamlit Session object" "Are you doing something fancy with threads?"
        )

    # Got the session object! Now let's attach some state into it.

    if not hasattr(this_session, "_custom_session_state"):
        this_session._custom_session_state = SessionState(**kwargs)

    return this_session._custom_session_state


def simple_transformers_model(model):
    return (type(model).__name__, model.args)


def get_color(i):
    # Colors taken from Sasha Trubetskoy's list of colors - https://sashamaps.net/docs/tools/20-colors/
    colors = [
        (60, 180, 75, 0.4),
        (255, 225, 25, 0.4),
        (0, 130, 200, 0.4),
        (245, 130, 48, 0.4),
        (145, 30, 180, 0.4),
        (70, 240, 240, 0.4),
        (240, 50, 230, 0.4),
        (210, 245, 60, 0.4),
        (250, 190, 212, 0.4),
        (0, 128, 128, 0.4),
        (220, 190, 255, 0.4),
        (170, 110, 40, 0.4),
        (255, 250, 200, 0.4),
        (128, 0, 0, 0.4),
        (170, 255, 195, 0.4),
        (128, 128, 0, 0.4),
        (255, 215, 180, 0.4),
        (0, 0, 128, 0.4),
        (128, 128, 128, 0.4),
        (255, 255, 255, 0.4),
        (0, 0, 0, 0.4),
        (230, 25, 75, 0.4),
    ]
    try:
        return str(colors[i])
    except IndexError:
        return str(tuple(np.random.rand(3,).tolist() + 0.7))
