%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

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

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",
        }

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)

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)

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