Update app/smart_recommendation.py

app/smart_recommendation.py

@@ -262,6 +262,28 @@ class SmartBudgetRecommender:
         result.sort(key=lambda x: x.average_monthly_expense, reverse=True)
         return result
 
+    def _get_category_name(self, category_id) -> str:
+        """Look up category name from categories collection"""
+        if not category_id:
+            return "Uncategorized"
+        
+        try:
+            # Try to find category in categories collection
+            if isinstance(category_id, ObjectId):
+                category_doc = self.db.categories.find_one({"_id": category_id})
+            else:
+                try:
+                    category_doc = self.db.categories.find_one({"_id": ObjectId(category_id)})
+                except:
+                    category_doc = self.db.categories.find_one({"_id": category_id})
+            
+            if category_doc:
+                return category_doc.get("name") or category_doc.get("title") or str(category_id)
+        except Exception:
+            pass
+        
+        return str(category_id) if category_id else "Uncategorized"
+    
     def _get_category_stats_from_budgets(
         self, user_id: str, month: int, year: int
     ) -> Dict:
@@ -270,40 +292,41 @@ class SmartBudgetRecommender:
 
         We treat each budget document (e.g. \"Office Maintenance\", \"LOGICGO\")
         as a spending category and derive an \"average\" from its amounts.
+        Also extracts categories from headCategories array.
         """
         budgets = []
         
-        # Try multiple query patterns to find budgets
-        # Pattern 1: Try with ObjectId (most common in WalletSync)
+        # Try multiple query patterns to find budgets (include both OPEN and CLOSE status)
+        # Pattern 1: Try with ObjectId (most common in WalletSync) - no status filter
         try:
-            query_objid = {"status": "OPEN", "createdBy": ObjectId(user_id)}
+            query_objid = {"createdBy": ObjectId(user_id)}
             budgets_objid = list(self.db.budgets.find(query_objid))
             if budgets_objid:
                 budgets.extend(budgets_objid)
         except Exception:
             pass
         
-        # Pattern 2: Try with string user_id
+        # Pattern 2: Try with string user_id - no status filter
         try:
-            query_str = {"status": "OPEN", "createdBy": user_id}
+            query_str = {"createdBy": user_id}
             budgets_str = list(self.db.budgets.find(query_str))
             if budgets_str:
                 budgets.extend(budgets_str)
         except Exception:
             pass
         
-        # Pattern 3: Try with user_id field (alternative field name)
+        # Pattern 3: Try with user_id field (alternative field name) - no status filter
         try:
-            query_userid = {"status": "OPEN", "user_id": user_id}
+            query_userid = {"user_id": user_id}
             budgets_userid = list(self.db.budgets.find(query_userid))
             if budgets_userid:
                 budgets.extend(budgets_userid)
         except Exception:
             pass
         
-        # Pattern 4: Try ObjectId with user_id field
+        # Pattern 4: Try ObjectId with user_id field - no status filter
         try:
-            query_objid_userid = {"status": "OPEN", "user_id": ObjectId(user_id)}
+            query_objid_userid = {"user_id": ObjectId(user_id)}
             budgets_objid_userid = list(self.db.budgets.find(query_objid_userid))
             if budgets_objid_userid:
                 budgets.extend(budgets_objid_userid)
@@ -329,14 +352,85 @@ class SmartBudgetRecommender:
 
         result: Dict[str, Dict] = {}
         for b in budgets:
-            # Use budget "name" as category label
-            category = b.get("name", "Uncategorized")
-            if not category or category == "Uncategorized":
-                # Try alternative field names
-                category = b.get("category") or b.get("title") or "Uncategorized"
+            # First, try to extract categories from headCategories array
+            head_categories = b.get("headCategories", [])
+            
+            if head_categories and isinstance(head_categories, list):
+                # Process nested categories from headCategories
+                for head_cat in head_categories:
+                    if not isinstance(head_cat, dict):
+                        continue
+                    
+                    # Get headCategory ID and amounts
+                    head_cat_id = head_cat.get("headCategory")
+                    head_cat_max = float(head_cat.get("maxAmount", 0) or 0)
+                    head_cat_spend = float(head_cat.get("spendAmount", 0) or 0)
+                    
+                    # Process nested categories within headCategory
+                    nested_categories = head_cat.get("categories", [])
+                    if nested_categories and isinstance(nested_categories, list):
+                        for nested_cat in nested_categories:
+                            if not isinstance(nested_cat, dict):
+                                continue
+                            
+                            nested_cat_id = nested_cat.get("category")
+                            nested_cat_max = float(nested_cat.get("maxAmount", 0) or 0)
+                            nested_cat_spend = float(nested_cat.get("spendAmount", 0) or 0)
+                            spend_limit_type = nested_cat.get("spendLimitType", "NO_LIMIT")
+                            
+                            # Only include categories with limits (maxAmount > 0 or spendLimitType != "NO_LIMIT")
+                            if nested_cat_max > 0 or (spend_limit_type != "NO_LIMIT" and nested_cat_spend > 0):
+                                # Look up actual category name
+                                nested_category_name = self._get_category_name(nested_cat_id)
+                                nested_base_amount = nested_cat_spend if nested_cat_spend > 0 else nested_cat_max
+                                
+                                if nested_category_name not in result:
+                                    result[nested_category_name] = {
+                                        "average_monthly": nested_base_amount,
+                                        "total": nested_base_amount,
+                                        "count": 1,
+                                        "months_analyzed": 1,
+                                        "std_dev": 0.0,
+                                        "monthly_values": [nested_base_amount],
+                                    }
+                                else:
+                                    result[nested_category_name]["total"] += nested_base_amount
+                                    result[nested_category_name]["count"] += 1
+                                    result[nested_category_name]["months_analyzed"] = result[nested_category_name]["count"]
+                                    result[nested_category_name]["average_monthly"] = (
+                                        result[nested_category_name]["total"] / result[nested_category_name]["count"]
+                                    )
+                                    result[nested_category_name]["monthly_values"].append(nested_base_amount)
+                    
+                    # Also include headCategory if it has amounts
+                    if head_cat_max > 0 or head_cat_spend > 0:
+                        head_category_name = self._get_category_name(head_cat_id)
+                        head_base_amount = head_cat_spend if head_cat_spend > 0 else head_cat_max
+                        
+                        if head_category_name not in result:
+                            result[head_category_name] = {
+                                "average_monthly": head_base_amount,
+                                "total": head_base_amount,
+                                "count": 1,
+                                "months_analyzed": 1,
+                                "std_dev": 0.0,
+                                "monthly_values": [head_base_amount],
+                            }
+                        else:
+                            result[head_category_name]["total"] += head_base_amount
+                            result[head_category_name]["count"] += 1
+                            result[head_category_name]["months_analyzed"] = result[head_category_name]["count"]
+                            result[head_category_name]["average_monthly"] = (
+                                result[head_category_name]["total"] / result[head_category_name]["count"]
+                            )
+                            result[head_category_name]["monthly_values"].append(head_base_amount)
+            
+            # Also include the main budget as a category (if it has amounts)
+            budget_name = b.get("name", "Uncategorized")
+            if not budget_name or budget_name == "Uncategorized":
+                budget_name = b.get("category") or b.get("title") or "Uncategorized"
 
             # Derive a base amount from WalletSync fields
-            # Try multiple field name variations
             max_amount = float(b.get("maxAmount", 0) or b.get("max_amount", 0) or b.get("amount", 0) or 0)
             spend_amount = float(b.get("spendAmount", 0) or b.get("spend_amount", 0) or b.get("spent", 0) or 0)
             budget_amount = float(b.get("budget", 0) or b.get("budgetAmount", 0) or 0)
@@ -349,26 +443,27 @@ class SmartBudgetRecommender:
             elif budget_amount > 0:
                 base_amount = budget_amount
             else:
-                continue  # Skip if no valid amount found
-
-            if category not in result:
-                result[category] = {
-                    "average_monthly": base_amount,
-                    "total": base_amount,
-                    "count": 1,
-                    "months_analyzed": 1,
-                    "std_dev": 0.0,
-                    "monthly_values": [base_amount],
-                }
-            else:
-                # If multiple budgets per category, average them
-                result[category]["total"] += base_amount
-                result[category]["count"] += 1
-                result[category]["months_analyzed"] = result[category]["count"]
-                result[category]["average_monthly"] = (
-                    result[category]["total"] / result[category]["count"]
-                )
-                result[category]["monthly_values"].append(base_amount)
+                base_amount = 0
+
+            # Only add main budget if it has an amount and we haven't processed categories
+            if base_amount > 0:
+                if budget_name not in result:
+                    result[budget_name] = {
+                        "average_monthly": base_amount,
+                        "total": base_amount,
+                        "count": 1,
+                        "months_analyzed": 1,
+                        "std_dev": 0.0,
+                        "monthly_values": [base_amount],
+                    }
+                else:
+                    result[budget_name]["total"] += base_amount
+                    result[budget_name]["count"] += 1
+                    result[budget_name]["months_analyzed"] = result[budget_name]["count"]
+                    result[budget_name]["average_monthly"] = (
+                        result[budget_name]["total"] / result[budget_name]["count"]
+                    )
+                    result[budget_name]["monthly_values"].append(base_amount)
 
         print(f"Processed {len(result)} budget categories for recommendations")
         return result
@@ -423,7 +518,6 @@ class SmartBudgetRecommender:
             return None
 
 
-
 # import json
 # import math
 # import os