Update refid grouping

1 files changed, 40 insertions, 23 deletions

diff --git a/bla.py b/bla.py
index 3d1cfb7..8bc6e26 100644
--- a/bla.py
+++ b/bla.py
@@ -1,12 +1,12 @@
 import csv
-from collections import deque
+from collections import defaultdict, deque
 
 class FIFOQueue:
     def __init__(self):
         self.queue = deque()
 
-    def add(self, amount, price):
-        self.queue.append((amount, price))
+    def add(self, amount, cost):
+        self.queue.append((amount, cost))
 
     def remove(self, amount):
         remaining = amount
@@ -15,47 +15,64 @@ class FIFOQueue:
         while remaining > 0:
             if not self.queue:
                 raise ValueError(f"Insufficient assets in queue to process sale of {amount}.")
-            quantity, price = self.queue[0]
+            quantity, cost = self.queue[0]
             if quantity > remaining:
-                cost_basis += remaining * price
-                self.queue[0] = (quantity - remaining, price)
+                cost_basis += remaining * cost
+                self.queue[0] = (quantity - remaining, cost)
                 remaining = 0
             else:
-                cost_basis += quantity * price
+                cost_basis += quantity * cost
                 remaining -= quantity
                 self.queue.popleft()
 
         return cost_basis
 
 def process_ledger(file_path):
-    fifo_queues = {}  # Separate FIFO queue per currency
+    fifo_queues = {}  # Separate FIFO queue per cryptocurrency
+    trades_by_refid = defaultdict(list)
     report = []
 
     with open(file_path, 'r') as file:
         reader = csv.DictReader(file)
         for row in reader:
-            currency = row["asset"]
-            fifo_queues.setdefault(currency, FIFOQueue())
+            # Group trades by refid
+            if row["type"] == "trade":
+                trades_by_refid[row["refid"]].append(row)
 
             # Handle deposits
-            if row["type"] == "deposit":
+            elif row["type"] == "deposit":
+                currency = row["asset"]
+                fifo_queues.setdefault(currency, FIFOQueue())
                 amount = float(row["amount"])
                 price = 0  # Deposits typically have no associated cost basis
                 fifo_queues[currency].add(amount, price)
 
-            # Handle trades - sales
-            elif row["type"] == "trade" and row["amount"].startswith('-'):
-                amount = -float(row["amount"])
-                cost_basis = fifo_queues[currency].remove(amount)
-                sale_proceeds = amount * float(row["fee"])  # Adjust fee as needed
-                profit_or_loss = sale_proceeds - cost_basis
-                report.append((row["time"], currency, profit_or_loss))
+        # Process grouped trades
+        for refid, trades in trades_by_refid.items():
+            if len(trades) == 2:  # Ensure we have two related rows (EUR + crypto)
+                eur_trade = next((trade for trade in trades if trade["asset"] == "EUR"), None)
+                crypto_trade = next((trade for trade in trades if trade["asset"] != "EUR"), None)
 
-            # Handle trades - purchases
-            elif row["type"] == "trade" and not row["amount"].startswith('-'):
-                amount = float(row["amount"])
-                price = float(row["fee"])  # Adjust fee as needed
-                fifo_queues[currency].add(amount, price)
+                if eur_trade and crypto_trade:
+                    crypto_asset = crypto_trade["asset"]
+                    eur_amount = float(eur_trade["amount"])
+                    crypto_amount = float(crypto_trade["amount"])
+                    fee = float(crypto_trade["fee"])
+                    fifo_queues.setdefault(crypto_asset, FIFOQueue())
+
+                    if eur_amount < 0:  # Purchase of cryptocurrency
+                        stake_amount = -eur_amount
+                        crypto_amount -= fee  # Adjust for fees
+                        fifo_queues[crypto_asset].add(crypto_amount, stake_amount)
+                    elif eur_amount > 0:  # Sale of cryptocurrency
+                        cost_basis = fifo_queues[crypto_asset].remove(crypto_amount)
+                        sale_proceeds = eur_amount
+                        profit_or_loss = sale_proceeds - cost_basis
+                        report.append((eur_trade["time"], crypto_asset, profit_or_loss))
+                else:
+                    raise ValueError(f"Unexpected trade grouping for refid {refid}")
+            else:
+                raise ValueError(f"Unexpected number of trades for refid {refid}")
 
     return report