class StrategicUser(User1559):
    
    epsilon = 0.1 # how much the user overbids by

    def __init__(self, wakeup_block, **kwargs):
        super().__init__(wakeup_block, cost_per_unit = 0, **kwargs)
        self.value = (2 + self.rng.pareto(2) * 20) * (10 ** 9)
        self.transacted = False

    def decide_parameters(self, env):
        if env["min_premium"] is None:
            min_premium = 1 * (10 ** 9)
        else:
            if env["close_to_full"]:
                min_premium = env["min_premium"] + self.epsilon * (10 ** 9)
            else:
                min_premium = 1e9

        gas_premium = min_premium
        max_fee = self.value

        return {
            "max_fee": max_fee, # in wei
            "gas_premium": gas_premium, # in wei
            "start_block": self.wakeup_block
        }
        
    def create_transaction(self, env):
        if self.transacted:
            return None
        
        tx = super().create_transaction(env)
        if not tx is None:
            tx.gas_used = constants["SIMPLE_TRANSACTION_GAS"]
            self.transacted = True
        return tx
    
    def export(self):
        return {
            **super().export(),
            "user_type": "strategic_user_1559",
        }

    def __str__(self):
        return f"1559 strategic affine user with value {self.value} and cost {self.cost_per_unit}"

class PatientUser(User1559):
    def __init__(self, wakeup_block, **kwargs):
        super().__init__(wakeup_block, cost_per_unit = 0, **kwargs)
        self.value = (2 + self.rng.pareto(2) * 20) * (10 ** 9)
        self.patience = 10
        self.transacted = False
        
    def update_patience(self):
        self.patience = self.patience - 1
        
    def create_transaction(self, env):
        if self.transacted:
            return None
        
        if self.patience == 0 or env["is_full"]:
            tx = super().create_transaction(env)
            if tx is not None:
                tx.gas_used = constants["SIMPLE_TRANSACTION_GAS"]
                self.transacted = True
            return tx
        else:
            self.update_patience()
            return None

class PatientUserPool(UserPool):
    
    def decide_transactions(self, users: Sequence[User], env: Dict, query_all: bool = False) -> Sequence[Transaction]:
        txs = []
        for user in users:
            self.users[user.pub_key] = user
        
        self.users = { pub_key: user for pub_key, user in self.users.items() if user.wakeup_block >= env["current_block"] - 20 }
        
        # We first ask non-patient users if they'd like to transact
        for user in self.users.values():
            if type(user) is StrategicUser:
                tx = user.transact(env)
                if not tx is None:
                    txs.append(tx)
        
        # Simulate a block being built by a miner to determine the gas used by the next block
        max_tx_in_block = int(constants["MAX_GAS_EIP1559"] / constants["SIMPLE_TRANSACTION_GAS"])

        valid_txs = [tx for tx in txs if tx.is_valid(env)]
        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]
        
        # If gas used is more than half of the block, patient users transact
        env["is_full"] = len(selected_txs) >= max_tx_in_block // 2
        
        for user in self.users.values():
            if type(user) is PatientUser:
                tx = user.transact(env)
                if not tx is None:
                    txs.append(tx)
                    
        return txs

MAX_TRANSACTIONS_IN_POOL = 500
MIN_ACCEPTABLE_TIP = 1e9

class MixedTxPool(TxPool):
    
    def add_txs(self, txs: Sequence[Transaction], env: dict) -> Sequence[Transaction]:
        for tx in txs:
            self.txs[tx.tx_hash] = tx
                
        if self.pool_length() > MAX_TRANSACTIONS_IN_POOL:
            sorted_txs = sorted(self.txs.values(), key = lambda tx: -tx.tip(env))
            self.empty_pool()
            self.add_txs(sorted_txs[0:MAX_TRANSACTIONS_IN_POOL], env)
            return sorted_txs[MAX_TRANSACTIONS_IN_POOL:]
        
        return []
    
    def select_transactions(self, env, user_pool=None, rng=rng) -> Sequence[Transaction]:
        # 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) >= MIN_ACCEPTABLE_TIP]
        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 update_basefee_additive(block: Block, basefee: int) -> int:
    """
    Basefee update rule

    Args:
        block (Block): The previous block
        basefee (int): The current basefee

    Returns:
        int: The new basefee
    """

    gas_used = sum([tx.gas_used for tx in block.txs])
    delta = gas_used - constants["TARGET_GAS_USED"]
    new_basefee = max(0, basefee + (2 ** 30) * delta // constants["TARGET_GAS_USED"] // constants["BASEFEE_MAX_CHANGE_DENOMINATOR"])
    return new_basefee

def simulate(demand_scenario, shares_scenario, extra_metrics = None, rng = rng, additive_rule = False):
    # Instantiate a couple of things
    txpool = MixedTxPool()
    chain = Chain()
    metrics = []
    user_pool = PatientUserPool()
    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,
        "min_premium": 1 * (10 ** 9),
        "is_full": False,
        "close_to_full": False
    }

    for t in tqdm(range(len(demand_scenario))):
        # 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(t)
        
        ### SIMULATION ###

        # We return some demand which on expectation yields `demand_scenario[t]` new users per round
        users = spawn_fixed_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, query_all=True)
        
        patient_sent = len([tx for tx in decided_txs if type(user_pool.get_user(tx.sender)) is PatientUser])

        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 = env["basefee"]
        )
        
        # Record the min premium in the block
        env["min_premium"] = block.min_premium()
        # If block is 90% full, strategic users kick in
        env["close_to_full"] = len(selected_txs) >= demand_equal_target * 2 * 0.9
        
        # The block is added to the chain
        chain.add_block(block)
        
        ### METRICS ###
                
        row_metrics = {
            "block": t,
            "users": len(users),
            "decided_txs": len(decided_txs),
            "included_txs": len(selected_txs),
            "basefee": env["basefee"] / (10 ** 9), # to Gwei
            "blk_min_premium": block.min_premium() / (10 ** 9), # to Gwei
            "blk_max_premium": block.max_premium() / (10 ** 9), # to Gwei
            "blk_min_tip": block.min_tip(env) / (10 ** 9), # to Gwei
            "blk_max_tip": block.max_tip(env) / (10 ** 9), # to Gwei
            "patient_sent": patient_sent,
        }
        
        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
        if additive_rule:
            env["basefee"] = update_basefee_additive(block, env["basefee"])
        else:
            env["basefee"] = update_basefee(block, env["basefee"])

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