Redo investment logic

commands/commands/rollcoin.py

@@ -37,14 +37,14 @@ class RollcoinState:
 @dataclass
 class RollcoinWallet:
     balance: float = 1
-    invested: float = 0
-    investment_value: float = 0
+    holdings: float = 0
+    cost_basis: float = 0
     nfts: list[str] = field(default_factory=list)
 
     def __str__(self):
         s = f"Wallet: {self.balance}\n"
-        if self.invested > 0 or self.investment_value > 0:
-            s += f"Invested: {self.investment_value}\n\t(cost basis {self.invested})\n"
+        if self.holdings > 0 or self.cost_basis > 0:
+            s += f"Investments: {self.holdings}\n\t(cost basis {self.cost_basis})\n"
         if len(self.nfts) > 0:
             s += f"NFTs:\n"
         for nft in self.nfts:
@@ -126,14 +126,11 @@ async def tip(
 
 
 async def evolve_market(state, transitions, multipliers, wallet_store, state_store):
-    weights = [row[state.market_state] for row in transitions]
-    state.market_state = random.choices(range(len(transitions)), weights=weights, k=1)[0]
-    # TODO increase treasury?
+    state.market_state = random.choices(range(len(transitions)), weights=transitions[state.market_state], k=1)[0]
     await state_store.save(GLOBAL_STATE_KEY, state)
-
     multiplier = multipliers[state.market_state]
     async for (wallet_id, wallet) in wallet_store.all():
-        wallet.investment_value *= multiplier
+        wallet.holdings *= multiplier
         await wallet_store.save(wallet_id, wallet)
 
 
@@ -155,22 +152,27 @@ async def bet(
         amount = amount(sender_wallet.balance)
 
     if sender_wallet.balance == 0:
-        return "Sorry! You don't have any rollcoins right now - try mining!"
+        yield "Sorry! You don't have any rollcoins right now - try mining!"
+        return
 
     if amount > sender_wallet.balance:
-        return f"Sorry! You only have {sender_wallet.balance} RollCoins available - try mining for more!"
+        yield f"Sorry! You only have {sender_wallet.balance} RollCoins available - try mining for more!"
+        return
 
     if amount <= 0:
         RollbotFailure.INVALID_ARGUMENTS.raise_exc(f"Amount must be positive, not {amount}")
 
     sender_wallet.balance -= amount
-    sender_wallet.invested += amount
-    sender_wallet.investment_value += amount
+    sender_wallet.holdings += amount
+    sender_wallet.cost_basis += amount
     await wallet_store.save(sender_id, sender_wallet)
 
+    # coins enter treasury, allow evolve_market to save state
+    state.treasury += amount
     await evolve_market(state, transitions, multipliers, wallet_store, state_store)
-    
-    return f"{messages[state.market_state]}\n{await wallet_store.load(sender_id)}", reply
+
+    yield f"Trade complete!\n{await wallet_store.load(sender_id)}", reply
+    yield f"Market status: {messages[state.market_state]}"
 
 
 @as_command
@@ -188,42 +190,37 @@ async def cash(
 ):
     if not isinstance(amount, float):
         # handle special converters
-        amount = amount(sender_wallet.investment_value)
+        amount = amount(sender_wallet.holdings)
 
-    if sender_wallet.investment_value == 0:
-        return "Sorry! You don't have any rollcoins invested right now - try betting!"
+    if sender_wallet.holdings == 0:
+        yield "Sorry! You don't have any rollcoins invested right now - try betting!"
+        return
 
-    if amount > sender_wallet.investment_value:
-        return f"Sorry! You only have {sender_wallet.investment_value} RollCoins in the market - try betting more first!"
+    if amount > sender_wallet.holdings:
+        yield f"Sorry! You only have {sender_wallet.holdings} RollCoins in the market - try betting more first!"
+        return
 
     if amount <= 0:
         RollbotFailure.INVALID_ARGUMENTS.raise_exc(f"Amount must be positive, not {amount}")
 
-    if amount <= sender_wallet.invested:
-        sender_wallet.balance += amount
-        sender_wallet.invested -= amount
-        sender_wallet.investment_value -= amount
-        await wallet_store.save(sender_id, sender_wallet)
+    if state.treasury == 0:
+        yield f"Sorry! The treasury is actually empty right now so uh... no one can sell..."
+        return
+
+    if amount <= state.treasury:
         response = "Successfully sold!"
-    elif amount <= sender_wallet.invested + state.treasury:
-        diff = amount - sender_wallet.invested
-        sender_wallet.balance += amount
-        sender_wallet.invested = 0
-        sender_wallet.investment_value -= amount
-        await wallet_store.save(sender_id, sender_wallet)
-        state.treasury -= diff
-        # market evolution will save state
-        response = f"Successfully sold! Took {diff} RollCoins from the treasury."
+        actual_amount = amount
     else:
-        given = sender_wallet.invested + state.treasury
-        sender_wallet.balance += given
-        sender_wallet.invested = 0
-        sender_wallet.investment_value -= given
-        await wallet_store.save(sender_id, sender_wallet)
-        state.treasury = 0
-        # market evolution will save state
-        response = f"Funny story, I'm actually out of coins! I gave you what I could, draining the treasury, which was {given}."
+        response = f"Funny story, I'm actually out of coins! I gave you what I could, draining the treasury, which was {state.treasury}."
+        actual_amount = state.treasury
 
+    sender_wallet.balance += actual_amount
+    sender_wallet.holdings -= actual_amount
+    sender_wallet.cost_basis = max(sender_wallet.cost_basis - actual_amount, 0)
+    await wallet_store.save(sender_id, sender_wallet)
+    # coins exit treasury, allow evolve_market to save state
+    state.treasury -= actual_amount
     await evolve_market(state, transitions, multipliers, wallet_store, state_store)
-    
-    return f"{response}\n{messages[state.market_state]}\n{await wallet_store.load(sender_id)}", reply
+
+    yield f"{response}\n{await wallet_store.load(sender_id)}", reply
+    yield f"Market status: {messages[state.market_state]}"