example.py

"""
Demonstrate randomly minting/burning and collecting data.  

Main actors:
- Exchange: A minter/burner (think coinbase)
- Trader: Buy/sell USDC from Exchange

Data on mint/burns is collected every X number of steps (as dataframe)
"""

import mesa
import polars as pl
from tqdm import tqdm
from usdc import load_usdc, load_evm_from_snapshot, fiat_to_usdc, MM
from simular import create_account, ether_to_wei, create_many_accounts

# USDC: default allocation quota
ALLOCATION_TO_MINTER = fiat_to_usdc(10_000)

# Trader: prob they trade, and amount they may buy/sell
TRADE_PROB = 0.6
TRADE_AMOUNT = fiat_to_usdc(2)


class Trader(mesa.Agent):
    """
    I randomly buy/sell USDC
    """

    def __init__(self, address, model):
        super().__init__(model)
        self.address = address

    def step(self):
        if self.model.random.random() < TRADE_PROB:
            choice = self.model.random.choice([0, 1])
            if 0 == choice:
                # buy
                self.model.exchange.buy(self.address, TRADE_AMOUNT)
            else:
                # sell
                bal = self.model.usdc.balanceOf.call(self.address)
                if bal >= TRADE_AMOUNT:
                    self.model.exchange.sell(self.address, TRADE_AMOUNT)


class Exchange:
    """
    I'm an exchange that interacts with the USDC contract. Traders use me
    to buy/sell USDC
    """

    def __init__(self, address, model):
        self.address = address
        self.model = model

        # data collected
        self.current_mints = 0
        self.current_burns = 0

        self.step = []
        self.total_mints = []
        self.total_burns = []
        self.total_supply = []

    def mint(self, amt):
        """
        Mint USDC to self.  May throw an exception IF I'm out of minting allocations
        """
        self.model.usdc.mint.transact(self.address, amt, caller=self.address)
        self.current_mints += amt

    def burn(self, amt):
        """
        Burn USDC. Only a minter can burn USDC
        """
        self.model.usdc.burn.transact(amt, caller=self.address)
        self.current_burns += amt

    def buy(self, to_address, amt):
        """
        Called by a trader to buy USDC
        """
        # Check the exchange has enough to fullfil the order
        bal = self.model.usdc.balanceOf.call(self.address)
        if bal <= amt:
            # If not, mint some
            self.mint(amt)

        # transfer the USDC to the buyer
        self.model.usdc.transfer.transact(to_address, amt, caller=self.address)

    def sell(self, onbehalf__address, amt):
        """
        Sell USDC:
        1. Trader transfers their USDC to me
        2. I burn the USDC.
        Note: only a minter can burn USDC
        """
        self.model.usdc.transfer.transact(self.address, amt, caller=onbehalf__address)
        self.burn(amt)

    def collect_data(self, step):
        """
        Collects model data. Called at the end of a step
        """
        self.step.append(step)
        self.total_mints.append(self.current_mints / 1e6)
        self.total_burns.append(self.current_burns / 1e6)
        self.total_supply.append(self.model.usdc.totalSupply.call() / 1e6)

        # cleared for step
        self.current_mints = 0
        self.current_burns = 0

    def to_data_frame(self):
        return pl.DataFrame(
            [
                self.step,
                self.total_mints,
                self.total_burns,
                self.total_supply,
            ],
            schema={
                "step": pl.Int32,
                "mint": pl.Float64,
                "burn": pl.Float64,
                "supply": pl.Float64,
            },
        )


class USDCModel(mesa.Model):
    """
    This model has 1 agent - a trader. The goal is to capture mint/burns:

    - Traders buy/sell USDC from the Exchange.
    - The Exchange needs to mint/burn to keep up with the traders.
    - At each step, the trader may randomly buy/sell triggering activity on the Exchange
    """

    def __init__(self, num_steps=1000, num_traders=10):
        super().__init__()
        self.num_steps = num_steps
        self.num_traders = num_traders

        # Load the EVM engine
        self.evm = load_evm_from_snapshot()
        # Load the USDC contracts for use
        self.usdc = load_usdc(self.evm)

        # Create a wallet for the exchange
        exchange_wallet = create_account(self.evm, value=ether_to_wei(1))
        # Create wallets for the traders
        trader_wallets = create_many_accounts(
            self.evm, num_traders, value=ether_to_wei(1)
        )
        # Master minter configures the exchange to be a minter with the default allocation
        self.usdc.configureMinter.transact(
            exchange_wallet, ALLOCATION_TO_MINTER, caller=MM
        )

        self.exchange = Exchange(exchange_wallet, self)
        for w in trader_wallets:
            self.agents.add(Trader(w, self))

    def step(self):
        self.agents.shuffle_do("step")
        self.exchange.collect_data(self.steps)

    def run_model(self):
        for _ in tqdm(range(self.num_steps)):
            self.step()

    def data(self):
        return self.exchange.to_data_frame()


if __name__ == "__main__":
    model = USDCModel()
    model.run_model()
    df = model.data()

    print(df)

    df.write_avro("./example_data.avro")