From bee908180c0ccd6aaab0bc9e8070b0e77ab968e9 Mon Sep 17 00:00:00 2001
From: "Bobby (aider)" <bobby@abellana.work>
Date: Tue, 11 Feb 2025 19:16:52 -0800
Subject: [PATCH] feat: Add functionality to unlink trades from trading plans

---
 src/streamlit_app.py        | 41 +++++++++++++++++++++++++++++++++++++
 src/trading/trading_plan.py | 19 +++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/src/streamlit_app.py b/src/streamlit_app.py
index bf0b0a2..4acb55b 100644
--- a/src/streamlit_app.py
+++ b/src/streamlit_app.py
@@ -1166,7 +1166,48 @@ def trading_plan_page():
                                         pl = (trade['exit_price'] - trade['entry_price']) * trade['shares']
                                         st.write(f"Exit: ${trade['exit_price']:.2f}")
                                         st.write(f"P/L: ${pl:.2f}")
+                                
+                                # Add unlink button for each trade
+                                if st.button("Unlink Trade", key=f"unlink_trade_{trade['id']}"):
+                                    try:
+                                        # Update the single trade
+                                        query = """
+                                            ALTER TABLE stock_db.trades
+                                            UPDATE plan_id = NULL
+                                            WHERE id = %(trade_id)s
+                                        """
+                                        with create_client() as client:
+                                            client.command(query, {'trade_id': trade['id']})
+                                        
+                                        # Recalculate metrics
+                                        metrics = calculate_plan_metrics(plan.id)
+                                        plan.win_rate = metrics['win_rate']
+                                        plan.average_return_per_trade = metrics['average_return']
+                                        plan.profit_factor = metrics['profit_factor']
+                                        update_trading_plan(plan)
+                                        
+                                        st.success(f"Trade unlinked successfully!")
+                                        st.query_params.update(rerun=True)
+                                    except Exception as e:
+                                        st.error(f"Error unlinking trade: {str(e)}")
                     
+                        # Add button to unlink all trades
+                        if st.button("Unlink All Trades", key=f"unlink_all_trades_{plan.id}"):
+                            try:
+                                if unlink_trades_from_plan(plan.id):
+                                    # Reset metrics
+                                    plan.win_rate = None
+                                    plan.average_return_per_trade = None
+                                    plan.profit_factor = None
+                                    update_trading_plan(plan)
+                                    
+                                    st.success("All trades unlinked successfully!")
+                                    st.query_params.update(rerun=True)
+                                else:
+                                    st.error("Error unlinking trades")
+                            except Exception as e:
+                                st.error(f"Error unlinking trades: {str(e)}")
+
                     # Get available trades
                     with create_client() as client:
                         query = """
diff --git a/src/trading/trading_plan.py b/src/trading/trading_plan.py
index 6071742..c31c890 100644
--- a/src/trading/trading_plan.py
+++ b/src/trading/trading_plan.py
@@ -312,10 +312,29 @@ def get_all_trading_plans(status: Optional[PlanStatus] = None) -> List[TradingPl
             print(f"Error retrieving trading plans: {e}")
             return []
 
+def unlink_trades_from_plan(plan_id: int) -> bool:
+    """Unlink all trades from a trading plan"""
+    with create_client() as client:
+        try:
+            query = """
+                ALTER TABLE stock_db.trades
+                UPDATE plan_id = NULL
+                WHERE plan_id = %(plan_id)s
+            """
+            client.command(query, {'plan_id': plan_id})
+            return True
+        except Exception as e:
+            print(f"Error unlinking trades from plan: {e}")
+            return False
+
 def delete_trading_plan(plan_id: int) -> bool:
     """Delete a trading plan by ID"""
     with create_client() as client:
         try:
+            # First unlink all trades
+            unlink_trades_from_plan(plan_id)
+            
+            # Then delete the plan
             query = "ALTER TABLE trading_plans DELETE WHERE id = %(id)s"
             client.command(query, {'id': plan_id})
             return True