MockDeFiProtocol.py

"""
MockDeFiProtocol.py
===================
Concrete implementation of DeFiProtocolInterface using a constant-product AMM.

This is the "Adapter" in the Adapter Pattern. It simulates a full DeFi
protocol in pure Python so agents have a fake-but-realistic blockchain to
interact with during development — before the real smart contracts are ready.

AMM invariant: reserve_usdc * reserve_eth = k  (Uniswap V2 style)
"""

import random
import math
import logging
from typing import Dict, List

from DeFiProtocolInterface import DeFiProtocolInterface
import config

logger = logging.getLogger(__name__)


class MockDeFiProtocol(DeFiProtocolInterface):
    """
    Simulates a DeFi protocol with:
      - Constant-product AMM for token swaps
      - Lending module with health factors
      - External oracle with random-walk price dynamics
    """

    def __init__(self, rng: random.Random | None = None, amm_config: Dict = None, liquidity_config: Dict = None):
        """
        Args:
            rng: Optional seeded Random instance for reproducibility.
            amm_config: Optional dict of AMM configuration parameters.
            liquidity_config: Optional dict of Liquidity Pool configuration parameters.
        """
        self._rng = rng or random.Random(config.RANDOM_SEED)

        self.amm_config = amm_config or {}
        self.liquidity_config = liquidity_config or {}

        # ------------------------------------------------------------------ #
        # ORACLE PRICE (external reference)
        # Sets up first so we can price initial TVL
        # ------------------------------------------------------------------ #
        self.oracle_price: float = config.ORACLE_INITIAL_PRICE

        # ------------------------------------------------------------------ #
        # AMM POOL STATE
        # ------------------------------------------------------------------ #
        if self.liquidity_config and "core_model" in self.liquidity_config:
            initial_tvl = self.liquidity_config["core_model"].get("initial_tvl", config.INITIAL_RESERVE_USDC * 2)
            self.reserve_usdc = initial_tvl / 2.0
            self.reserve_eth = (initial_tvl / 2.0) / self.oracle_price
        else:
            self.reserve_usdc: float = config.INITIAL_RESERVE_USDC
            self.reserve_eth:  float = config.INITIAL_RESERVE_ETH

        self.k: float = self.reserve_usdc * self.reserve_eth  # invariant

        # AMM behavior params
        amm_behavior = self.amm_config.get("behavioral_parameters", {})
        self.swap_fee = amm_behavior.get("swap_fee_percent", 0.3) / 100.0
        
        amm_core = self.amm_config.get("core_model", {})
        self.amm_type = amm_core.get("type", "xyk")
        if self.amm_type == "stableswap":
            self.amp_factor = amm_core.get("parameters", {}).get("A", 85)
        elif self.amm_type == "concentrated_liquidity":
            self.price_range = amm_core.get("parameters", {}).get("range_percent", 10.0)

        # ------------------------------------------------------------------ #
        # LENDING MODULE — borrower positions
        # {address: {"collateral_eth": float, "debt_usdc": float}}
        # ------------------------------------------------------------------ #
        self._borrowers: Dict[str, Dict[str, float]] = {}

        # Seed with pre-built borrower positions (close to liquidation)
        self._seed_borrowers()

        logger.info(
            "MockDeFiProtocol initialised [%s] | fee=%.2f%% | "
            "pool=%.0f USDC / %.1f ETH | price=%.2f | oracle=%.2f",
            self.amm_type, self.swap_fee * 100,
            self.reserve_usdc, self.reserve_eth,
            self.get_token_price(), self.oracle_price,
        )

    # ====================================================================== #
    # PRIVATE HELPERS                                                          #
    # ====================================================================== #

    def _seed_borrowers(self) -> None:
        """Creates NUM_BORROWERS pre-seeded positions near the liquidation edge."""
        for i in range(config.NUM_BORROWERS):
            addr = f"borrower_{i:04d}"
            # Collateral: 1–3 ETH per borrower
            collateral_eth = self._rng.uniform(1.0, 3.0)
            # Debt is chosen to give the target health factor defined in config
            # health_factor = (collateral_eth * price * LTV) / debt_usdc
            # => debt_usdc = (collateral_eth * price * LTV) / health_factor
            collateral_usdc_value = collateral_eth * self.oracle_price
            debt_usdc = (
                collateral_usdc_value
                * config.LIQUIDATION_THRESHOLD
                / config.BORROWER_START_HEALTH_FACTOR
            )
            self._borrowers[addr] = {
                "collateral_eth": collateral_eth,
                "debt_usdc":      debt_usdc,
            }
            logger.debug(
                "Seeded borrower %s | collateral=%.4f ETH | debt=%.2f USDC | HF=%.3f",
                addr, collateral_eth, debt_usdc,
                self.get_user_health_factor(addr),
            )

    def _update_k(self) -> None:
        """Recalculate the invariant after adding liquidity."""
        self.k = self.reserve_usdc * self.reserve_eth

    # ====================================================================== #
    # READ-ONLY QUERIES                                                        #
    # ====================================================================== #

    def get_token_price(self) -> float:
        """AMM-derived ETH price in USDC: reserve_usdc / reserve_eth."""
        if self.reserve_eth == 0:
            return float("inf")
        return self.reserve_usdc / self.reserve_eth

    def get_oracle_price(self) -> float:
        """Current external oracle price (independent of AMM state)."""
        return self.oracle_price

    def get_pool_liquidity(self) -> Dict[str, float]:
        """Returns the full pool state snapshot."""
        return {
            "reserve_usdc": self.reserve_usdc,
            "reserve_eth":  self.reserve_eth,
            "k":            self.k,
        }

    def get_user_health_factor(self, address: str) -> float:
        """
        health_factor = (collateral_eth * oracle_price * LTV) / debt_usdc
        Returns inf if address has no debt.
        """
        pos = self._borrowers.get(address)
        if pos is None or pos["debt_usdc"] <= 0:
            return float("inf")
        collateral_value = pos["collateral_eth"] * self.oracle_price
        return (collateral_value * config.LIQUIDATION_THRESHOLD) / pos["debt_usdc"]

    def has_bad_debt(self) -> bool:
        """Returns True if any borrower's collateral value is less than their debt."""
        for address, pos in self._borrowers.items():
            if pos["debt_usdc"] > 0:
                collateral_value = pos["collateral_eth"] * self.oracle_price
                if collateral_value < pos["debt_usdc"]:
                    return True
        return False

    def get_all_borrower_addresses(self) -> List[str]:
        """Returns all known borrower addresses."""
        return list(self._borrowers.keys())

    # ====================================================================== #
    # STATE-CHANGING ACTIONS                                                   #
    # ====================================================================== #

    def execute_swap(
        self,
        token_in: str,
        amount_in: float,
        caller_address: str,
    ) -> Dict[str, float]:
        """
        Swap tokens via the selected AMM model and fee percent.
        """
        price_before = self.get_token_price()
        amount_in_after_fee = amount_in * (1 - self.swap_fee)

        if token_in.upper() == "USDC":
            amount_out = self._calculate_out_amount(
                self.reserve_usdc, self.reserve_eth, amount_in_after_fee, "USDC"
            )
            amount_out = min(amount_out, self.reserve_eth * 0.99)
            self.reserve_usdc += amount_in
            self.reserve_eth  -= amount_out
            token_out = "ETH"

        elif token_in.upper() == "ETH":
            amount_out = self._calculate_out_amount(
                self.reserve_eth, self.reserve_usdc, amount_in_after_fee, "ETH"
            )
            amount_out = min(amount_out, self.reserve_usdc * 0.99)
            self.reserve_eth  += amount_in
            self.reserve_usdc -= amount_out
            token_out = "USDC"

        else:
            raise ValueError(f"Unknown token_in: {token_in!r}. Must be 'USDC' or 'ETH'.")

        # Clamp reserves (safety floor)
        self.reserve_usdc = max(self.reserve_usdc, 1.0)
        self.reserve_eth  = max(self.reserve_eth,  0.001)

        price_after  = self.get_token_price()
        price_impact = abs(price_after - price_before) / price_before if price_before else 0.0

        logger.info(
            "SWAP [%s] | caller=%s | in=%.4f %s | out=%.4f %s | "
            "price %.2f -> %.2f (impact %.2f%%)",
            self.amm_type, caller_address, amount_in, token_in,
            amount_out, token_out,
            price_before, price_after, price_impact * 100,
        )

        return {
            "amount_out":   amount_out,
            "new_price":    price_after,
            "price_impact": price_impact,
            "token_out":    token_out,
        }

    def _calculate_out_amount(self, reserve_in: float, reserve_out: float, amount_in: float, token_in: str) -> float:
        if self.amm_type == "constant_sum":
            # Linear price mock: treats ETH/USDC as 1:P where P is initial/current peg price
            current_price = self.get_token_price()
            if token_in == "USDC":
                return amount_in / current_price
            else:
                return amount_in * current_price
                
        elif self.amm_type == "stableswap":
            # Emulated stableswap (Curve) curve blending xyk with constant sum based on A
            A = getattr(self, "amp_factor", 85)
            cp_out = (reserve_out * amount_in) / (reserve_in + amount_in)
            current_price = reserve_in / reserve_out if token_in == "ETH" else reserve_out / reserve_in
            cs_out = amount_in * current_price
            weight_cs = min(A / (A + 10), 0.99)
            return (cs_out * weight_cs) + (cp_out * (1 - weight_cs))
            
        elif self.amm_type == "concentrated_liquidity":
            # Virtual reserves emulation based on concentrated range
            concentration = getattr(self, "price_range", 10.0) 
            virtual_multiplier = max(1.0, 100.0 / max(concentration, 1.0))
            v_reserve_in = reserve_in * virtual_multiplier
            v_reserve_out = reserve_out * virtual_multiplier
            return (v_reserve_out * amount_in) / (v_reserve_in + amount_in)
            
        else: # "xyk" and fallback
            return (reserve_out * amount_in) / (reserve_in + amount_in)

    def execute_liquidation(
        self,
        borrower_address: str,
        liquidator_address: str,
    ) -> Dict[str, float]:
        """
        Liquidates an unhealthy borrower.

        The liquidator seizes the borrower's collateral + a bonus.
        The borrower's full debt is wiped.
        """
        pos = self._borrowers.get(borrower_address)
        if pos is None:
            logger.warning("Liquidation called on unknown address: %s", borrower_address)
            return {"debt_repaid": 0.0, "collateral_seized": 0.0, "liquidation_bonus": 0.0}

        hf = self.get_user_health_factor(borrower_address)
        if hf >= 1.0:
            logger.warning(
                "Liquidation rejected for %s (HF=%.3f >= 1.0)", borrower_address, hf
            )
            return {"debt_repaid": 0.0, "collateral_seized": 0.0, "liquidation_bonus": 0.0}

        debt_usdc       = pos["debt_usdc"]
        collateral_eth  = pos["collateral_eth"]
        bonus_eth       = collateral_eth * config.LIQUIDATION_BONUS
        seized_eth      = collateral_eth + bonus_eth

        # Wipe the position
        self._borrowers[borrower_address] = {"collateral_eth": 0.0, "debt_usdc": 0.0}

        logger.info(
            "LIQUIDATION | liquidator=%s | borrower=%s | "
            "debt_repaid=%.2f USDC | seized=%.4f ETH (bonus=%.4f ETH)",
            liquidator_address, borrower_address,
            debt_usdc, seized_eth, bonus_eth,
        )

        return {
            "debt_repaid":       debt_usdc,
            "collateral_seized": seized_eth,
            "liquidation_bonus": bonus_eth,
        }

    def execute_borrow(
        self,
        borrower_address: str,
        amount_usdc: float,
    ) -> Dict[str, float]:
        """
        Borrows USDC. The caller must have some ETH collateral already
        (for simplicity, we auto-add collateral if the address is new).
        """
        if borrower_address not in self._borrowers:
            # New borrower: auto-collateralise with 2 ETH
            self._borrowers[borrower_address] = {
                "collateral_eth": 2.0,
                "debt_usdc":      0.0,
            }

        pos = self._borrowers[borrower_address]
        
        # Check against LTV max
        collateral_value = pos["collateral_eth"] * self.oracle_price
        max_borrow = collateral_value * config.LIQUIDATION_LTV_RATIO
        
        if pos["debt_usdc"] + amount_usdc > max_borrow:
            amount_usdc = max(0.0, max_borrow - pos["debt_usdc"])
            
        pos["debt_usdc"] += amount_usdc

        new_hf = self.get_user_health_factor(borrower_address)

        logger.info(
            "BORROW | address=%s | amount=%.2f USDC | new_HF=%.3f",
            borrower_address, amount_usdc, new_hf,
        )

        return {
            "borrowed_amount":   amount_usdc,
            "new_health_factor": new_hf,
        }

    def execute_deposit_collateral(
        self,
        borrower_address: str,
        amount_eth: float,
    ) -> Dict[str, float]:
        if borrower_address not in self._borrowers:
            self._borrowers[borrower_address] = {"collateral_eth": 0.0, "debt_usdc": 0.0}
        
        self._borrowers[borrower_address]["collateral_eth"] += amount_eth
        
        new_hf = self.get_user_health_factor(borrower_address)
        
        logger.info(
            "DEPOSIT_COLLATERAL | address=%s | amount=%.4f ETH | new_HF=%.3f",
            borrower_address, amount_eth, new_hf
        )
        return {
            "deposited_amount": amount_eth,
            "new_health_factor": new_hf,
        }

    def execute_withdraw_collateral(
        self,
        borrower_address: str,
        amount_eth: float,
    ) -> Dict[str, float]:
        pos = self._borrowers.get(borrower_address)
        if not pos or pos["collateral_eth"] < amount_eth:
            amount_eth = pos["collateral_eth"] if pos else 0.0
            
        if pos and pos["debt_usdc"] > 0:
            # Check LTV constraint -> you can't withdraw collateral if it pushes you above max LTV
            min_collateral_value = pos["debt_usdc"] / config.LIQUIDATION_LTV_RATIO
            min_collateral_eth = min_collateral_value / self.oracle_price
            max_withdraw = max(0.0, pos["collateral_eth"] - min_collateral_eth)
            
            amount_eth = min(amount_eth, max_withdraw)
            
        if pos:
            pos["collateral_eth"] -= amount_eth
            
        new_hf = self.get_user_health_factor(borrower_address)
        
        logger.info(
            "WITHDRAW_COLLATERAL | address=%s | amount=%.4f ETH | new_HF=%.3f",
            borrower_address, amount_eth, new_hf
        )
        return {
            "withdrawn_amount": amount_eth,
            "new_health_factor": new_hf,
        }

    def execute_repay(
        self,
        borrower_address: str,
        amount_usdc: float,
    ) -> Dict[str, float]:
        pos = self._borrowers.get(borrower_address)
        if not pos:
            return {"repaid_amount": 0.0, "new_health_factor": float("inf")}
            
        amount_usdc = min(amount_usdc, pos["debt_usdc"])
        pos["debt_usdc"] -= amount_usdc
        
        new_hf = self.get_user_health_factor(borrower_address)
        
        logger.info(
            "REPAY | address=%s | amount=%.2f USDC | new_HF=%.3f",
            borrower_address, amount_usdc, new_hf
        )
        return {
            "repaid_amount": amount_usdc,
            "new_health_factor": new_hf,
        }

    def execute_supply(
        self,
        supplier_address: str,
        token: str,
        amount: float,
    ) -> Dict[str, float]:
        """
        Supplies liquidity to the pool (grows reserves proportionally).
        LP tokens minted are proportional to the share of the pool added.
        """
        if token.upper() == "USDC":
            proportion   = amount / (self.reserve_usdc + amount) if self.reserve_usdc else 1.0
            lp_minted    = proportion * math.sqrt(self.k)  # simplified LP token model
            self.reserve_usdc += amount
            new_reserve  = self.reserve_usdc
        elif token.upper() == "ETH":
            proportion   = amount / (self.reserve_eth + amount) if self.reserve_eth else 1.0
            lp_minted    = proportion * math.sqrt(self.k)
            self.reserve_eth  += amount
            new_reserve  = self.reserve_eth
        else:
            raise ValueError(f"Unknown token: {token!r}")

        self._update_k()  # k grows when liquidity is added

        logger.info(
            "SUPPLY | address=%s | %.4f %s | LP_minted=%.4f",
            supplier_address, amount, token, lp_minted,
        )

        return {
            "lp_tokens_minted": lp_minted,
            "new_reserve":      new_reserve,
        }

    def execute_withdraw(
        self,
        withdrawer_address: str,
        fraction: float,
    ) -> Dict[str, float]:
        """
        Withdraws liquidity equivalent to a fraction of the total pool.
        """
        fraction = max(0.0, min(1.0, fraction))
        
        usdc_amount = self.reserve_usdc * fraction
        eth_amount = self.reserve_eth * fraction
        
        self.reserve_usdc -= usdc_amount
        self.reserve_eth -= eth_amount
        
        # Clamp reserves
        self.reserve_usdc = max(self.reserve_usdc, 1.0)
        self.reserve_eth  = max(self.reserve_eth,  0.001)
        self._update_k()

        logger.info(
            "WITHDRAW | address=%s | fraction=%.2f%% | USDC=%.2f | ETH=%.4f",
            withdrawer_address, fraction * 100, usdc_amount, eth_amount
        )

        return {
            "usdc_withdrawn": usdc_amount,
            "eth_withdrawn": eth_amount,
        }

    def get_lending_state(self) -> Dict:
        """
        Returns a full snapshot of the lending module for the current step.
        Covers all active positions: collateral, debt, health factors,
        utilization, and bad debt.
        """
        active = {
            addr: pos for addr, pos in self._borrowers.items()
            if pos["debt_usdc"] > 0 or pos["collateral_eth"] > 0
        }

        total_collateral_eth = sum(pos["collateral_eth"] for pos in active.values())
        total_debt_usdc      = sum(pos["debt_usdc"]      for pos in active.values())
        total_collateral_val = total_collateral_eth * self.oracle_price

        health_factors: List[float] = []
        liquidatable   = 0
        bad_debt_usdc  = 0.0

        for addr, pos in active.items():
            if pos["debt_usdc"] > 0:
                hf = self.get_user_health_factor(addr)
                health_factors.append(hf)
                if hf < 1.0:
                    liquidatable += 1
                col_val = pos["collateral_eth"] * self.oracle_price
                if col_val < pos["debt_usdc"]:
                    bad_debt_usdc += pos["debt_usdc"] - col_val

        min_hf = min(health_factors) if health_factors else 9999.0
        avg_hf = sum(health_factors) / len(health_factors) if health_factors else 9999.0

        # Utilization: total debt vs pool USDC (lending side proxy)
        utilization = (
            total_debt_usdc / self.reserve_usdc
            if self.reserve_usdc > 0 else 0.0
        )

        return {
            "total_collateral_eth":        round(total_collateral_eth, 4),
            "total_collateral_value_usdc": round(total_collateral_val, 2),
            "total_debt_usdc":             round(total_debt_usdc, 2),
            "active_borrower_count":       len(active),
            "liquidatable_count":          liquidatable,
            "min_health_factor":           round(min_hf, 4),
            "avg_health_factor":           round(avg_hf, 4),
            "bad_debt_usdc":               round(bad_debt_usdc, 2),
            "utilization_rate":            round(utilization, 4),
            "has_bad_debt":                bad_debt_usdc > 0,
        }

    def advance_oracle_price(self) -> float:
        """
        Random walk: oracle_price *= (1 + drift + volatility * N(0,1))
        Clipped to never go negative or below $10.
        """
        shock = self._rng.gauss(0, 1)
        pct_change = config.ORACLE_DRIFT + config.ORACLE_VOLATILITY * shock
        self.oracle_price *= (1.0 + pct_change)
        self.oracle_price = max(self.oracle_price, 10.0)  # floor at $10

        logger.debug("Oracle price advanced to %.2f", self.oracle_price)
        return self.oracle_price