Work on interest system

database/crud/currency.py

@@ -1,25 +1,29 @@
 from datetime import date
-from typing import Optional, Union
+from typing import List, Optional, Union
 
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.crud import users
 from database.exceptions import currency as exceptions
-from database.schemas import Bank, NightlyData
+from database.schemas import Bank, BankSavings, NightlyData
 from database.utils.math.currency import (
     capacity_upgrade_price,
+    interest_rate,
     interest_upgrade_price,
     rob_upgrade_price,
+    savings_cap,
 )
 
 __all__ = [
     "add_dinks",
+    "apply_daily_interest",
     "claim_nightly",
     "deduct_dinks",
     "gamble_dinks",
     "get_bank",
     "get_nightly_data",
-    "invest",
+    "save",
     "upgrade_capacity",
     "upgrade_interest",
     "upgrade_rob",
@@ -36,15 +40,30 @@ async def get_bank(session: AsyncSession, user_id: int) -> Bank:
     return user.bank
 
 
+async def get_savings(session: AsyncSession, user_id: int) -> BankSavings:
+    """Get a user's savings info"""
+    user = await users.get_or_add_user(session, user_id)
+    return user.savings
+
+
 async def get_nightly_data(session: AsyncSession, user_id: int) -> NightlyData:
     """Get a user's nightly info"""
     user = await users.get_or_add_user(session, user_id)
     return user.nightly_data
 
 
-async def invest(session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None) -> int:
+async def save(
+    session: AsyncSession,
+    user_id: int,
+    amount: Union[str, int],
+    *,
+    bank: Optional[Bank] = None,
+    savings: Optional[BankSavings] = None
+) -> int:
     """Invest some of your Dinks"""
     bank = bank or await get_bank(session, user_id)
+    savings = savings or await get_savings(session, user_id)
+
     if amount == "all":
         amount = bank.dinks
 
@@ -54,29 +73,44 @@ async def invest(session: AsyncSession, user_id: int, amount: Union[str, int], *
     # Don't allow investing more dinks than you own
     amount = min(bank.dinks, int(amount))
 
+    # Don't allow exceeding the limit
+    limit = savings_cap(bank.capacity_level)
+
+    if savings.saved >= limit:
+        raise exceptions.SavingsCapExceeded
+
+    if savings.saved + amount > limit:
+        amount = max(0, limit - savings.saved)
+
     bank.dinks -= amount
-    bank.invested += amount
+    savings.saved += amount
 
     session.add(bank)
+    session.add(savings)
     await session.commit()
 
     return amount
 
 
 async def withdraw(session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None) -> int:
-    """Withdraw your invested Dinks"""
+    """Withdraw your saved Dinks"""
     bank = bank or await get_bank(session, user_id)
+    savings = await get_savings(session, user_id)
+
     if amount == "all":
-        amount = bank.invested
+        amount = savings.saved
 
     # Don't allow withdrawing more dinks than you own
-    amount = min(bank.invested, int(amount))
+    amount = min(savings.saved, int(amount))
 
     bank.dinks += amount
-    bank.invested -= amount
+    savings.saved -= amount
+    savings.daily_minimum = min(savings.daily_minimum, savings.saved)
 
     session.add(bank)
+    session.add(savings)
     await session.commit()
+
     return amount
 
 
@@ -218,3 +252,22 @@ async def rob(
     session.add(robber)
     session.add(robbed)
     await session.commit()
+
+
+async def apply_daily_interest(session: AsyncSession):
+    """Apply daily interest rates to all accounts with saved Dinks"""
+    statement = select(BankSavings)
+    all_savings: List[BankSavings] = list((await session.execute(statement)).scalars().all())
+
+    for savings_account in all_savings:
+        if savings_account.saved == 0:
+            continue
+
+        bank = await get_bank(session, savings_account.user_id)
+        rate = interest_rate(bank.interest_level)
+
+        savings_account.saved = float(savings_account.saved * rate)
+        savings_account.daily_minimum = savings_account.saved
+        session.add(savings_account)
+
+    await session.commit()

database/crud/users.py

@@ -3,7 +3,7 @@ from typing import Optional
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.schemas import Bank, NightlyData, User
+from database.schemas import Bank, BankSavings, NightlyData, User
 
 __all__ = [
     "get_or_add_user",
@@ -37,8 +37,12 @@ async def get_or_add_user(session: AsyncSession, user_id: int, *, options: Optio
     user.bank = bank
     user.nightly_data = nightly_data
 
+    savings = BankSavings(user_id=user_id)
+    user.savings = savings
+
     session.add(bank)
     session.add(nightly_data)
+    session.add(savings)
     session.add(user)
 
     await session.commit()

database/exceptions/currency.py

@@ -7,3 +7,7 @@ class DoubleNightly(Exception):
 
 class NotEnoughDinks(Exception):
     """Exception raised when trying to do something you don't have the Dinks for"""
+
+
+class SavingsCapExceeded(Exception):
+    """Exception raised when trying to save more Dinks than the cap allows"""

database/schemas.py

@@ -12,6 +12,7 @@ from database import enums
 __all__ = [
     "Base",
     "Bank",
+    "BankSavings",
     "Birthday",
     "Bookmark",
     "CommandStats",
@@ -52,7 +53,6 @@ class Bank(Base):
     user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id"))
 
     dinks: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False)
-    invested: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False)
 
     # Interest rate
     interest_level: Mapped[int] = mapped_column(server_default="1", nullable=False)
@@ -66,6 +66,20 @@ class Bank(Base):
     user: Mapped[User] = relationship(uselist=False, back_populates="bank", lazy="selectin")
 
 
+class BankSavings(Base):
+    """Savings information for a user's bank"""
+
+    __tablename__ = "savings"
+
+    savings_id: Mapped[int] = mapped_column(primary_key=True)
+    user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id"))
+
+    saved: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False)
+    daily_minimum: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False)
+
+    user: Mapped[User] = relationship(uselist=False, back_populates="savings", lazy="selectin")
+
+
 class Birthday(Base):
     """A user's birthday"""
 
@@ -350,3 +364,6 @@ class User(Base):
     reminders: Mapped[List[Reminder]] = relationship(
         back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan"
     )
+    savings: Mapped[List[BankSavings]] = relationship(
+        back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan"
+    )

database/utils/math/currency.py

@@ -4,6 +4,8 @@ __all__ = [
     "capacity_upgrade_price",
     "interest_upgrade_price",
     "rob_upgrade_price",
+    "interest_rate",
+    "savings_cap",
     "jail_chance",
     "jail_time",
     "rob_amount",
@@ -16,7 +18,7 @@ def interest_upgrade_price(level: int) -> int:
     base_cost = 600
     growth_rate = 1.8
 
-    return math.floor(base_cost * (growth_rate**level))
+    return math.floor(base_cost * (growth_rate ** (level - 1)))
 
 
 def capacity_upgrade_price(level: int) -> int:
@@ -24,7 +26,7 @@ def capacity_upgrade_price(level: int) -> int:
     base_cost = 800
     growth_rate = 1.6
 
-    return math.floor(base_cost * (growth_rate**level))
+    return math.floor(base_cost * (growth_rate ** (level - 1)))
 
 
 def rob_upgrade_price(level: int) -> int:
@@ -32,7 +34,23 @@ def rob_upgrade_price(level: int) -> int:
     base_cost = 950
     growth_rate = 1.9
 
-    return math.floor(base_cost * (growth_rate**level))
+    return math.floor(base_cost * (growth_rate ** (level - 1)))
+
+
+def interest_rate(level: int) -> float:
+    """Calculate the amount of interest you will receive"""
+    base_rate = 1.025
+    growth_rate = 0.03
+
+    return base_rate + (growth_rate * (level - 1))
+
+
+def savings_cap(level: int) -> int:
+    """Calculate the maximum amount you can save"""
+    base_limit = 1000
+    growth_rate = 1.10
+
+    return math.floor(base_limit * (growth_rate ** (level - 1)))
 
 
 def jail_chance(level: int) -> float: