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

# In the current, first-price auction-based paradigm, users are presented with three wallet defaults to specify their gas price bids (typically, "slow", "medium" or "fast"), while experienced users can set their own price with advanced settings.
# 
# Since [EIP 1559 equilibrates to a market price](https://ethereum.github.io/abm1559/notebooks/stationary1559.html), we look into defaults that embody the price-taking behaviour of users under 1559. Their wallets presents them with a "current transaction price", and a simple **binary option**: transact or not.
# 
# This default is likely sufficient for periods where the basefee is relatively stable, but fails to accurately represent user demand when that demand is shifting upwards, since [strategic users](https://ethereum.github.io/abm1559/notebooks/strategicUser.html) enter bidding wars to get ahead. In this case, a wallet could revert to the well-known UX presenting three price levels, with differentiated guarantees for inclusion and delay, which we'll call **expert mode**.
# 
# Finding the correct price levels to offer as well as the criteria determining when the UI should switch from the "price-taking" UI to the "shifting demand" UI is discussed via simulations of various demand behaviours. We first import classes from our agent-based transaction market simulation.

# In[1]:


import os, sys
sys.path.insert(1, os.path.realpath(os.path.pardir))
from typing import Sequence

from abm1559.utils import constants

# Target at most 100 transactions per block
gamma = 250000
constants["SIMPLE_TRANSACTION_GAS"] = gamma

from abm1559.txpool import TxPool

from abm1559.users import User1559

from abm1559.userpool import UserPool

from abm1559.chain import (
    Chain,
    Block1559,
)

from abm1559.simulator import (
    spawn_poisson_heterogeneous_demand,
    update_basefee,
)

from abm1559.txs import Tx1559
from abm1559.config import rng

import pandas as pd
import numpy as np
import seaborn as sns
from tqdm.notebook import tqdm
import plotly.express as px
import plotly.io as pio
pd.options.plotting.backend = "plotly"
pio.renderers.default = "plotly_mimetype+notebook_connected"


# ## Simulation components
# 
# ### Wallet oracle
# 
# We introduce a new element sitting between the user and the market, a `WalletOracle` providing the user with the two UIs discussed above. The wallet features two major parameters:
# 
# - A criterion deciding when to switch from the binary UI to the expert mode. The criterion takes the form of a parameterised expression:
# 
# > Switch to the expert mode whenever the current basefee is $r$% above its moving average of the last $B$ blocks.
# 
# - The price levels to suggest the user in expert mode. Although we settle on _three_ price levels, mimicking the current wallets UI, it may be possible to consider two or four or any other number we think reasonable. For each price level, the wallet suggests a different gas premium, following the rules:
# 
# >  - The max priority fee $g_i$ of price level $i$ is $c_i$ times the miner minimum fee $\delta$, so $g_i = c_i \cdot \delta$ with $c_i < c_j, \, \forall i < j$.
# > - The max fee $m_i$ of price level $i$ is $c_i$ times the current basefee.
# 
# In the simple setting, we'll first look at $c_1 = 1$ (slow users send the smallest acceptable priority fee), $c_2 = 2$ and $c_3 = 3$. Later on, we'll investigate more complex rules:
#   
# > The gas premium $g_i$ of price level $i$ is a function of the recorded premiums in the $B$ last blocks, $g_i = f_i(B_{t-B}, \dots, B_{t-1})$, with $g_i < g_j, \, \forall i < j$.

# In[2]:


class WalletOracle:
    
    def __init__(self, sensitivity=0.05, ma_window=5):
        self.sensitivity = sensitivity
        self.ma_window = ma_window
        self.basefees = []
        self.ui_mode = "posted"
        self.expert_defaults = {
            key: { "max_fee": 0, "gas_premium": 0 } for key in ["slow", "medium", "fast"]
        }
        self.min_premium = 1 * (10 ** 9) # in wei
        
    def update_defaults(self, basefee):
        self.basefees = self.basefees[-(self.ma_window - 1):] + [basefee]
        ma_basefee = np.mean(self.basefees)
        
        if (1 + self.sensitivity) * ma_basefee >= basefee:
            # basefee broadly in line with or under its average value in the last blocks
            self.ui_mode = "posted"
            
        else:
            self.ui_mode = "expert"
            self.expert_defaults = {
                "slow": { "max_fee": 2*basefee, "gas_premium": self.min_premium },
                "medium": { "max_fee": 3*basefee, "gas_premium": 2*self.min_premium },
                "fast": { "max_fee": 4*basefee, "gas_premium": 3*self.min_premium },
            }
            
    def get_expert_defaults(self, max_fee):
        return {
            key: {
                "max_fee": min(max_fee, value["max_fee"]),
                "gas_premium": value["gas_premium"],
            } for key, value in self.expert_defaults.items()
        }
            
    def get_defaults(self, speed, max_fee):
        if self.ui_mode == "posted":
            return (self.ui_mode, {
                "max_fee": min([max_fee, max([1.5 * self.basefees[-1], self.basefees[-1] + self.min_premium])]),
                "gas_premium": self.min_premium
            })
        else:
            return (self.ui_mode, self.get_expert_defaults(max_fee)[speed])


# ### User behaviour with wallets

# We define a new `User` class who uses the wallet defined above to make their pricing decision. This user is an `AffineUser` which experiences a cost for waiting, drawn from some distribution. We assume that users know their quantile in this distribution. For instance, if we draw costs from a uniform distribution between 0 and 1 Gwei per gas unit, then a user with cost 0.85 knows they are in the top 15% of the most impatient users.
# 
# Users who care a lot about fast inclusion tend to choose the fast default. We'll assume that since users know their quantile, they can also identify whether the slow, medium, or fast default is appropriate for them. Assuming the wallet suggests three price levels, we have two numbers $0 < q_1 < q_2 < 0$ such that:
# 
# - Users below quantile $q_1$ choose the slow default.
# - Users between quantiles $q_1$ and $q_2$ choose the medium default.
# - Users above quantile $q_2$ choose the fast default.
# 
# We'll leave these two numbers $q_1, q_2$ as parameters for the simulation, arbitrarily picking $q_1 = 0.4$ and $q_2 = 0.75$.

# In[3]:


class UserWallet(User1559):
    
    def __init__(self, wakeup_block, **kwargs):
        super().__init__(wakeup_block, **kwargs)
        self.value = self.rng.uniform(low=0, high=20) * (10 ** 9)
        self.cost_per_unit = self.rng.uniform(low=0, high=1) * (10 ** 9)
    
    def expected_time(self, params):
        return 0
    
    def decide_parameters(self, env):
        wallet = env["wallet"]
        
        if self.cost_per_unit <= env["cost_quantiles"][0]:
            speed = "slow"
        elif self.cost_per_unit <= env["cost_quantiles"][1]:
            speed = "medium"
        else:
            speed = "fast"
        
        defaults = wallet.get_defaults(speed, self.value)[1]
        
        return {
            **defaults,
            "start_block": self.wakeup_block,
        }


# ### Transaction pool
# 
# We also define the following transaction pool behaviour:
# 
# - The maximum number of pending transactions in the pool is limited to 500. With our setting of at most 100 transactions in the block (and a target of 50 transactions per block), the pool holds up to 5 full blocks worth of transactions.
# - Miners select the highest tip-paying transactions from the pool, where `tip = min(premium, fee_cap - basefee)`, as long as `tip >= MIN_PREMIUM`, with `MIN_PREMIUM` set to 1 Gwei.
# - Whenever too many transactions are added to the pool, the pool is resorted by `tip`, with the lowest tip-paying transactions excluded until the pool limit size is recovered.

# In[4]:


MAX_TX_POOL = 500
MIN_PREMIUM = 1e9

class TxPoolLimit(TxPool):
    
    def __init__(self, max_txs=MAX_TX_POOL, min_premium=MIN_PREMIUM, **kwargs):
        super().__init__(**kwargs)
        self.max_txs = max_txs
        self.min_premium = MIN_PREMIUM
    
    def add_txs(self, txs: Sequence[Tx1559], env: dict) -> Sequence[Tx1559]:
        for tx in txs:
            self.txs[tx.tx_hash] = tx
                    
        if self.pool_length() > self.max_txs:
            sorted_txs = sorted(self.txs.values(), key = lambda tx: -tx.tip(env))
            self.empty_pool()
            self.add_txs(sorted_txs[0:self.max_txs], env)
            return sorted_txs[self.max_txs:]
        
        return []

class TxPool1559(TxPoolLimit):
    
    def select_transactions(self, env, user_pool=None, rng=rng) -> Sequence[Tx1559]:
        # Miner side
        max_tx_in_block = int(constants["MAX_GAS_EIP1559"] / constants["SIMPLE_TRANSACTION_GAS"])

        valid_txs = [tx for tx in self.txs.values() if tx.is_valid(env) and tx.tip(env) >= self.min_premium]
        rng.shuffle(valid_txs)

        sorted_valid_demand = sorted(
            valid_txs,
            key = lambda tx: -tx.tip(env)
        )
        selected_txs = sorted_valid_demand[0:max_tx_in_block]

        return selected_txs


# ### Running the simulation
# 
# Open the code snippet below to see the simulation steps. Every round, we spawn a certain quantity of users, some of them are discouraged by the price level and balk, while others send their transactions in. The transaction pool receives the transactions and the miner includes the highest tip-paying transactions first, breaking ties arbitrarily.

# In[5]:


def simulate(demand_scenario, shares_scenario, rng=np.random.default_rng(1)):
    # Instantiate a couple of things
    txpool = TxPool1559()
    basefee = constants["INITIAL_BASEFEE"]
    chain = Chain()
    metrics = []
    user_pool = UserPool()
    wallet = WalletOracle()

    for t in tqdm(range(len(demand_scenario)), desc="simulation loop", leave=False):
        
        # Update oracle
        wallet.update_defaults(basefee)
        
        # We return some demand which on expectation yields demand_scenario[t] new users per round
        users = spawn_poisson_heterogeneous_demand(t, demand_scenario[t], shares_scenario[t], rng=rng)
        cost_quantiles = np.quantile([u.cost_per_unit for u in users], q = [0.4, 0.75])
        
        # `params` are the "environment" of the simulation
        env = {
            "basefee": basefee,
            "current_block": t,
            "wallet": wallet,
            "cost_quantiles": cost_quantiles,
        }
        
        # Add users to the pool and check who wants to transact
        # We query each new user with the current basefee value
        # Users either return a transaction or None if they prefer to balk
        decided_txs = user_pool.decide_transactions(users, env)

        # New transactions are added to the transaction pool
        txpool.add_txs(decided_txs, env)

        # The best valid transactions are taken out of the pool for inclusion
        selected_txs = txpool.select_transactions(env)
        txpool.remove_txs([tx.tx_hash for tx in selected_txs])

        # We create a block with these transactions
        block = Block1559(txs = selected_txs, parent_hash = chain.current_head, height = t, basefee = basefee)
        
        # The block is added to the chain
        chain.add_block(block)

        row_metrics = {
            "block": t,
            "basefee": basefee / (10 ** 9),
            "ui_mode": wallet.ui_mode,
            "users": len(users),
            "decided_txs": len(decided_txs),
            "included_txs": len(selected_txs),
            "blk_min_premium": block.min_premium() / (10 ** 9), # to Gwei
            "blk_avg_gas_price": block.average_gas_price(),
            "blk_avg_tip": block.average_tip(),
            "pool_length": txpool.pool_length(),
        }
        metrics.append(row_metrics)

        # Finally, basefee is updated and a new round starts
        basefee = update_basefee(block, basefee)

    return (pd.DataFrame(metrics), user_pool, chain)


# In[6]:


blocks = 100
demand_scenario = [250 for i in range(blocks)]

shares_scenario = [{
    UserWallet: 1
} for i in range(blocks)]

(df, user_pool, chain) = simulate(demand_scenario, shares_scenario)


# ## Easy setting: Fixed demand per step, simple wallet rules
# 
# We start with a setting where a constant number of users spawns between every two blocks, and all users send their transactions via the `WalletOracle`.
# 
# Since basefee is too low at the start, we see the `ui_mode` switching to `expert` mode from block 1 onwards, picking up a demand shift important enough to allow users to choose their price level. Users are now presented with the three options and choose their preferred one.
# 
# By block 25, when the basefee stabilises, the UI switches back to the `posted` price binary option.

# In[7]:


df[(df.block <= 30)]


# In[8]:


fig = df.plot("block", ["basefee", "blk_avg_tip"])
fig.update_layout(
    title = "Basefee and average tip per block",
    xaxis_title = "Block",
    yaxis_title = "Value (Gwei)"
)
fig


# The plot shows clearly that users are bidding above the minimum miner fee (1 Gwei) at the beginning, while basefee catches up to the congestion level.

# In[9]:


# Obtain the pool of users (all users spawned by the simulation)
user_pool_df = user_pool.export().rename(columns={ "pub_key": "sender" })

# Export the trace of the chain, all transactions included in blocks
chain_df = chain.export()

# Join the two to associate transactions with their senders
user_txs_df = chain_df.join(user_pool_df.set_index("sender"), on="sender")


# In[10]:


bins = 10
user_txs_df["user_sw"] = user_txs_df.apply(
    lambda row: row.value - (row.block_height - row.start_block) * row.cost_per_unit - (row.basefee + row.tip / (10 ** 9)),
    axis=1
)
user_txs_df["value_bin"] = pd.cut(user_txs_df["value"], bins=np.linspace(0, 20, bins+1))
user_txs_df["cost_bin"] = pd.cut(user_txs_df["cost_per_unit"], bins=np.linspace(0, 1, bins+1))
user_txs_df["user_impatience"] = user_txs_df.cost_per_unit.apply(
    lambda c: "slow" if c < 0.4 else "medium" if c < 0.75 else "fast"
)
grouped = user_txs_df.groupby(["block_height", "value_bin", "cost_bin"]).mean().reset_index()


# ### Plots
# 
# We call _utility_ a measure of user satisfaction. Users have a value for their transaction (in Gwei per gas unit) and a cost for waiting (in Gwei per gas unit per block waited for). When a user with value $v$ and cost for waiting $c$ waits for $t$ blocks and pays $f$, their utility is given by $v - ct - f$.
# 
# In the plot below, we show the utility obtained by users with various value/cost per unit combinations at block 2, with the utility depicted by the colour of the tile: the more dark red, the higher the utility. Higher value users have higher utility, since they are getting a good deal for inclusion! Basefee is low so the total fee is low too, and they don't have to wait for long.
# 
# But note that users with cost for waiting lower than 0.4 are totally absent from the plot. These users choose the "slow" option, and so are outbid by more impatient users picking the "medium" and "fast" options.

# In[12]:


gp = grouped[grouped.block_height == 2].copy()
gp['cost_bin'] = gp['cost_bin'].astype('str')
gp['value_bin'] = gp['value_bin'].astype('str')
gp = gp.pivot("value_bin", "cost_bin", "user_sw").reindex(
    np.flip([str(pd.Interval(x, x+2, closed="right")) for x in np.linspace(0, 18, bins)])
)
fig = px.imshow(gp)
fig.update_layout(
    title = "Utility received by users included at block 2",
    xaxis_title = "Cost per block waited (Gwei per gas unit)",
    yaxis_title = "Value (Gwei per gas unit)"
)
fig


# Over time, low value users become priced out. Observe too that between block 1 and 14, users with low cost per unit of time are also priced out: given the high congestion, users with high cost per unit of time who use the "fast" wallet default are included in priority.
# 
# After block 14, once basefee is high enough, all higher value users, regardless of their cost per time unit, have a chance at inclusion. It doesn't matter if they are impatient or not. Since basefee prices out everyone except for a demand equal to the block target size, as long as your value is high enough, you'll be included in the next block.

# In[13]:


def draw_heatmap(*args, **kwargs):
    data = kwargs.pop('data')
    d = data.pivot(index=args[1], columns=args[0], values=args[2])
    ax = sns.heatmap(d, **kwargs, vmin=0, vmax=20, cmap="plasma", center=5, cbar=False)
    ax.invert_yaxis()

grid_size = 6
fg = sns.FacetGrid(grouped[grouped.block_height < grid_size ** 2], col='block_height', col_wrap = grid_size)
fg.map_dataframe(draw_heatmap, 'cost_bin', 'value_bin', 'user_sw')


# Here is another way to look at it. Each round, we are spawning:
# 
# - 40% of very patient users, who choose the "slow" option in expert mode
# - 35% of medium patient users, who choose the "medium" option in expert mode
# - 25% of impatient users, who choose the "fast" option in expert mode
# 
# If these users all had as much chances of being included, we'd expect to see them in roughly the proportions described above in each block. Yet as the next plot shows us this is not the case at the very beginning.

# In[14]:


import matplotlib.pyplot as plt
counts_df = user_txs_df[["block_height", "user_impatience", "basefee"]].groupby(
    ["block_height", "user_impatience"]
).count().unstack().fillna(0)
counts_df["total"] = sum([counts_df.iloc[:,i] for i in range(3)])
for i in range(3):
    counts_df.iloc[:,i] = counts_df.iloc[:,i] / counts_df["total"] * 100
expected_percent = [25, 35, 40]
for i in range(3):
    counts_df.iloc[:,i] = counts_df.iloc[:,i] - expected_percent[i]
pl = counts_df.iloc[:,0:3].rolling(window=4).mean()
pl.columns = pl.columns.droplevel(0)

fig = px.line(pl)
fig.update_layout(
    title = "Inclusion difference relative to expectation",
    xaxis_title = "Block",
    yaxis_title = "Inclusion difference"
)
fig


# We smooth the plot slightly for readability.
# 
# At the start, fast users dominate, unsurprisingly. As basefee rises and more and more users are priced out (including some fast users who have a low value for their transaction, although they are very impatient), there is room for medium patient users to be included. This keeps going until basefee reaches around its equilibrium point and all users are included in their expected proportions.
Wallet defaults for EIP 1559// References + footnotes

// Authors
let authorData = ["barnabe"];