Update main.py

main.py

@@ -94,98 +94,92 @@ def marketing_rec():
 
     return jsonify(str(response['text']))
 
+
+# Profit/Customer Engagement Prediction endpoint
 @app.route("/predict_metric", methods=["POST"])
 @cross_origin()
 def predict_metric():
-    try:
-        request_data = request.json
-        logging.info(f"Received request data: {request_data}")
-
-        user_id = request_data.get("user_id")
-        interval = request_data.get("interval", 30)
-        metric_type = request_data.get("metric_type", "Profit")  # "Profit" or "Customer Engagement"
-
-        # Check if user_id is provided
-        if not user_id:
-            logging.error("User ID is missing in the request")
-            return jsonify({"error": "User ID is missing"}), 400
-
-        transactions_ref = db.collection("system_users").document(user_id).collection("transactions")
-        data = []
-        
-        # Fetch Income and Expense for Profit calculation
-        if metric_type == "Profit":
-            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()
-                logging.info(f"Processing income transaction: {transaction}")
-                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()
-                logging.info(f"Processing expense transaction: {transaction}")
-                date_str = transaction["date"]
-                amount = transaction["amountDue"]
-                expense_data[date_str] = expense_data.get(date_str, 0) + amount
-
-            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":
-            income_query = transactions_ref.where("transactionType", "==", "Income").stream()
-            engagement_data = {}
-
-            for doc in income_query:
-                transaction = doc.to_dict()
-                logging.info(f"Processing engagement transaction: {transaction}")
-                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})
-
-        # Log final aggregated data before processing
-        logging.info(f"Aggregated data: {data}")
-
-        # Data processing with Prophet
-        df = pd.DataFrame(data)
-        df['date'] = pd.to_datetime(df['date'])
-        df['date'] = df['date'].dt.tz_localize(None)
-        df = df.sort_values("date").set_index("date")
-        df = df.resample("D").sum().reset_index()
-        df.columns = ["ds", "y"]
-
-        # Check if data is sufficient
-        if df.shape[0] < 10:
-            logging.warning("Not enough data to train the Prophet model")
-            return jsonify({"error": "Not enough data for prediction"}), 400
-
-        # Train Prophet model
-        model = Prophet(daily_seasonality=True, yearly_seasonality=True)
-        model.fit(df)
-
-        # Make future predictions
-        future_dates = model.make_future_dataframe(periods=interval)
-        forecast = model.predict(future_dates)
-
-        # Prepare and log predictions
-        forecast_data = forecast[['ds', 'yhat']].tail(interval)
-        predictions = [{"date": row['ds'].strftime('%Y-%m-%d'), "value": row['yhat']} for _, row in forecast_data.iterrows()]
-        logging.info(f"Predictions: {predictions}")
-
-        return jsonify({"predictedData": predictions})
-
-    except Exception as e:
-        logging.error(f"Error in /predict_metric endpoint: {e}")
-        return jsonify({"error": str(e)}), 500
+    request_data = request.json
+    user_id = request_data.get("user_id")
+    interval = request_data.get("interval", 30)
+    metric_type = request_data.get("metric_type", "Profit")  # "Profit" or "Customer Engagement"
+
+    transactions_ref = db.collection("system_users").document(user_id).collection("transactions")
+
+    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
+    df['date'] = pd.to_datetime(df['date'])
+    df['date'] = df['date'].dt.tz_localize(None)
+
+    # Set 'date' as index
+    df = df.sort_values("date").set_index("date")
+
+    # Resample daily to ensure regular intervals (fill missing dates)
+    df = df.resample("D").sum().reset_index()
+
+    df.columns = ["ds", "y"]  # ds: date, y: target
+
+    # Check if there's enough data to train the model
+    if df.shape[0] < 10:
+        return jsonify({"error": "Not enough data for prediction"})
+
+    # Initialize and fit the Prophet model
+    model = Prophet(daily_seasonality=True, yearly_seasonality=True)
+    model.fit(df)
+
+    # 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)
+    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})
+