unfuck the calculator

staking/stakecalc2.py

@@ -20,16 +20,26 @@ class StakingPool:
         total_fee = amount * self.fee_rate
 
         # Distribute fee
-        # TODO: Handle edge case where pool has no initial users == Community funds gets the reward
         # 1. Community funds
         community_fee = total_fee * self.community_pct
+        self.community_fund += community_fee
 
         # 2. Developer funds
         developer_fee = total_fee * self.developer_pct
+        self.developer_fund += developer_fee
 
         # 3. Reward stakers
+        # If there are no users in the pool, the reward is placed in the community fund
         reward_amount = total_fee - community_fee - developer_fee
-        self.__distribute_reward(reward_amount)
+
+        ########
+        # Test no community fund (distribute equally)
+        self.community_distribute()
+
+        if len(self.user_balances) < 1:
+            self.community_fund += reward_amount
+        else:
+            self.__distribute_reward(reward_amount)
 
         # Deposit amount minus fee to contract_fund (wallet)
         self.contract_fund += amount - total_fee
@@ -42,19 +52,40 @@ class StakingPool:
 
     def unstake(self, user: str):
         # Fetch amount from user_balances
+        if user not in self.user_balances:
+            return 0.0
+        elif self.user_balances[user] <= 0:
+            return 0.0
+
+        user_balance = self.user_balances[user]
 
         # Remove user from pool to prevent self-reward leaving a hanging balance
+        del self.user_balances[user]
 
         # Calculate fee
+        total_fee = user_balance * self.fee_rate
+        # Adjust working balance
+        user_balance -= total_fee
 
         # Distribute fee
-        # TODO: User needs to be removed from the pool first to avoid hanging balances after rewarding stakers
         # 1. Community funds
+        community_fee = total_fee * self.community_pct
+        self.community_fund += community_fee
+
+        #########
+        # TEST DISTRIBUTE TO COMMUNITY
+        self.community_distribute()
+
         # 2. Developer funds
+        developer_fee = total_fee * self.developer_pct
+        self.developer_fund += developer_fee
+
         # 3. Reward stakers
+        reward_amount = total_fee - community_fee - developer_fee
+        self.__distribute_reward(reward_amount)
 
         # Transfer tokens from contract_funds (wallet) to user's wallet address
-        pass
+        return user_balance
 
     def dev_withdrawal(self, destination: str, amount: float):
         pass
@@ -62,23 +93,51 @@ class StakingPool:
     def community_withdrawal(self, destination: str, amount: float):
         pass
 
-    def __distribute_reward(self, amount):
-        # # # #
-        # Lock user_balances mutex
-        # # # #
+    def community_distribute(self):
+        # Take all funds in community stash and distribute it to stakers
+        amount = self.community_fund
+        self.__distribute_reward(amount)
+        self.community_fund = 0.0
 
+    def __distribute_reward(self, amount):
         # Get total staked amount
+        staked_total = sum(self.user_balances.values())
+
         # Determine distribution amongst users
         # Update self.user_balances
+        for user, user_stake in self.user_balances.items():
+            pct = user_stake / staked_total
+            user_reward = pct * amount
+            self.user_balances[user] += user_reward
 
-        # # # #
-        # Unlock user_balances mutex
-        # # # #
-        pass
+    def stats(self):
+        print('\n--- POOL STATS ---')
+        print(f'staked:         {sum(self.user_balances.values())}')
+        print(f'community fund: {self.community_fund}')
+        print(f'developer fund: {self.developer_fund}')
+        print('------------------')
+        print('--- USER TOTALS ---')
+        for user, balance in self.user_balances.items():
+            print(f'{user}: {balance}')
+        print('------------------')
 
 
 def main():
-    pass
+    pool = StakingPool()
+    pool.stake("owner", 50000)
+    pool.stake("user1", 50000000)
+    # pool.stake("user2", 5000)
+    # pool.stake("user3", 50000)
+    # pool.stake("user3", 50000)
+    # pool.stake("user3", 50000)
+    # pool.stake("user3", 50000)
+    # pool.stake("user3", 50000)
+    # pool.stake("user3", 50000)
+    pool.stats()
+    # pool.community_distribute()
+
+    print(pool.unstake("owner"))
+    pool.stats()
 
 
 if __name__ == '__main__':