model.py

"""
model.py
========
DeFiSimulationModel — The Mesa Model orchestrating the entire simulation.

Responsibilities:
  - Initialise the MockDeFiProtocol (the shared fake blockchain)
  - Spawn all agent types and register them with the scheduler
  - Collect aggregate metrics (pool liquidity, prices) every step via DataCollector
  - Advance the oracle price before agents act each step
  - Pipe all agent actions to the SQLite database via SimulationLogger
"""

import logging
import mesa
from mesa.datacollection import DataCollector

from DeFiProtocolInterface import DeFiProtocolInterface
from MockDeFiProtocol       import MockDeFiProtocol
from database               import SimulationLogger
from agents                 import (
    RetailTraderAgent,
    MidTierTraderAgent,
    CryptoWhaleAgent,
    ArbitrageTraderAgent,
    MEVSearcherAgent,
    MarketMakerAgent,
    InstitutionalParticipantAgent,
)
import config

logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
# DataCollector helper functions (must be module-level for pickling) #
# ------------------------------------------------------------------ #

def _collect_reserve_usdc(model: "DeFiSimulationModel") -> float:
    return model.protocol.get_pool_liquidity()["reserve_usdc"]

def _collect_reserve_eth(model: "DeFiSimulationModel") -> float:
    return model.protocol.get_pool_liquidity()["reserve_eth"]

def _collect_token_price(model: "DeFiSimulationModel") -> float:
    return model.protocol.get_token_price()

def _collect_oracle_price(model: "DeFiSimulationModel") -> float:
    return model.protocol.get_oracle_price()

def _collect_price_spread_pct(model: "DeFiSimulationModel") -> float:
    """Signed % spread between AMM and oracle price — key anomaly signal."""
    amm    = model.protocol.get_token_price()
    oracle = model.protocol.get_oracle_price()
    if oracle == 0:
        return 0.0
    return (amm - oracle) / oracle * 100.0


class DeFiSimulationModel(mesa.Model):
    """
    Top-level Mesa Model for the DeFi Risk Simulation Lab.
    """

    def __init__(
        self,
        seed:               int = config.RANDOM_SEED,
        db_path:            str = config.DB_PATH,
    ):
        super().__init__(seed=seed)

        logger.info("=" * 60)
        logger.info("DeFiSimulationModel initialising ...")

        # ---------------------------------------------------------------- #
        # 1. PROTOCOL — the shared fake blockchain (passed as interface)   #
        # ---------------------------------------------------------------- #
        self.protocol: DeFiProtocolInterface = MockDeFiProtocol(rng=None)

        # ---------------------------------------------------------------- #
        # 2. DATA PIPELINE — shared SQLite logger                          #
        # ---------------------------------------------------------------- #
        self.db = SimulationLogger(db_path=db_path)

        # ---------------------------------------------------------------- #
        # 3. SCHEDULER — RandomActivation: agents act in random order      #
        # ---------------------------------------------------------------- #
        self.schedule = mesa.time.RandomActivation(self)

        # ---------------------------------------------------------------- #
        # 4. SPAWN AGENTS                                                  #
        # ---------------------------------------------------------------- #
        all_accounts = [f"account_{i:03d}" for i in range(config.TOTAL_ACCOUNTS)]
        
        idx = 0
        def get_slice(count: int):
            nonlocal idx
            chunk = all_accounts[idx:idx+count]
            idx += count
            return chunk

        retail_accs = get_slice(config.ACCOUNTS_RETAIL)
        mid_tier_accs = get_slice(config.ACCOUNTS_MID_TIER)
        whale_accs = get_slice(config.ACCOUNTS_WHALE)
        arb_accs = get_slice(config.ACCOUNTS_ARBITRAGE)
        mev_accs = get_slice(config.ACCOUNTS_MEV)
        mm_accs = get_slice(config.ACCOUNTS_MARKET_MAKER)
        inst_accs = get_slice(config.ACCOUNTS_INSTITUTIONAL)

        logger.info(
            "Accounts distributed: %d Retail, %d MidTier, %d Whale, %d Arb, %d MEV, %d MM, %d Inst",
            len(retail_accs), len(mid_tier_accs), len(whale_accs), len(arb_accs),
            len(mev_accs), len(mm_accs), len(inst_accs)
        )

        agent_id = 0
        
        ag_retail = RetailTraderAgent(agent_id, self, self.protocol, self.db, accounts=retail_accs)
        self.schedule.add(ag_retail)
        agent_id += 1
        
        ag_mid = MidTierTraderAgent(agent_id, self, self.protocol, self.db, accounts=mid_tier_accs)
        self.schedule.add(ag_mid)
        agent_id += 1
        
        ag_whale = CryptoWhaleAgent(agent_id, self, self.protocol, self.db, accounts=whale_accs)
        self.schedule.add(ag_whale)
        agent_id += 1
        
        ag_arb = ArbitrageTraderAgent(agent_id, self, self.protocol, self.db, accounts=arb_accs)
        self.schedule.add(ag_arb)
        agent_id += 1
        
        ag_mev = MEVSearcherAgent(agent_id, self, self.protocol, self.db, accounts=mev_accs)
        self.schedule.add(ag_mev)
        agent_id += 1
        
        ag_mm = MarketMakerAgent(agent_id, self, self.protocol, self.db, accounts=mm_accs)
        self.schedule.add(ag_mm)
        agent_id += 1
        
        ag_inst = InstitutionalParticipantAgent(agent_id, self, self.protocol, self.db, accounts=inst_accs)
        self.schedule.add(ag_inst)
        agent_id += 1

        self.total_agents = agent_id
        logger.info("Total agents spawned: %d managing %d accounts", self.total_agents, config.TOTAL_ACCOUNTS)

        # ---------------------------------------------------------------- #
        # 5. DATA COLLECTOR — Mesa built-in, one row per step              #
        # ---------------------------------------------------------------- #
        self.datacollector = DataCollector(
            model_reporters={
                "pool_reserve_usdc":  _collect_reserve_usdc,
                "pool_reserve_eth":   _collect_reserve_eth,
                "token_price":        _collect_token_price,
                "oracle_price":       _collect_oracle_price,
                "price_spread_pct":   _collect_price_spread_pct,
            }
        )

        # Collect initial state (step 0)
        self.datacollector.collect(self)
        logger.info(
            "Initial state | AMM price=%.2f | oracle=%.2f | "
            "pool=%.0f USDC / %.1f ETH",
            self.protocol.get_token_price(),
            self.protocol.get_oracle_price(),
            self.protocol.get_pool_liquidity()["reserve_usdc"],
            self.protocol.get_pool_liquidity()["reserve_eth"],
        )
        logger.info("=" * 60)

    # ------------------------------------------------------------------ #
    # STEP                                                                 #
    # ------------------------------------------------------------------ #

    def step(self) -> None:
        """
        Advance the simulation by one tick:
          1. Advance oracle price (before agents perceive it)
          2. Activate all agents in random order (RandomActivation)
          3. Process Liquidity Pool Behavioral Triggers
          4. Collect DataCollector metrics for this step
        """
        current_step = self.schedule.steps + 1
        logger.info("--- Step %d / %d ---", current_step, config.SIMULATION_STEPS)

        # Step 1: Oracle price random-walk (agents will see new price)
        new_oracle = self.protocol.advance_oracle_price()
        logger.debug("Oracle advanced: %.2f", new_oracle)

        # Step 2: All agents act (RandomActivation handles random ordering)
        self.schedule.step()

        # Step 3: Process Liquidity Pool Behavioral Triggers
        current_price = self.protocol.get_token_price()
        if not hasattr(self, "initial_token_price"):
            self.initial_token_price = current_price

        if self.initial_token_price > 0:
            drop_pct = (self.initial_token_price - current_price) / self.initial_token_price * 100.0
            
            # Check for protocol config
            if hasattr(self.protocol, "liquidity_config"):
                liq_config = getattr(self.protocol, "liquidity_config", {}) or getattr(self.protocol, "_mock", {}).liquidity_config if hasattr(self.protocol, "_mock") else {}
            elif hasattr(self.protocol, "_mock") and hasattr(self.protocol._mock, "liquidity_config"):
                liq_config = self.protocol._mock.liquidity_config
            else:
                liq_config = {}

            if liq_config:
                behavior = liq_config.get("behavioral_parameters", {})
                trigger_drop = behavior.get("withdrawal_trigger_price_drop_percent")
                
                if trigger_drop and drop_pct >= trigger_drop:
                    if not getattr(self, "has_bank_run", False):
                        fraction = behavior.get("withdrawal_fraction_on_trigger", 0.5)
                        self.protocol.execute_withdraw("LP_MIGRATION", fraction)
                        self.has_bank_run = True
                        logger.warning("Bank run triggered! LP withdrew %.2f%% of liquidity.", fraction * 100)

                        # Log this event into the database as a system event
                        self.db.log(
                            step=current_step,
                            agent_id=-1,
                            agent_type="LIQUIDITY_POOL",
                            action="BANK_RUN_WITHDRAWAL",
                            amount=0.0,
                            token="ALL",
                            pool_reserve_usdc=self.protocol.get_pool_liquidity()["reserve_usdc"],
                            pool_reserve_eth=self.protocol.get_pool_liquidity()["reserve_eth"],
                            token_price=current_price,
                            oracle_price=new_oracle,
                            reasoning=f"Price dropped {drop_pct:.1f}% >= trigger {trigger_drop}%"
                        )

        # Step 4: Collect aggregate metrics
        self.datacollector.collect(self)

    # ------------------------------------------------------------------ #
    # TEARDOWN                                                             #
    # ------------------------------------------------------------------ #

    def close(self) -> None:
        """Call this after the simulation finishes to flush the DB."""
        self.db.close()
        logger.info("DeFiSimulationModel closed.")