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

# In[1]:


get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'svg'")

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

from typing import Sequence, Dict, Tuple

from abm1559.utils import (
    constants,
)

gamma = 76000 # obtained from stats analysis of sample txs
constants["SIMPLE_TRANSACTION_GAS"] = gamma

from abm1559.txpool import TxPool

from abm1559.users import (
    User1559,
    User
)

from abm1559.config import rng

from abm1559.txs import Tx1559

from abm1559.userpool import UserPool

from abm1559.chain import (
    Chain,
    Block1559,
)

from abm1559.simulator import (
    spawn_poisson_heterogeneous_demand,
    update_basefee,
    generate_gbm,
    apply_block_time_variance,
    generate_poisson_process,
    generate_abm,
    generate_jump_process,
)

import matplotlib.pyplot as plt
import pandas as pd
pd.set_option('display.max_rows', 1000)
import numpy as np
import time
from tqdm.notebook import tqdm
from itertools import product


# In[2]:


MAX_TX_POOL = 4096
MIN_PREMIUM = 1e9

class TxPool1559(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 []
    
    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
    
class TxPoolTrendPicker(TxPool1559):
    
    # window = 1 <=> Fixed band policy
    
    def __init__(self, band_width=0.1, window=1, **kwargs):
        super().__init__(**kwargs)
        self.band_width = band_width
        self.window = window
        
    def should_evict(self, tx: Tx1559, basefee: int, delta: float) -> bool:
        # band_width = 1 <=> never evict
        return (1+delta) * (1-self.band_width) * basefee > tx.max_fee
    
    def apply_eviction_policy(self, txs: Sequence[Tx1559], env: dict) -> Tuple[Sequence[Tx1559], Sequence[Tx1559]]:
        current_basefee = env["basefees"][-1]
        moving_average = sum(env["basefees"]) / len(env["basefees"]) if len(env["basefees"]) < self.window else sum([env["basefees"][-i] for i in range(1, self.window+1)]) // self.window
        delta = (current_basefee - moving_average) / current_basefee
        evicted_txs = [tx for tx in txs if self.should_evict(tx, current_basefee, delta)]
        accepted_txs = [tx for tx in txs if not self.should_evict(tx, current_basefee, delta)]
        return (evicted_txs, accepted_txs)
    
    def purge_pool_after_basefee_update(self, env: dict) -> Sequence[Tx1559]:
        evicted_txs, _ = self.apply_eviction_policy(self.txs.values(), env)
        self.remove_txs([tx.tx_hash for tx in evicted_txs])
        return evicted_txs


# In[3]:


class StrategicUser(User1559):
    
    def __init__(self, wakeup_block, **kwargs):
        super().__init__(wakeup_block, cost_per_unit = 0, **kwargs)
#         self.value = (1 + self.rng.pareto(2) * 20) * 1e9
        self.value = self.rng.uniform(10, 100) * 1e9

    def decide_parameters(self, env):
        # If previous block was close to full, consider being strategic
        if env["wallet_mode"] == "posted":
            gas_premium = MIN_PREMIUM
            max_fee = min(max(env["basefee"] * 3, MIN_PREMIUM * 3), self.value)
        else:
            max_fee = min(max(env["basefee"] * 3, MIN_PREMIUM * 3), self.value)
            gas_premium = min(max_fee, env["previous_avg_tip"] + MIN_PREMIUM)
            
        return {
            "max_fee": max_fee, # in wei
            "gas_premium": gas_premium, # in wei
            "start_block": self.wakeup_block,
        }

    def create_transaction(self, env):
        
        tx_params = self.decide_parameters(env)

        tx = Tx1559(
            sender = self.pub_key,
            tx_params = tx_params,
        )

        return tx

    def export(self):
        return {
            **super().export(),
            "user_type": "strategic_user",
        }


# ---
# 
# ## Demand generation

# In[7]:


def generate_gbm_demand(S0: float, duration: int, blocks: int, volatility: float, mean_ia_time: float, rng: np.random.Generator) -> Sequence[int]:
    mu = 0.5 * volatility**2
    gbm = list(generate_gbm(S0, duration, paths=1, mu=mu, sigma=volatility, rng=rng).flatten())
    return apply_block_time_variance(gbm, blocks, mean_ia_time=mean_ia_time, rng=rng)


# In[8]:


def generate_decaying_abm_demand(S0: float, duration: int, blocks: int, volatility: float, mean_ia_time: float, rng: np.random.Generator) -> Sequence[int]:
    pp = poisson_process(0.001, duration, rng)
    jp = jump_process(pp, duration, S0, rng, discount = 0.01)
    abm = list(generate_abm(S0, duration, paths=1, mu=0, sigma=volatility, rng=rng).flatten())
    return apply_block_time_variance(abm + jp, blocks, mean_ia_time=mean_ia_time, rng=rng)


# ---

# In[9]:


def simulate(demand_scenario, demand_id, shares_scenario, txpool, extra_metrics = None, rng = rng, silent=False):
    # Instantiate a couple of things
    chain = Chain()
    metrics = []
    user_pool = UserPool()
    max_basefee_window = 30
    start_time = time.time()
    block_target = int(constants["MAX_GAS_EIP1559"] / constants["SIMPLE_TRANSACTION_GAS"] / 2.0)
    
    # `env` is the "environment" of the simulation
    env = {
        "basefee": constants["INITIAL_BASEFEE"],
        "current_block": None,
        "basefees": [constants["INITIAL_BASEFEE"]],
        "previous_gas_used": 50, # in percent
        "previous_avg_tip": 1e9,
        "wallet_mode": "posted",
    }

    for t in tqdm(range(len(demand_scenario)), disable=silent, desc="simulation loop", leave=False):
        if demand_scenario[t] > 5 * MAX_TX_POOL:
            break
        
        # Sets current block
        env["current_block"] = t
        
        # Reset the random number generator with new seed to generate users with same values across runs
        rng = np.random.default_rng((2 ** t) * (3 ** demand_id))
        
        ### SIMULATION ###
        
        # The transaction pool applies its eviction policy after basefee update
        basefee_update_evicted_txs = txpool.purge_pool_after_basefee_update(env)

        # 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)
        
        # Add new users to the pool
        # 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)
        
        # Divide incoming transactions between accepted and rejected ones
        incoming_evicted_txs, accepted_txs = txpool.apply_eviction_policy(decided_txs, env)

        # New transactions are added to the transaction pool
        # `evicted_txs` holds the transactions removed from the pool for lack of space
        pool_limit_evicted_txs = txpool.add_txs(accepted_txs, env)

        # The best valid transactions are taken out of the pool for inclusion
        selected_txs = txpool.select_transactions(env, rng = rng)
        
        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 = env["basefee"]
        )
        
        # Record the gas used and avg tip in the block
        env["previous_gas_used"] = int(block.gas_used() / constants["MAX_GAS_EIP1559"] * 100)
        env["previous_avg_tip"] = int(block.average_tip() * 1e9)
        env["wallet_mode"] = "posted" if env["previous_gas_used"] <= 90 else "expert"
        
        # The block is added to the chain
        chain.add_block(block)
        
        ### METRICS ###
        user_efficiency = sum([user_pool.get_user(tx.sender).value for tx in selected_txs])
                
        row_metrics = {
            "block": t,
            "basefee_update_evictions": len(basefee_update_evicted_txs),
            "users": len(users),
            "decided_txs": len(decided_txs),
            "incoming_evictions": len(incoming_evicted_txs),
            "pool_limit_evictions": len(pool_limit_evicted_txs),
            "included_txs": len(selected_txs),
            "basefee": env["basefee"] / 1e9, # to Gwei
            "gas_used": env["previous_gas_used"],
            "avg_tip": env["previous_avg_tip"] / 1e9,
            "pool_length": txpool.pool_length(),
            "user_efficiency": user_efficiency,
        }
        
        if not extra_metrics is None:
            row_metrics = {
                **row_metrics,
                **extra_metrics(env, users, user_pool, txpool),
            }
        
        metrics.append(row_metrics)
        
        # Finally, basefee is updated and a new round starts
        env["basefee"] = update_basefee(block, env["basefee"])
        env["basefees"] = env["basefees"][-(max_basefee_window-1):] + [env["basefee"]]

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

volatilities = [0.1, 0.5, 1.0]
# volatilities = [1.0]
demand_paths = range(20)
band_widths = [0.0, 1.0/3, 2.0/3, 1.0]
# band_widths = [0.0]
blocks = 600
metrics = {}
dfs = []
run = 0
gamma = 76000
mean_ia_time = 13
max_k = int(2 * blocks * mean_ia_time)
max_txs_in_block = int(constants["MAX_GAS_EIP1559"] / gamma)
S0s = [int(k * max_txs_in_block / mean_ia_time) for k in [1, 1.5, 2]]

for params in tqdm(list(product(
    demand_paths, volatilities, band_widths, S0s
)), desc="session loop", leave=False):
    (demand_path, volatility, band_width, S0) = params
    
    # Generate demand
    rng = np.random.default_rng(demand_path)
    demand_scenario = np.maximum(generate_decaying_abm_demand(S0, max_k, blocks, volatility, mean_ia_time, rng), 0)
   
    # Shares of new users per time step
    shares_scenario = [{
        StrategicUser: 1.00,
    } for t in range(blocks)]

    # Transaction pool
    txpool = TxPoolTrendPicker(band_width = band_width)
    
    # Simulate
    (df, user_pool, chain) = simulate(demand_scenario, demand_path, shares_scenario, txpool, silent=True)
    df["run"] = run
    df["volatility"] = volatility
    df["band_width"] = band_width
    df["demand_path"] = demand_path
    df["S0"] = S0
    run += 1
    dfs += [df]


# In[10]:


df = pd.concat(dfs)
df.to_csv("../data/sim_runs.csv", index=False, index_label=False)

Dynamical Analysis of the EIP-1559 Ethereum Fee Market (code)// Authors
let authorData = ["barnabe"];