Update main.py

main.py

@@ -93,28 +93,58 @@ def marketing_rec():
 
     return jsonify(str(response['text']))
 
-# Income/Expenses Prediction endpoint
-@app.route("/predict_revenue", methods=["POST"])
+# Profit/Customer Engagement Prediction endpoint
+@app.route("/predict_metric", methods=["POST"])
 @cross_origin()
-def predict_revenue():
+def predict_metric():
     request_data = request.json
     user_id = request_data.get("user_id")
     interval = request_data.get("interval", 30)
-    transaction_type = request_data.get("transaction_type", "Income") 
+    metric_type = request_data.get("metric_type", "Profit")  # "Profit" or "Customer Engagement"
 
-    # Fetch transaction data based on user and transaction type
     transactions_ref = db.collection("system_users").document(user_id).collection("transactions")
-    query = transactions_ref.where("transactionType", "==", transaction_type).stream()
-    
+
     data = []
-    for doc in query:
-        transaction = doc.to_dict()
-        data.append({
-            "date": transaction["date"], 
-            "amountDue": transaction["amountDue"]
-        })
-
-    # Create DataFrame from transaction data
+    
+    if metric_type == "Profit":
+        # Fetch both Income and Expense transactions for Profit calculation
+        income_query = transactions_ref.where("transactionType", "==", "Income").stream()
+        expense_query = transactions_ref.where("transactionType", "==", "Expense").stream()
+
+        income_data = {}
+        expense_data = {}
+
+        for doc in income_query:
+            transaction = doc.to_dict()
+            date_str = transaction["date"]
+            amount = transaction["amountDue"]
+            income_data[date_str] = income_data.get(date_str, 0) + amount
+
+        for doc in expense_query:
+            transaction = doc.to_dict()
+            date_str = transaction["date"]
+            amount = transaction["amountDue"]
+            expense_data[date_str] = expense_data.get(date_str, 0) + amount
+
+        # Calculate net profit for each date
+        for date, income in income_data.items():
+            expense = expense_data.get(date, 0)
+            data.append({"date": date, "amountDue": income - expense})
+
+    elif metric_type == "Customer Engagement":
+        # Use count of Income transactions per day as Customer Engagement
+        income_query = transactions_ref.where("transactionType", "==", "Income").stream()
+
+        engagement_data = {}
+        for doc in income_query:
+            transaction = doc.to_dict()
+            date_str = transaction["date"]
+            engagement_data[date_str] = engagement_data.get(date_str, 0) + 1
+
+        for date, count in engagement_data.items():
+            data.append({"date": date, "amountDue": count})
+
+    # Create DataFrame from the aggregated data
     df = pd.DataFrame(data)
 
     # Ensure 'date' column is datetime
@@ -124,7 +154,7 @@ def predict_revenue():
     # Set 'date' as index
     df = df.sort_values("date").set_index("date")
 
-    # Resample daily to ensure regular intervals
+    # Resample daily to ensure regular intervals (fill missing dates)
     df = df.resample("D").sum().reset_index()
 
     df.columns = ["ds", "y"]  # ds: date, y: target
@@ -137,17 +167,18 @@ def predict_revenue():
     model = Prophet(daily_seasonality=True, yearly_seasonality=True)
     model.fit(df)
 
-    # dataframe for future predictions
+    # DataFrame for future predictions
     future_dates = model.make_future_dataframe(periods=interval)
     forecast = model.predict(future_dates)
 
     # Extract the forecast for the requested interval
-    forecast_data = forecast[['ds', 'yhat']].tail(interval)  
+    forecast_data = forecast[['ds', 'yhat']].tail(interval)
     predictions = [{"date": row['ds'].strftime('%Y-%m-%d'), "value": row['yhat']} for _, row in forecast_data.iterrows()]
 
     # Return predictions in JSON format
     return jsonify({"predictedData": predictions})
 
 
+
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)