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"

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])

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

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

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)

# 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")

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()