Update forecast_model.py

forecast_model.py

@@ -0,0 +1,89 @@
+import pandas as pd
+from prophet import Prophet
+import json
+
+def train_and_forecast(data):
+    try:
+        # Convert input data to DataFrame
+        if not data:
+            return {"error": "Input data is empty"}
+        df = pd.DataFrame(data)
+        
+        # Validate required columns
+        required_cols = ['Date', 'Attendance']
+        if not all(col in df.columns for col in required_cols):
+            return {"error": "Input must contain 'Date' and 'Attendance' columns"}
+        
+        # Validate and process Date and Attendance
+        df['Date'] = pd.to_datetime(df['Date'], errors='coerce')
+        if df['Date'].isna().any():
+            return {"error": "Invalid 'Date' format. Use YYYY-MM-DD"}
+        df['Attendance'] = pd.to_numeric(df['Attendance'], errors='coerce')
+        if df['Attendance'].isna().any():
+            return {"error": "'Attendance' must be numeric"}
+        
+        # Check for sufficient data
+        if len(df) < 2:
+            return {"error": "At least 2 data points are required for forecasting"}
+        
+        # Preserve Risk and Alert for historical dates
+        df_risk_alert = None
+        if 'Risk' in df.columns and 'Alert' in df.columns:
+            df_risk_alert = df[['Date', 'Risk', 'Alert']].drop_duplicates(subset='Date')
+        
+        # Rename columns for Prophet
+        df_prophet = df.rename(columns={'Date': 'ds', 'Attendance': 'y'})
+        
+        # Initialize and train Prophet model
+        model = Prophet(
+            yearly_seasonality=True,
+            weekly_seasonality=True,
+            daily_seasonality=True
+        )
+        model.fit(df_prophet)
+
+        # Create future dates (30 days)
+        future = model.make_future_dataframe(periods=30, freq='D')
+        forecast = model.predict(future)
+
+        # Select relevant columns and rename back
+        forecast = forecast[['ds', 'yhat']].rename(columns={'ds': 'Date', 'yhat': 'Attendance'})
+        forecast['Date'] = forecast['Date'].astype(str)
+
+        # Estimate Risk and Alert for all dates
+        avg_attendance = df['Attendance'].mean()
+        def estimate_risk(attendance):
+            if attendance >= avg_attendance * 0.8:
+                return round(10.0 + (avg_attendance - attendance) * 2, 1)
+            elif attendance >= avg_attendance * 0.5:
+                return round(15.0 + (avg_attendance - attendance) * 3, 1)
+            else:
+                return round(20.0 + (avg_attendance - attendance) * 4, 1)
+
+        def estimate_alert(attendance):
+            if attendance >= avg_attendance * 0.8:
+                return "Low"
+            elif attendance >= avg_attendance * 0.5:
+                return "Medium"
+            else:
+                return "High"
+
+        forecast['Risk'] = forecast['Attendance'].apply(estimate_risk)
+        forecast['Alert'] = forecast['Attendance'].apply(estimate_alert)
+
+        # Merge historical Risk/Alert
+        if df_risk_alert is not None:
+            df_risk_alert['Date'] = df_risk_alert['Date'].astype(str)
+            forecast = forecast.merge(df_risk_alert, on='Date', how='left', suffixes=('', '_hist'))
+            forecast['Risk'] = forecast['Risk_hist'].combine_first(forecast['Risk'])
+            forecast['Alert'] = forecast['Alert_hist'].combine_first(forecast['Alert'])
+            forecast = forecast.drop(columns=['Risk_hist', 'Alert_hist'])
+
+        # Round numeric values
+        forecast['Attendance'] = forecast['Attendance'].round(1)
+
+        # Output relevant columns
+        forecast = forecast[['Date', 'Attendance', 'Risk', 'Alert']]
+        return forecast.to_dict('records')
+    except Exception as e:
+        return {"error": f"Forecasting failed: {str(e)}"}