#!/usr/bin/env python
# coding: utf-8

# In[45]:


from jupyterhub import orm

def init_db():
    db = orm.new_session_factory("sqlite:///:memory:")()

    user = orm.User(name="test")
    db.add(user)

    db.add(orm.OAuthClient(identifier="jupyterhub"))
    db.commit()
    return db

def add_tokens(db, n, all_tokens=None):
    user = db.query(orm.User).one()
    all_tokens = all_tokens or []
    for i in range(n):
        all_tokens.append(user.new_api_token())
    return all_tokens
    
db = init_db()
all_tokens = add_tokens(db, 2);


# In[46]:


db.get_bind().echo = True


# In[47]:


orm.APIToken.find_prefix = orm.APIToken.find_prefix_startswith
orm.APIToken.find(db, all_tokens[0])


# In[48]:


orm.APIToken.find_prefix = orm.APIToken.find_prefix_equal
orm.APIToken.find(db, all_tokens[0])


# In[49]:


db.get_bind().echo = False


# In[50]:


import random


# In[51]:


get_ipython().run_line_magic('timeit', 'random.choice(all_tokens)')


# In[52]:


no_match = 'x' * len(all_tokens[0])


# In[74]:


start = time.perf_counter()
deadline = start + 60
records = []

db = init_db()
all_tokens = []


def add_rows(tr, match, found):
    n = len(all_tokens)
    for run in tr.all_runs:
        per_loop = run / tr.loops
        records.append((n, per_loop, match, found, tr.loops))


while time.perf_counter() < deadline:
    all_tokens = add_tokens(db, len(all_tokens) or 100, all_tokens)
    print(len(all_tokens))
    for match in ('equal', 'startswith'):
        orm.APIToken.find_prefix = getattr(orm.APIToken, f"find_prefix_{match}")
        for found in (True, False):
            if found:
                tr = get_ipython().run_line_magic('timeit', '-o -n 500 orm.APIToken.find(db, random.choice(all_tokens))')
            else:
                tr = get_ipython().run_line_magic('timeit', '-o -n 500 orm.APIToken.find(db, no_match)')
            add_rows(tr, match, found)


# In[75]:


import pandas as pd
df = pd.DataFrame(records, columns=["n", "t", "match", "found", "loops"])
df["


# In[148]:


import altair as alt
df["ms"] = df.t * 1e3
chart = (
    alt.Chart(df)
    .mark_point()
    .encode(
        x=alt.X(
            "n",
            axis=alt.Axis(title="number of tokens"),
        ),
        y=alt.Y(
            "ms",
            scale=alt.Scale(type="log"),
            axis=alt.Axis(title="lookup time (ms)"),
        ),
        color="match",
        shape="found",
        opacity=alt.value(0.25),
    )
)
chart


# In[146]:


(
    chart
    + chart.transform_regression(
        "n",
        "t",
        groupby=["match"],
        method="exp",
    ).mark_line()
)