Update app.py

app.py

@@ -2,8 +2,9 @@ import streamlit as st
 import streamlit.components.v1 as components
 import pandas as pd
 import matplotlib.pyplot as plt
+import plotly.figure_factory as ff
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 import json
 import requests
 import base64
@@ -68,7 +69,7 @@ if 'weather_data' not in st.session_state:
 def fetch_weather_data(project_location, date):
     if not WEATHER_API_KEY:
         logger.error("WEATHER_API_KEY not set")
-        return None, "Weather API key not set. Please provide a valid API key."
+        return None, {"error": "Weather API key not set. Please provide a valid API key."}
     try:
         params = {
             "q": project_location,
@@ -79,7 +80,7 @@ def fetch_weather_data(project_location, date):
         response.raise_for_status()
         data = response.json()
         
-        # Find the closest forecast to the specified date
+        # Find the closest forecast to the target date
         target_date = datetime.strptime(date, "%Y-%m-%d")
         closest_forecast = None
         min_time_diff = float('inf')
@@ -92,31 +93,30 @@ def fetch_weather_data(project_location, date):
                 closest_forecast = forecast
         
         if not closest_forecast:
-            return None, "No forecast available for the specified date."
+            return None, {"error": "No forecast available for the specified date."}
         
         # Map weather conditions to impact score
-        weather_main = closest_forecast['weather'][0]['main'].lower()
+        weather_main = forecast['weather'][0]['main'].lower()
+        impact_score = 50  # Default
         if 'clear' in weather_main:
             impact_score = 10
         elif 'clouds' in weather_main:
-            impact_score = 30 if closest_forecast['clouds']['all'] < 50 else 50
+            impact_score = 30 if forecast['clouds']['all'] < 50 else 50
         elif 'rain' in weather_main:
-            impact_score = 70 if closest_forecast['rain'].get('3h', 0) < 2.5 else 85
+            impact_score = 70 if forecast['rain'].get('3h', 0) < 2.5 else 85
         elif 'storm' in weather_main or 'thunderstorm' in weather_main:
             impact_score = 90
-        else:
-            impact_score = 50  # Default for other conditions (e.g., fog, snow)
         
         weather_condition = get_weather_condition(impact_score)
         return {
             "weather_impact_score": impact_score,
             "weather_condition": weather_condition,
-            "temperature": closest_forecast['main']['temp'],
-            "humidity": closest_forecast['main']['humidity']
+            "temperature": forecast['main']['temp'],
+            "humidity": forecast['main']['humidity']
         }, None
     except Exception as e:
         logger.error(f"Failed to fetch weather data: {str(e)}")
-        return None, f"Failed to fetch weather data for {project_location}: {str(e)}"
+        return None, {"error": f"Failed to fetch weather data for {project_location}: {str(e)}"}
 
 # Function to format high_risk_phases with flag and alert
 def format_high_risk_phases(high_risk_phases):
@@ -127,8 +127,57 @@ def format_high_risk_phases(high_risk_phases):
         formatted.append(f"{flag} {phase['phase']}: {phase['task']} (Risk: {phase['risk']:.1f}%) {alert}")
     return formatted
 
+# Function to generate Gantt chart
+def generate_gantt_chart(input_data, prediction):
+    try:
+        phase = input_data["phase"]
+        task = input_data["task"]
+        expected_duration = input_data["task_expected_duration"]
+        actual_duration = input_data["task_actual_duration"]
+        forecast_date = datetime.strptime(input_data["weather_forecast_date"], "%Y-%m-%d")
+        delay_risk = prediction["delay_probability"]
+
+        # Calculate start and end dates
+        start_date = forecast_date - timedelta(days=max(expected_duration, actual_duration))
+        expected_end = start_date + timedelta(days=expected_duration)
+        actual_end = start_date + timedelta(days=actual_duration) if actual_duration > 0 else expected_end
+
+        # Prepare Gantt chart data
+        df = [
+            dict(Task=f"{phase}: {task} (Expected)", Start=start_date.strftime("%Y-%m-%d"), Finish=expected_end.strftime("%Y-%m-%d"), Resource="Expected", Risk=delay_risk),
+            dict(Task=f"{phase}: {task} (Actual)", Start=start_date.strftime("%Y-%m-%d"), Finish=actual_end.strftime("%Y-%m-%d"), Resource="Actual", Risk=delay_risk)
+        ]
+
+        # Color based on delay risk
+        colors = {
+            "Expected": "rgb(0, 255, 0)" if delay_risk <= 50 else "rgb(255, 255, 0)" if delay_risk <= 75 else "rgb(255, 0, 0)",
+            "Actual": "rgb(0, 200, 0)" if delay_risk <= 50 else "rgb(200, 200, 0)" if delay_risk <= 75 else "rgb(200, 0, 0)"
+        }
+
+        # Create Gantt chart
+        fig = ff.create_gantt(
+            df,
+            colors=colors,
+            index_col="Resource",
+            title=f"Gantt Chart for {phase}: {task}",
+            show_colorbar=True,
+            bar_width=0.4,
+            showgrid_x=True,
+            showgrid_y=True
+        )
+        fig.update_layout(
+            xaxis_title="Timeline",
+            yaxis_title="Task",
+            height=300,
+            margin=dict(l=150)
+        )
+        return fig
+    except Exception as e:
+        logger.error(f"Failed to generate Gantt chart: {str(e)}")
+        return None
+
 # Function to generate PDF
-def generate_pdf(input_data, prediction, heatmap_fig):
+def generate_pdf(input_data, prediction, heatmap_fig, gantt_fig):
     buffer = BytesIO()
     doc = SimpleDocTemplate(buffer, pagesize=letter)
     styles = getSampleStyleSheet()
@@ -186,17 +235,33 @@ def generate_pdf(input_data, prediction, heatmap_fig):
     heatmap_fig.savefig(img_buffer, format='png', bbox_inches='tight')
     img_buffer.seek(0)
     story.append(Image(img_buffer, width=6*inch, height=2*inch))
-    plt.close(heatmap_fig)
+    story.append(Spacer(1, 12))
+
+    # Gantt Chart
+    if gantt_fig:
+        story.append(Paragraph("Gantt Chart", styles['Heading2']))
+        gantt_buffer = BytesIO()
+        try:
+            gantt_fig.write_image(gantt_buffer, format='PNG')
+            gantt_buffer.seek(0)
+            story.append(Image(gantt_buffer, width=6*inch, height=3*inch))
+        except Exception as e:
+            logger.error(f"Failed to include Gantt chart in PDF: {str(e)}")
+            story.append(Paragraph("Gantt Chart unavailable due to rendering issues.", styles['Normal']))
+        story.append(Spacer(1, 12))
 
     doc.build(story)
     buffer.seek(0)
     return buffer
 
-# Function to save data to Salesforce, including PDF
+# Function to save data to Salesforce, including PDF and Status__c
 def save_to_salesforce(input_data, prediction, pdf_buffer):
     if sf is None:
         return "Salesforce connection not established."
     try:
+        # Determine Status__c based on delay probability
+        status = "Flagged" if prediction["delay_probability"] > 75 else "Running"
+
         # Prepare data for Delay_Predictor__c object
         sf_data = {
             "Project_Name__c": input_data["project_name"],
@@ -214,7 +279,8 @@ def save_to_salesforce(input_data, prediction, pdf_buffer):
             "Project_Location__c": input_data["project_location"],
             "Delay_Probability__c": prediction["delay_probability"],
             "AI_Insights__c": prediction["ai_insights"],
-            "High_Risk_Phases__c": "; ".join(format_high_risk_phases(prediction["high_risk_phases"]))
+            "High_Risk_Phases__c": "; ".join(format_high_risk_phases(prediction["high_risk_phases"])),
+            "Status__c": status
         }
         logger.info(f"Attempting to save to Salesforce Delay_Predictor__c: {sf_data}")
         
@@ -324,8 +390,8 @@ if submit_button:
         if project_location and weather_forecast_date:
             weather_data, weather_error = fetch_weather_data(project_location, input_data["weather_forecast_date"])
             if weather_error:
-                st.error(weather_error)
-                logger.error(weather_error)
+                st.error(weather_error.get("error", "Unknown weather error"))
+                logger.error(weather_error.get("error", "Unknown weather error"))
                 input_data["weather_impact_score"] = 50  # Fallback value
                 input_data["weather_condition"] = "Unknown"
             else:
@@ -399,7 +465,13 @@ if submit_button:
                 ax.set_title("Delay Risk Heatmap")
                 plt.tight_layout()
                 
-                pdf_buffer = generate_pdf(input_data, prediction, fig)
+                # Generate Gantt chart
+                gantt_fig = generate_gantt_chart(input_data, prediction)
+                if gantt_fig:
+                    st.plotly_chart(gantt_fig, use_container_width=True)
+                    logger.info("Gantt chart rendered")
+                
+                pdf_buffer = generate_pdf(input_data, prediction, fig, gantt_fig)
                 plt.close(fig)
                 st.download_button(
                     label="Download Prediction Report (PDF)",