Update app.py

app.py

@@ -1,59 +1,67 @@
-import streamlit as st
-import pandas as pd
-from typing import Dict, List, Optional, Any
-from pydantic import BaseModel, Field
+import os
+import json
 import base64
 import io
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Any
+
+import numpy as np
+import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
-from abc import ABC, abstractmethod
+import streamlit as st
+
+from scipy.stats import ttest_ind, f_oneway
 from sklearn.model_selection import train_test_split
 from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import accuracy_score
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
 from statsmodels.tsa.seasonal import seasonal_decompose
 from statsmodels.tsa.stattools import adfuller
+
+from pydantic import BaseModel, Field
+from Bio import Entrez
+
 from langchain.prompts import PromptTemplate
 from groq import Groq
-import os
-import numpy as np
-from scipy.stats import ttest_ind, f_oneway
-import json
-from Bio import Entrez
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
 
-# Initialize Groq Client
+# ---------------------- Initialize External Clients ---------------------------
+# Initialize Groq Client with API Key from environment variables
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
 # ---------------------- Base Classes and Schemas ---------------------------
+
 class ResearchInput(BaseModel):
-    """Base schema for research tool inputs"""
+    """Base schema for research tool inputs."""
     data_key: str = Field(..., description="Session state key containing DataFrame")
     columns: Optional[List[str]] = Field(None, description="List of columns to analyze")
 
 class TemporalAnalysisInput(ResearchInput):
-    """Schema for temporal analysis"""
+    """Schema for temporal analysis."""
     time_col: str = Field(..., description="Name of timestamp column")
     value_col: str = Field(..., description="Name of value column to analyze")
 
 class HypothesisInput(ResearchInput):
-    """Schema for hypothesis testing"""
+    """Schema for hypothesis testing."""
     group_col: str = Field(..., description="Categorical column defining groups")
     value_col: str = Field(..., description="Numerical column to compare")
 
 class ModelTrainingInput(ResearchInput):
-    """Schema for model training"""
+    """Schema for model training."""
     target_col: str = Field(..., description="Name of target column")
 
 class DataAnalyzer(ABC):
-    """Abstract base class for data analysis modules"""
+    """Abstract base class for data analysis modules."""
     @abstractmethod
-    def invoke(self, data:pd.DataFrame, **kwargs) -> Dict[str, Any]:
+    def invoke(self, data: pd.DataFrame, **kwargs) -> Dict[str, Any]:
         pass
 
 # ---------------------- Concrete Analyzer Implementations ---------------------------
+
 class AdvancedEDA(DataAnalyzer):
-    """Comprehensive Exploratory Data Analysis"""
+    """Comprehensive Exploratory Data Analysis."""
     def invoke(self, data: pd.DataFrame, **kwargs) -> Dict[str, Any]:
         try:
             analysis = {
@@ -84,44 +92,44 @@ class AdvancedEDA(DataAnalyzer):
             return {"error": f"EDA Failed: {str(e)}"}
 
 class DistributionVisualizer(DataAnalyzer):
-    """Distribution visualizations"""
+    """Distribution visualizations."""
     def invoke(self, data: pd.DataFrame, columns: List[str], **kwargs) -> str:
-      try:
-          plt.figure(figsize=(12, 6))
-          for i, col in enumerate(columns, 1):
-              plt.subplot(1, len(columns), i)
-              sns.histplot(data[col], kde=True, stat="density")
-              plt.title(f'Distribution of {col}', fontsize=10)
-              plt.xticks(fontsize=8)
-              plt.yticks(fontsize=8)
-          plt.tight_layout()
-        
-          buf = io.BytesIO()
-          plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
-          plt.close()
-          return base64.b64encode(buf.getvalue()).decode()
-      except Exception as e:
-          return f"Visualization Error: {str(e)}"
+        try:
+            plt.figure(figsize=(12, 6))
+            for i, col in enumerate(columns, 1):
+                plt.subplot(1, len(columns), i)
+                sns.histplot(data[col], kde=True, stat="density")
+                plt.title(f'Distribution of {col}', fontsize=10)
+                plt.xticks(fontsize=8)
+                plt.yticks(fontsize=8)
+            plt.tight_layout()
+
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
+            plt.close()
+            return base64.b64encode(buf.getvalue()).decode()
+        except Exception as e:
+            return f"Visualization Error: {str(e)}"
 
 class TemporalAnalyzer(DataAnalyzer):
-    """Time series analysis"""
+    """Time series analysis."""
     def invoke(self, data: pd.DataFrame, time_col: str, value_col: str, **kwargs) -> Dict[str, Any]:
         try:
             ts_data = data.set_index(pd.to_datetime(data[time_col]))[value_col]
             decomposition = seasonal_decompose(ts_data, period=365)
-            
+
             plt.figure(figsize=(12, 8))
             decomposition.plot()
             plt.tight_layout()
-            
+
             buf = io.BytesIO()
             plt.savefig(buf, format='png')
             plt.close()
             plot_data = base64.b64encode(buf.getvalue()).decode()
-            
+
             return {
                 "trend_statistics": {
-                    "stationarity": adfuller(ts_data)[1],
+                    "stationarity_p_value": adfuller(ts_data)[1],
                     "seasonality_strength": max(decomposition.seasonal)
                 },
                 "visualization": plot_data
@@ -130,628 +138,866 @@ class TemporalAnalyzer(DataAnalyzer):
             return {"error": f"Temporal Analysis Failed: {str(e)}"}
 
 class HypothesisTester(DataAnalyzer):
-    """Statistical hypothesis testing"""
+    """Statistical hypothesis testing."""
     def invoke(self, data: pd.DataFrame, group_col: str, value_col: str, **kwargs) -> Dict[str, Any]:
-      try:
-        groups = data[group_col].unique()
-      
-        if len(groups) < 2:
-            return {"error": "Insufficient groups for comparison"}
-        
-        if len(groups) == 2:
+        try:
+            groups = data[group_col].unique()
+
+            if len(groups) < 2:
+                return {"error": "Insufficient groups for comparison"}
+
             group_data = [data[data[group_col] == g][value_col] for g in groups]
-            stat, p = ttest_ind(*group_data)
-            test_type = "Independent t-test"
-        else:
-          group_data = [data[data[group_col] == g][value_col] for g in groups]
-          stat, p = f_oneway(*group_data)
-          test_type = "ANOVA"
-        
-        return {
-            "test_type": test_type,
-            "test_statistic": stat,
-            "p_value": p,
-            "effect_size": {
-              "cohens_d": abs(group_data[0].mean() - group_data[1].mean())/np.sqrt(
-                    (group_data[0].var() + group_data[1].var())/2
-              ) if len(groups) == 2 else None
-            },
-            "interpretation": self.interpret_p_value(p)
-        }
-      except Exception as e:
-        return {"error": f"Hypothesis Testing Failed: {str(e)}"}
 
-    def interpret_p_value(self, p: float) -> str:
-      if p < 0.001: return "Very strong evidence against H0"
-      elif p < 0.01: return "Strong evidence against H0"
-      elif p < 0.05: return "Evidence against H0"
-      elif p < 0.1: return "Weak evidence against H0"
-      else: return "No significant evidence against H0"
+            if len(groups) == 2:
+                stat, p = ttest_ind(*group_data)
+                test_type = "Independent t-test"
+                effect_size = self.calculate_cohens_d(group_data[0], group_data[1])
+            else:
+                stat, p = f_oneway(*group_data)
+                test_type = "ANOVA"
+                effect_size = None
+
+            return {
+                "test_type": test_type,
+                "test_statistic": stat,
+                "p_value": p,
+                "effect_size": effect_size,
+                "interpretation": self.interpret_p_value(p)
+            }
+        except Exception as e:
+            return {"error": f"Hypothesis Testing Failed: {str(e)}"}
+
+    @staticmethod
+    def calculate_cohens_d(x: pd.Series, y: pd.Series) -> Optional[float]:
+        """Calculate Cohen's d for effect size."""
+        try:
+            mean_diff = abs(x.mean() - y.mean())
+            pooled_std = np.sqrt((x.var() + y.var()) / 2)
+            return mean_diff / pooled_std
+        except Exception:
+            return None
+
+    @staticmethod
+    def interpret_p_value(p: float) -> str:
+        """Interpret the p-value."""
+        if p < 0.001:
+            return "Very strong evidence against H0"
+        elif p < 0.01:
+            return "Strong evidence against H0"
+        elif p < 0.05:
+            return "Evidence against H0"
+        elif p < 0.1:
+            return "Weak evidence against H0"
+        else:
+            return "No significant evidence against H0"
 
 class LogisticRegressionTrainer(DataAnalyzer):
-    """Logistic Regression Model Trainer"""
+    """Logistic Regression Model Trainer."""
     def invoke(self, data: pd.DataFrame, target_col: str, columns: List[str], **kwargs) -> Dict[str, Any]:
-      try:
-        X = data[columns]
-        y = data[target_col]
-        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
-        model = LogisticRegression(max_iter=1000)
-        model.fit(X_train, y_train)
-        y_pred = model.predict(X_test)
-        accuracy = accuracy_score(y_test, y_pred)
-        return {
-          "model_type": "Logistic Regression",
-           "accuracy": accuracy,
-           "model_params": model.get_params()
-         }
-      except Exception as e:
-         return {"error": f"Logistic Regression Model Error: {str(e)}"}
+        try:
+            X = data[columns]
+            y = data[target_col]
+            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+            model = LogisticRegression(max_iter=1000)
+            model.fit(X_train, y_train)
+            y_pred = model.predict(X_test)
+            accuracy = accuracy_score(y_test, y_pred)
+            return {
+                "model_type": "Logistic Regression",
+                "accuracy": accuracy,
+                "model_params": model.get_params()
+            }
+        except Exception as e:
+            return {"error": f"Logistic Regression Model Error: {str(e)}"}
+
 # ---------------------- Business Logic Layer ---------------------------
 
 class ClinicalRule(BaseModel):
-    """Defines a clinical rule"""
+    """Defines a clinical rule."""
     name: str
     condition: str
     action: str
-    severity: str # low, medium or high
+    severity: str  # low, medium, or high
 
-class ClinicalRulesEngine():
+class ClinicalRulesEngine:
     """Executes rules against patient data."""
     def __init__(self):
-       self.rules: Dict[str, ClinicalRule] = {}
-    
+        self.rules: Dict[str, ClinicalRule] = {}
+
     def add_rule(self, rule: ClinicalRule):
         self.rules[rule.name] = rule
 
-    def execute_rules(self, data: pd.DataFrame):
+    def execute_rules(self, data: pd.DataFrame) -> Dict[str, Any]:
         results = {}
         for rule_name, rule in self.rules.items():
             try:
-              if eval(rule.condition, {}, {"df":data}):
-                results[rule_name] = {"rule_matched": True,
-                                       "action": rule.action,
-                                      "severity": rule.severity
-                                     }
-              else:
-                results[rule_name] = {"rule_matched": False, "action": None, "severity": None}
+                # Evaluate the condition using the dataframe 'df'
+                rule_matched = eval(rule.condition, {}, {"df": data})
+                results[rule_name] = {
+                    "rule_matched": rule_matched,
+                    "action": rule.action if rule_matched else None,
+                    "severity": rule.severity if rule_matched else None
+                }
             except Exception as e:
-                 results[rule_name] = {"rule_matched": False, "error": str(e), "severity": None}
+                results[rule_name] = {
+                    "rule_matched": False,
+                    "error": str(e),
+                    "severity": None
+                }
         return results
 
 class ClinicalKPI(BaseModel):
-  """Define a clinical KPI"""
-  name: str
-  calculation: str
-  threshold: Optional[float] = None
+    """Define a clinical KPI."""
+    name: str
+    calculation: str
+    threshold: Optional[float] = None
 
-class ClinicalKPIMonitoring():
-    """Calculates KPIs based on data"""
+class ClinicalKPIMonitoring:
+    """Calculates KPIs based on data."""
     def __init__(self):
-      self.kpis : Dict[str, ClinicalKPI] = {}
+        self.kpis: Dict[str, ClinicalKPI] = {}
 
-    def add_kpi(self, kpi:ClinicalKPI):
-      self.kpis[kpi.name] = kpi
+    def add_kpi(self, kpi: ClinicalKPI):
+        self.kpis[kpi.name] = kpi
 
-    def calculate_kpis(self, data: pd.DataFrame):
+    def calculate_kpis(self, data: pd.DataFrame) -> Dict[str, Any]:
         results = {}
         for kpi_name, kpi in self.kpis.items():
             try:
-                results[kpi_name] = eval(kpi.calculation, {}, {"df": data})
+                kpi_value = eval(kpi.calculation, {}, {"df": data})
+                results[kpi_name] = {
+                    "value": kpi_value,
+                    "threshold": kpi.threshold,
+                    "status": self.evaluate_threshold(kpi_value, kpi.threshold)
+                }
             except Exception as e:
                 results[kpi_name] = {"error": str(e)}
         return results
 
+    @staticmethod
+    def evaluate_threshold(value: Any, threshold: Optional[float]) -> Optional[str]:
+        if threshold is None:
+            return None
+        try:
+            return "Above Threshold" if value > threshold else "Below Threshold"
+        except TypeError:
+            return "Threshold Evaluation Not Applicable"
+
 class DiagnosisSupport(ABC):
-  """Abstract class for implementing clinical diagnoses."""
-  @abstractmethod
-  def diagnose(self, data: pd.DataFrame, target_col: str, columns: List[str], diagnosis_key : str = "diagnosis" , **kwargs) -> pd.DataFrame:
-         pass
+    """Abstract class for implementing clinical diagnoses."""
+    @abstractmethod
+    def diagnose(
+        self, 
+        data: pd.DataFrame, 
+        target_col: str, 
+        columns: List[str], 
+        diagnosis_key: str = "diagnosis", 
+        **kwargs
+    ) -> pd.DataFrame:
+        pass
 
 class SimpleDiagnosis(DiagnosisSupport):
-    """Provides a simple diagnosis example, based on the Logistic regression model"""
+    """Provides a simple diagnosis example, based on the Logistic regression model."""
     def __init__(self):
-      self.model : LogisticRegressionTrainer = LogisticRegressionTrainer()
-    
-    def diagnose(self, data: pd.DataFrame, target_col: str, columns: List[str], diagnosis_key : str = "diagnosis", **kwargs) -> pd.DataFrame:
+        self.model: LogisticRegressionTrainer = LogisticRegressionTrainer()
+
+    def diagnose(
+        self, 
+        data: pd.DataFrame, 
+        target_col: str, 
+        columns: List[str], 
+        diagnosis_key: str = "diagnosis", 
+        **kwargs
+    ) -> pd.DataFrame:
         try:
-             result = self.model.invoke(data, target_col=target_col, columns = columns)
-             if "accuracy" in result:
-               return pd.DataFrame({diagnosis_key: [f"Accuracy {result['accuracy']}"],
-                                   "model": result["model_type"]})
-             else:
-                return pd.DataFrame({diagnosis_key: [f"Diagnosis failed: {result}"]})
-
+            result = self.model.invoke(data, target_col=target_col, columns=columns)
+            if "accuracy" in result:
+                return pd.DataFrame({
+                    diagnosis_key: [f"Model Accuracy: {result['accuracy']:.2%}"],
+                    "model": [result["model_type"]]
+                })
+            else:
+                return pd.DataFrame({
+                    diagnosis_key: [f"Diagnosis failed: {result.get('error', 'Unknown error')}"]
+                })
         except Exception as e:
-                return pd.DataFrame({diagnosis_key:[f"Error during diagnosis {e}"]})
-
+            return pd.DataFrame({
+                diagnosis_key: [f"Error during diagnosis: {e}"]
+            })
 
 class TreatmentRecommendation(ABC):
-  """Abstract class for treatment recommendations"""
-  @abstractmethod
-  def recommend(self, data: pd.DataFrame,  condition_col: str, treatment_col:str, recommendation_key: str = "recommendation", **kwargs) -> pd.DataFrame:
-     pass
+    """Abstract class for treatment recommendations."""
+    @abstractmethod
+    def recommend(
+        self, 
+        data: pd.DataFrame,  
+        condition_col: str, 
+        treatment_col: str, 
+        recommendation_key: str = "recommendation", 
+        **kwargs
+    ) -> pd.DataFrame:
+        pass
 
 class BasicTreatmentRecommendation(TreatmentRecommendation):
-  """A placeholder class for basic treatment recommendations"""
-  def recommend(self, data: pd.DataFrame, condition_col: str, treatment_col:str, recommendation_key: str = "recommendation", **kwargs) -> pd.DataFrame:
-       if condition_col not in data.columns or treatment_col not in data.columns:
-          return pd.DataFrame({recommendation_key: ["Condition or Treatment columns not found!"]})
-       treatment = data[data[condition_col] == "High"][treatment_col].to_list()
-       if len(treatment)>0:
-         return pd.DataFrame({recommendation_key: [f"Treatment recommended for High risk patients: {treatment}"]})
-       else:
-           return pd.DataFrame({recommendation_key: [f"No treatment recommendation found!"]})
-
-
-class MedicalKnowledgeBase():
-  """Abstract class for Medical Knowledge"""
-  @abstractmethod
-  def search_medical_info(self, query: str, pub_email:str="") -> str:
-    pass
+    """A placeholder class for basic treatment recommendations."""
+    def recommend(
+        self, 
+        data: pd.DataFrame, 
+        condition_col: str, 
+        treatment_col: str, 
+        recommendation_key: str = "recommendation", 
+        **kwargs
+    ) -> pd.DataFrame:
+        if condition_col not in data.columns or treatment_col not in data.columns:
+            return pd.DataFrame({
+                recommendation_key: ["Condition or Treatment columns not found!"]
+            })
+        
+        treatment = data[data[condition_col] == "High"][treatment_col].to_list()
+        if treatment:
+            return pd.DataFrame({
+                recommendation_key: [f"Treatment recommended for High risk patients: {treatment}"]
+            })
+        else:
+            return pd.DataFrame({
+                recommendation_key: ["No treatment recommendation found!"]
+            })
+
+class MedicalKnowledgeBase(ABC):
+    """Abstract class for Medical Knowledge."""
+    @abstractmethod
+    def search_medical_info(self, query: str, pub_email: str = "") -> str:
+        pass
 
 class SimpleMedicalKnowledge(MedicalKnowledgeBase):
-    """Simple Medical Knowledge Class with TF-IDF and PubMed"""
+    """Simple Medical Knowledge Class with TF-IDF and PubMed."""
     def __init__(self):
-       self.knowledge_base = {
+        self.knowledge_base = {
             "diabetes": "The recommended treatment for diabetes includes lifestyle changes, medication, and monitoring.",
             "heart disease": "Risk factors for heart disease include high blood pressure, high cholesterol, and smoking.",
-             "fever": "For a fever, you can consider over-the-counter medications like acetaminophen or ibuprofen. Rest and hydration are also important.",
+            "fever": "For a fever, you can consider over-the-counter medications like acetaminophen or ibuprofen. Rest and hydration are also important.",
             "headache": "For a headache, try rest, hydration, and over-the-counter pain relievers. Consult a doctor if it is severe or persistent.",
             "cold": "For a cold, get rest, drink plenty of fluids, and use over-the-counter remedies like decongestants."
-       }
-       self.vectorizer = TfidfVectorizer()
-       self.tfidf_matrix = self.vectorizer.fit_transform(self.knowledge_base.values())
-       
+        }
+        self.vectorizer = TfidfVectorizer()
+        self.tfidf_matrix = self.vectorizer.fit_transform(self.knowledge_base.values())
+
     def search_pubmed(self, query: str, email: str) -> str:
-      try:
-          Entrez.email = email
-          handle = Entrez.esearch(db="pubmed", term=query, retmax=1)
-          record = Entrez.read(handle)
-          handle.close()
-          if record["IdList"]:
-            handle = Entrez.efetch(db="pubmed", id=record["IdList"][0], rettype="abstract", retmode="text")
-            abstract = handle.read()
+        """Search PubMed for abstracts related to the query."""
+        try:
+            Entrez.email = email
+            handle = Entrez.esearch(db="pubmed", term=query, retmax=1)
+            record = Entrez.read(handle)
             handle.close()
-            return abstract
-          else:
-            return "No abstracts found for this query on PubMed"
-      except Exception as e:
-            return f"Error searching pubmed {e}"
-
+            if record["IdList"]:
+                handle = Entrez.efetch(db="pubmed", id=record["IdList"][0], rettype="abstract", retmode="text")
+                abstract = handle.read()
+                handle.close()
+                return abstract
+            else:
+                return "No abstracts found for this query on PubMed."
+        except Exception as e:
+            return f"Error searching PubMed: {e}"
 
     def search_medical_info(self, query: str, pub_email: str = "") -> str:
+        """Search the medical knowledge base and PubMed for relevant information."""
         try:
-          query_vector = self.vectorizer.transform([query])
-          similarities = cosine_similarity(query_vector, self.tfidf_matrix)
-          best_match_index = np.argmax(similarities)
-          best_match_keyword = list(self.knowledge_base.keys())[best_match_index]
-          best_match_info =  list(self.knowledge_base.values())[best_match_index]
-            
-          pubmed_result = self.search_pubmed(query, pub_email)
-          if "No abstracts found for this query on PubMed" not in pubmed_result:
-             return f"Based on the query provided, I found this: {best_match_info} \n\nFrom Pubmed I also found the following abstract: \n {pubmed_result}"
-          else:
-            return f"Based on the query provided, I found this: {best_match_info} \n\n{pubmed_result}"
+            query_vector = self.vectorizer.transform([query])
+            similarities = cosine_similarity(query_vector, self.tfidf_matrix)
+            best_match_index = np.argmax(similarities)
+            best_match_keyword = list(self.knowledge_base.keys())[best_match_index]
+            best_match_info = list(self.knowledge_base.values())[best_match_index]
+
+            pubmed_result = self.search_pubmed(query, pub_email)
+            if "No abstracts found" not in pubmed_result:
+                return (
+                    f"**Based on your query:** {best_match_info}\n\n"
+                    f"**PubMed Abstract:**\n{pubmed_result}"
+                )
+            else:
+                return (
+                    f"**Based on your query:** {best_match_info}\n\n"
+                    f"{pubmed_result}"
+                )
         except Exception as e:
-            return f"Medical Knowledge Search Failed {e}"
-
+            return f"Medical Knowledge Search Failed: {e}"
 
 class ForecastingEngine(ABC):
+    """Abstract class for forecasting."""
     @abstractmethod
     def predict(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
-         pass
+        pass
 
 class SimpleForecasting(ForecastingEngine):
-  def predict(self, data: pd.DataFrame, period: int = 7, **kwargs) -> pd.DataFrame:
-      #Placeholder for actual forecasting
-      return pd.DataFrame({"forecast":[f"Forecast for the next {period} days"]})
+    """Simple forecasting engine."""
+    def predict(self, data: pd.DataFrame, period: int = 7, **kwargs) -> pd.DataFrame:
+        # Placeholder for actual forecasting logic
+        return pd.DataFrame({"forecast": [f"Forecast for the next {period} days"]})
 
 # ---------------------- Insights and Reporting Layer ---------------------------
-class AutomatedInsights():
+
+class AutomatedInsights:
+    """Generates automated insights based on selected analyses."""
     def __init__(self):
-      self.analyses : Dict[str, DataAnalyzer] = {
-        "EDA": AdvancedEDA(),
-        "temporal": TemporalAnalyzer(),
-        "distribution": DistributionVisualizer(),
-        "hypothesis": HypothesisTester(),
-        "model": LogisticRegressionTrainer()
-      }
-
-    def generate_insights(self, data: pd.DataFrame, analysis_names: List[str], **kwargs):
-       results = {}
-       for name in analysis_names:
-           if name in self.analyses:
-             analyzer = self.analyses[name]
-             try:
-                results[name] = analyzer.invoke(data=data, **kwargs)
-             except Exception as e:
-                results[name] = {"error": str(e)}
-           else:
-               results[name] = {"error": "Analysis not found"}
-       return results
-
-class Dashboard():
+        self.analyses: Dict[str, DataAnalyzer] = {
+            "EDA": AdvancedEDA(),
+            "temporal": TemporalAnalyzer(),
+            "distribution": DistributionVisualizer(),
+            "hypothesis": HypothesisTester(),
+            "model": LogisticRegressionTrainer()
+        }
+
+    def generate_insights(self, data: pd.DataFrame, analysis_names: List[str], **kwargs) -> Dict[str, Any]:
+        results = {}
+        for name in analysis_names:
+            analyzer = self.analyses.get(name)
+            if analyzer:
+                try:
+                    results[name] = analyzer.invoke(data=data, **kwargs)
+                except Exception as e:
+                    results[name] = {"error": str(e)}
+            else:
+                results[name] = {"error": "Analysis not found"}
+        return results
+
+class Dashboard:
+    """Handles the creation and display of the dashboard."""
     def __init__(self):
-        self.layout: Dict[str,str] = {}
-    
+        self.layout: Dict[str, str] = {}
+
     def add_visualisation(self, vis_name: str, vis_type: str):
         self.layout[vis_name] = vis_type
-  
-    def display_dashboard(self, data_dict: Dict[str,pd.DataFrame]):
-      st.header("Dashboard")
-      for vis_name, vis_type in self.layout.items():
-          st.subheader(vis_name)
-          if vis_type == "table":
-              if vis_name in data_dict:
-                st.table(data_dict[vis_name])
-              else:
-                st.write("Data Not Found")
-          elif vis_type == "plot":
-            if vis_name in data_dict:
-              df = data_dict[vis_name]
-              if len(df.columns) > 1:
-                fig = plt.figure()
-                sns.lineplot(data=df)
-                st.pyplot(fig)
-              else:
-                 st.write("Please have more than 1 column")
-            else:
-                 st.write("Data not found")
-class AutomatedReports():
-  def __init__(self):
-     self.report_definition: Dict[str,str] = {}
-
-  def create_report_definition(self, report_name: str, definition: str):
-       self.report_definition[report_name] = definition
-  
-  def generate_report(self, report_name: str, data:Dict[str, pd.DataFrame]):
-       if report_name not in self.report_definition:
-          return {"error":"Report name not found"}
-       st.header(f"Report : {report_name}")
-       st.write(f"Report Definition: {self.report_definition[report_name]}")
-       for df_name, df in data.items():
-         st.subheader(f"Data: {df_name}")
-         st.table(df)
+
+    def display_dashboard(self, data_dict: Dict[str, pd.DataFrame]):
+        st.header("Dashboard")
+        for vis_name, vis_type in self.layout.items():
+            st.subheader(vis_name)
+            if vis_type == "table":
+                df = data_dict.get(vis_name)
+                if df is not None:
+                    st.table(df)
+                else:
+                    st.write("Data Not Found")
+            elif vis_type == "plot":
+                df = data_dict.get(vis_name)
+                if df is not None:
+                    if len(df.columns) > 1:
+                        fig = plt.figure()
+                        sns.lineplot(data=df)
+                        st.pyplot(fig)
+                    else:
+                        st.write("Please select a DataFrame with more than 1 column for plotting.")
+                else:
+                    st.write("Data not found")
+
+class AutomatedReports:
+    """Manages automated report definitions and generation."""
+    def __init__(self):
+        self.report_definitions: Dict[str, str] = {}
+
+    def create_report_definition(self, report_name: str, definition: str):
+        self.report_definitions[report_name] = definition
+
+    def generate_report(self, report_name: str, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
+        if report_name not in self.report_definitions:
+            return {"error": "Report name not found"}
+        report_content = {
+            "Report Name": report_name,
+            "Report Definition": self.report_definitions[report_name],
+            "Data": {df_name: df.to_dict() for df_name, df in data.items()}
+        }
+        return report_content
 
 # ---------------------- Data Acquisition Layer ---------------------------
+
 class DataSource(ABC):
     """Base class for data sources."""
     @abstractmethod
     def connect(self) -> None:
         """Connect to the data source."""
         pass
-    
+
     @abstractmethod
     def fetch_data(self, query: str, **kwargs) -> pd.DataFrame:
-         """Fetch the data based on a specific query."""
-         pass
-
+        """Fetch the data based on a specific query."""
+        pass
 
 class CSVDataSource(DataSource):
     """Data source for CSV files."""
-    def __init__(self, file_path: str):
+    def __init__(self, file_path: io.BytesIO):
         self.file_path = file_path
         self.data: Optional[pd.DataFrame] = None
-    
+
     def connect(self):
         self.data = pd.read_csv(self.file_path)
-    
+
     def fetch_data(self, query: str = None, **kwargs) -> pd.DataFrame:
-      if self.data is None:
-        raise Exception("No connection is made, call connect()")
-      return self.data
-    
+        if self.data is None:
+            raise Exception("No connection is made, call connect()")
+        return self.data
+
 class DatabaseSource(DataSource):
-   def __init__(self, connection_string: str, database_type: str):
+    """Data source for SQL Databases."""
+    def __init__(self, connection_string: str, database_type: str):
         self.connection_string = connection_string
-        self.database_type = database_type
+        self.database_type = database_type.lower()
         self.connection = None
-    
-   def connect(self):
-     if self.database_type.lower() == "sql":
-        #Placeholder for the actual database connection
-        self.connection = "Connected to SQL Database"
-     else:
-        raise Exception(f"Database type '{self.database_type}' is not supported")
-      
-   def fetch_data(self, query: str, **kwargs) -> pd.DataFrame:
+
+    def connect(self):
+        if self.database_type == "sql":
+            # Placeholder for actual SQL connection logic
+            self.connection = "Connected to SQL Database"
+        else:
+            raise Exception(f"Database type '{self.database_type}' is not supported.")
+
+    def fetch_data(self, query: str, **kwargs) -> pd.DataFrame:
         if self.connection is None:
             raise Exception("No connection is made, call connect()")
-        #Placeholder for the data fetching
-        return pd.DataFrame({"result":[f"Fetched data based on query: {query}"]})
-     
+        # Placeholder for data fetching logic
+        return pd.DataFrame({"result": [f"Fetched data based on query: {query}"]})
 
 class DataIngestion:
+    """Handles data ingestion from various sources."""
     def __init__(self):
-      self.sources : Dict[str, DataSource] = {}
-    
+        self.sources: Dict[str, DataSource] = {}
+
     def add_source(self, source_name: str, source: DataSource):
-         self.sources[source_name] = source
-    
+        self.sources[source_name] = source
+
     def ingest_data(self, source_name: str, query: str = None, **kwargs) -> pd.DataFrame:
-       if source_name not in self.sources:
-         raise Exception(f"Source '{source_name}' not found")
-       source = self.sources[source_name]
-       source.connect()
-       return source.fetch_data(query, **kwargs)
-       
+        if source_name not in self.sources:
+            raise Exception(f"Source '{source_name}' not found.")
+        source = self.sources[source_name]
+        source.connect()
+        return source.fetch_data(query, **kwargs)
+
 class DataModel(BaseModel):
-  name : str
-  kpis : List[str] = Field(default_factory=list)
-  dimensions : List[str] = Field(default_factory=list)
-  custom_calculations : Optional[Dict[str, str]] = None
-  relations: Optional[Dict[str,str]] = None #Example {table1: table2}
-  
-  def to_json(self):
-    return json.dumps(self.dict())
-  
-  @staticmethod
-  def from_json(json_str):
-    return DataModel(**json.loads(json_str))
-  
-class DataModelling():
+    """Defines a data model."""
+    name: str
+    kpis: List[str] = Field(default_factory=list)
+    dimensions: List[str] = Field(default_factory=list)
+    custom_calculations: Optional[Dict[str, str]] = None
+    relations: Optional[Dict[str, str]] = None  # Example: {"table1": "table2"}
+
+    def to_json(self) -> str:
+        return json.dumps(self.dict())
+
+    @staticmethod
+    def from_json(json_str: str) -> 'DataModel':
+        return DataModel(**json.loads(json_str))
+
+class DataModelling:
+    """Manages data models."""
     def __init__(self):
-      self.models : Dict[str, DataModel] = {}
-    
-    def add_model(self, model:DataModel):
-       self.models[model.name] = model
-    
+        self.models: Dict[str, DataModel] = {}
+
+    def add_model(self, model: DataModel):
+        self.models[model.name] = model
+
     def get_model(self, model_name: str) -> DataModel:
-      if model_name not in self.models:
-          raise Exception(f"Model '{model_name}' not found")
-      return self.models[model_name]
+        if model_name not in self.models:
+            raise Exception(f"Model '{model_name}' not found.")
+        return self.models[model_name]
+
 # ---------------------- Main Streamlit Application ---------------------------
+
 def main():
+    """Main function to run the Streamlit app."""
     st.set_page_config(page_title="AI Clinical Intelligence Hub", layout="wide")
     st.title("🏥 AI-Powered Clinical Intelligence Hub")
 
-    # Session State
+    # Initialize Session State
+    initialize_session_state()
+
+    # Sidebar for Data Management
+    with st.sidebar:
+        data_management_section()
+
+    # Main Content
+    if st.session_state.data:
+        col1, col2 = st.columns([1, 3])
+
+        with col1:
+            dataset_metadata_section()
+
+        with col2:
+            main_tabs_section()
+
+def initialize_session_state():
+    """Initialize necessary components in Streamlit's session state."""
     if 'data' not in st.session_state:
-        st.session_state.data = {}  # store pd.DataFrame under a name
+        st.session_state.data = {}  # Store pd.DataFrame under a name
     if 'data_ingestion' not in st.session_state:
-      st.session_state.data_ingestion = DataIngestion()
+        st.session_state.data_ingestion = DataIngestion()
     if 'data_modelling' not in st.session_state:
-      st.session_state.data_modelling = DataModelling()
+        st.session_state.data_modelling = DataModelling()
     if 'clinical_rules' not in st.session_state:
-      st.session_state.clinical_rules = ClinicalRulesEngine()
+        st.session_state.clinical_rules = ClinicalRulesEngine()
     if 'kpi_monitoring' not in st.session_state:
-       st.session_state.kpi_monitoring = ClinicalKPIMonitoring()
+        st.session_state.kpi_monitoring = ClinicalKPIMonitoring()
     if 'forecasting_engine' not in st.session_state:
-       st.session_state.forecasting_engine = SimpleForecasting()
+        st.session_state.forecasting_engine = SimpleForecasting()
     if 'automated_insights' not in st.session_state:
-      st.session_state.automated_insights = AutomatedInsights()
+        st.session_state.automated_insights = AutomatedInsights()
     if 'dashboard' not in st.session_state:
-       st.session_state.dashboard = Dashboard()
+        st.session_state.dashboard = Dashboard()
     if 'automated_reports' not in st.session_state:
-      st.session_state.automated_reports = AutomatedReports()
+        st.session_state.automated_reports = AutomatedReports()
     if 'diagnosis_support' not in st.session_state:
-       st.session_state.diagnosis_support = SimpleDiagnosis()
+        st.session_state.diagnosis_support = SimpleDiagnosis()
     if 'treatment_recommendation' not in st.session_state:
-       st.session_state.treatment_recommendation = BasicTreatmentRecommendation()
+        st.session_state.treatment_recommendation = BasicTreatmentRecommendation()
     if 'knowledge_base' not in st.session_state:
-      st.session_state.knowledge_base = SimpleMedicalKnowledge()
+        st.session_state.knowledge_base = SimpleMedicalKnowledge()
     if 'pub_email' not in st.session_state:
         st.session_state.pub_email = st.secrets.get("PUB_EMAIL", "")  # Load PUB_EMAIL from secrets
-  
 
-    # Sidebar for Data Management
-    with st.sidebar:
-        st.header("⚙️ Data Management")
-        data_source_selection = st.selectbox("Select Data Source Type",["CSV","SQL Database"])
-        if data_source_selection == "CSV":
-           uploaded_file = st.file_uploader("Upload research dataset (CSV)", type=["csv"])
-           if uploaded_file:
-              source_name = st.text_input("Data Source Name")
-              if source_name:
+def data_management_section():
+    """Handles the data management section in the sidebar."""
+    st.header("⚙️ Data Management")
+    data_source_selection = st.selectbox("Select Data Source Type", ["CSV", "SQL Database"])
+
+    if data_source_selection == "CSV":
+        handle_csv_upload()
+    elif data_source_selection == "SQL Database":
+        handle_sql_database()
+
+    if st.button("Ingest Data"):
+        ingest_data_action()
+
+def handle_csv_upload():
+    """Handles CSV file uploads."""
+    uploaded_file = st.file_uploader("Upload research dataset (CSV)", type=["csv"])
+    if uploaded_file:
+        source_name = st.text_input("Data Source Name")
+        if source_name:
+            try:
+                csv_source = CSVDataSource(file_path=uploaded_file)
+                st.session_state.data_ingestion.add_source(source_name, csv_source)
+                st.success(f"Uploaded {uploaded_file.name} as '{source_name}'.")
+            except Exception as e:
+                st.error(f"Error loading dataset: {e}")
+
+def handle_sql_database():
+    """Handles SQL database connections."""
+    conn_str = st.text_input("Enter connection string for SQL DB")
+    if conn_str:
+        source_name = st.text_input("Data Source Name")
+        if source_name:
+            try:
+                sql_source = DatabaseSource(connection_string=conn_str, database_type="sql")
+                st.session_state.data_ingestion.add_source(source_name, sql_source)
+                st.success(f"Added SQL DB Source '{source_name}'.")
+            except Exception as e:
+                st.error(f"Error loading database source: {e}")
+
+def ingest_data_action():
+    """Performs data ingestion from the selected source."""
+    if st.session_state.data_ingestion.sources:
+        source_name_to_fetch = st.selectbox("Select Data Source to Ingest", list(st.session_state.data_ingestion.sources.keys()))
+        query = st.text_area("Optional Query to Fetch data")
+        if source_name_to_fetch:
+            with st.spinner("Ingesting data..."):
                 try:
-                    csv_source = CSVDataSource(file_path=uploaded_file)
-                    st.session_state.data_ingestion.add_source(source_name,csv_source)
-                    st.success(f"Uploaded {uploaded_file.name}")
+                    data = st.session_state.data_ingestion.ingest_data(source_name_to_fetch, query)
+                    st.session_state.data[source_name_to_fetch] = data
+                    st.success(f"Ingested data from '{source_name_to_fetch}'.")
                 except Exception as e:
-                  st.error(f"Error loading dataset: {e}")
-        elif data_source_selection == "SQL Database":
-          conn_str = st.text_input("Enter connection string for SQL DB")
-          if conn_str:
-            source_name = st.text_input("Data Source Name")
-            if source_name:
-              try:
-                 sql_source = DatabaseSource(connection_string=conn_str, database_type="sql")
-                 st.session_state.data_ingestion.add_source(source_name, sql_source)
-                 st.success(f"Added SQL DB Source {source_name}")
-              except Exception as e:
-                 st.error(f"Error loading database source {e}")
-
-
-        if st.button("Ingest Data"):
-           if st.session_state.data_ingestion.sources:
-               source_name_to_fetch = st.selectbox("Select Data Source to Ingest", list(st.session_state.data_ingestion.sources.keys()))
-               query = st.text_area("Optional Query to Fetch data")
-               if source_name_to_fetch:
-                    with st.spinner("Ingesting data..."):
-                      try:
-                          data = st.session_state.data_ingestion.ingest_data(source_name_to_fetch, query)
-                          st.session_state.data[source_name_to_fetch] = data
-                          st.success(f"Ingested data from {source_name_to_fetch}")
-                      except Exception as e:
-                        st.error(f"Ingestion failed: {e}")
-           else:
-             st.error("No data source added, please add data source")
+                    st.error(f"Ingestion failed: {e}")
+    else:
+        st.error("No data source added. Please add a data source.")
+
+def dataset_metadata_section():
+    """Displays metadata for the selected dataset."""
+    st.subheader("Dataset Metadata")
+    data_source_keys = list(st.session_state.data.keys())
+    selected_data_key = st.selectbox("Select Dataset", data_source_keys)
+
+    if selected_data_key:
+        data = st.session_state.data[selected_data_key]
+        metadata = {
+            "Variables": list(data.columns),
+            "Time Range": {
+                col: {
+                    "min": data[col].min(),
+                    "max": data[col].max()
+                } for col in data.select_dtypes(include='datetime').columns
+            },
+            "Size": f"{data.memory_usage().sum() / 1e6:.2f} MB"
+        }
+        st.json(metadata)
+
+def main_tabs_section():
+    """Creates and manages the main tabs in the application."""
+    analysis_tab, clinical_logic_tab, insights_tab, reports_tab, knowledge_tab = st.tabs([
+        "Data Analysis",
+        "Clinical Logic",
+        "Insights",
+        "Reports",
+        "Medical Knowledge"
+    ])
+
+    with analysis_tab:
+        data_analysis_section()
+
+    with clinical_logic_tab:
+        clinical_logic_section()
+
+    with insights_tab:
+        insights_section()
+
+    with reports_tab:
+        reports_section()
+
+    with knowledge_tab:
+        medical_knowledge_section()
+
+def data_analysis_section():
+    """Handles the Data Analysis tab."""
+    selected_data_key = st.sidebar.selectbox("Select Dataset for Analysis", list(st.session_state.data.keys()))
+    if not selected_data_key:
+        st.warning("Please select a dataset to perform analysis.")
+        return
+
+    data = st.session_state.data[selected_data_key]
+    analysis_type = st.selectbox("Select Analysis Mode", [
+        "Exploratory Data Analysis",
+        "Temporal Pattern Analysis",
+        "Comparative Statistics",
+        "Distribution Analysis",
+        "Train Logistic Regression Model"
+    ])
+
+    if analysis_type == "Exploratory Data Analysis":
+        perform_eda(data)
+    elif analysis_type == "Temporal Pattern Analysis":
+        perform_temporal_analysis(data)
+    elif analysis_type == "Comparative Statistics":
+        perform_comparative_statistics(data)
+    elif analysis_type == "Distribution Analysis":
+        perform_distribution_analysis(data)
+    elif analysis_type == "Train Logistic Regression Model":
+        perform_logistic_regression_training(data)
+
+def perform_eda(data: pd.DataFrame):
+    """Performs Exploratory Data Analysis."""
+    analyzer = AdvancedEDA()
+    eda_result = analyzer.invoke(data=data)
+    st.subheader("Data Quality Report")
+    st.json(eda_result)
+
+def perform_temporal_analysis(data: pd.DataFrame):
+    """Performs Temporal Pattern Analysis."""
+    time_cols = data.select_dtypes(include='datetime').columns
+    num_cols = data.select_dtypes(include=np.number).columns
+
+    time_col = st.selectbox("Select Temporal Variable", time_cols)
+    value_col = st.selectbox("Select Analysis Variable", num_cols)
+
+    if time_col and value_col:
+        analyzer = TemporalAnalyzer()
+        result = analyzer.invoke(data=data, time_col=time_col, value_col=value_col)
+        if "visualization" in result:
+            st.image(f"data:image/png;base64,{result['visualization']}")
+        st.json(result)
+
+def perform_comparative_statistics(data: pd.DataFrame):
+    """Performs Comparative Statistics."""
+    categorical_cols = data.select_dtypes(include='category').columns
+    numeric_cols = data.select_dtypes(include=np.number).columns
+
+    group_col = st.selectbox("Select Grouping Variable", categorical_cols)
+    value_col = st.selectbox("Select Metric Variable", numeric_cols)
+
+    if group_col and value_col:
+        analyzer = HypothesisTester()
+        result = analyzer.invoke(data=data, group_col=group_col, value_col=value_col)
+        st.subheader("Statistical Test Results")
+        st.json(result)
+
+def perform_distribution_analysis(data: pd.DataFrame):
+    """Performs Distribution Analysis."""
+    numeric_cols = data.select_dtypes(include=np.number).columns.tolist()
+    selected_cols = st.multiselect("Select Variables for Distribution Analysis", numeric_cols)
+
+    if selected_cols:
+        analyzer = DistributionVisualizer()
+        img_data = analyzer.invoke(data=data, columns=selected_cols)
+        if "Visualization Error" not in img_data:
+            st.image(f"data:image/png;base64,{img_data}")
+        else:
+            st.error(img_data)
+
+def perform_logistic_regression_training(data: pd.DataFrame):
+    """Trains a Logistic Regression model."""
+    numeric_cols = data.select_dtypes(include=np.number).columns.tolist()
+    target_col = st.selectbox("Select Target Variable", data.columns.tolist())
+    selected_cols = st.multiselect("Select Feature Variables", numeric_cols)
+
+    if selected_cols and target_col:
+        analyzer = LogisticRegressionTrainer()
+        result = analyzer.invoke(data=data, target_col=target_col, columns=selected_cols)
+        st.subheader("Logistic Regression Model Results")
+        st.json(result)
+
+def clinical_logic_section():
+    """Handles the Clinical Logic tab."""
+    st.header("Clinical Logic")
+    
+    # Clinical Rules Management
+    st.subheader("Clinical Rules")
+    rule_name = st.text_input("Enter Rule Name")
+    condition = st.text_area("Enter Rule Condition (use 'df' for DataFrame)", 
+                             help="Example: df['blood_pressure'] > 140")
+    action = st.text_area("Enter Action to be Taken on Rule Match")
+    severity = st.selectbox("Enter Severity for the Rule", ["low", "medium", "high"])
+
+    if st.button("Add Clinical Rule"):
+        try:
+            rule = ClinicalRule(
+                name=rule_name, 
+                condition=condition, 
+                action=action, 
+                severity=severity
+            )
+            st.session_state.clinical_rules.add_rule(rule)
+            st.success("Added Clinical Rule successfully.")
+        except Exception as e:
+            st.error(f"Error in rule definition: {e}")
 
-    if st.session_state.data:
-        col1, col2 = st.columns([1, 3])
+    # Clinical KPI Management
+    st.subheader("Clinical KPI Definition")
+    kpi_name = st.text_input("Enter KPI Name")
+    kpi_calculation = st.text_area("Enter KPI Calculation (use 'df' for DataFrame)", 
+                                    help="Example: df['patient_count'].sum()")
+    threshold = st.text_input("Enter Threshold for KPI", help="Optional")
 
-        with col1:
-            st.subheader("Dataset Metadata")
-            
-            data_source_keys = list(st.session_state.data.keys())
-            selected_data_key = st.selectbox("Select Dataset", data_source_keys)
-            
-            if selected_data_key:
-                data = st.session_state.data[selected_data_key]
-                st.json({
-                    "Variables": list(data.columns),
-                    "Time Range": {
-                        col: {
-                            "min": data[col].min(),
-                            "max": data[col].max()
-                        } for col in data.select_dtypes(include='datetime').columns
-                    },
-                    "Size": f"{data.memory_usage().sum() / 1e6:.2f} MB"
-                })
-        with col2:
-            analysis_tab, clinical_logic_tab, insights_tab, reports_tab, knowledge_tab = st.tabs([
-              "Data Analysis",
-              "Clinical Logic",
-               "Insights",
-               "Reports",
-              "Medical Knowledge"
-              ])
-
-            with analysis_tab:
-                if selected_data_key:
-                    analysis_type = st.selectbox("Select Analysis Mode", [
-                     "Exploratory Data Analysis",
-                     "Temporal Pattern Analysis",
-                      "Comparative Statistics",
-                      "Distribution Analysis",
-                       "Train Logistic Regression Model"
-                    ])
-                    data = st.session_state.data[selected_data_key]
-                    if analysis_type == "Exploratory Data Analysis":
-                        analyzer = AdvancedEDA()
-                        eda_result = analyzer.invoke(data=data)
-                        st.subheader("Data Quality Report")
-                        st.json(eda_result)
-
-                    elif analysis_type == "Temporal Pattern Analysis":
-                        time_col = st.selectbox("Temporal Variable",
-                            data.select_dtypes(include='datetime').columns)
-                        value_col = st.selectbox("Analysis Variable",
-                            data.select_dtypes(include=np.number).columns)
-
-                        if time_col and value_col:
-                          analyzer = TemporalAnalyzer()
-                          result = analyzer.invoke(data=data, time_col=time_col, value_col=value_col)
-                          if "visualization" in result:
-                              st.image(f"data:image/png;base64,{result['visualization']}")
-                          st.json(result)
-
-                    elif analysis_type == "Comparative Statistics":
-                        group_col = st.selectbox("Grouping Variable",
-                             data.select_dtypes(include='category').columns)
-                        value_col = st.selectbox("Metric Variable",
-                             data.select_dtypes(include=np.number).columns)
-
-                        if group_col and value_col:
-                          analyzer = HypothesisTester()
-                          result = analyzer.invoke(data=data, group_col=group_col, value_col=value_col)
-                          st.subheader("Statistical Test Results")
-                          st.json(result)
-
-                    elif analysis_type == "Distribution Analysis":
-                        num_cols = data.select_dtypes(include=np.number).columns.tolist()
-                        selected_cols = st.multiselect("Select Variables", num_cols)
-                        if selected_cols:
-                           analyzer = DistributionVisualizer()
-                           img_data = analyzer.invoke(data=data, columns=selected_cols)
-                           st.image(f"data:image/png;base64,{img_data}")
-                       
-                    elif analysis_type == "Train Logistic Regression Model":
-                       num_cols = data.select_dtypes(include=np.number).columns.tolist()
-                       target_col = st.selectbox("Select Target Variable",
-                                                   data.columns.tolist())
-                       selected_cols = st.multiselect("Select Feature Variables", num_cols)
-                       if selected_cols and target_col:
-                         analyzer = LogisticRegressionTrainer()
-                         result = analyzer.invoke(data=data, target_col=target_col, columns=selected_cols)
-                         st.subheader("Logistic Regression Model Results")
-                         st.json(result)
-            with clinical_logic_tab:
-              st.header("Clinical Logic")
-              st.subheader("Clinical Rules")
-              rule_name = st.text_input("Enter Rule Name")
-              condition = st.text_area("Enter Rule Condition (use 'df' for data frame), Example df['blood_pressure'] > 140")
-              action = st.text_area("Enter Action to be Taken on Rule Match")
-              severity = st.selectbox("Enter Severity for the Rule", ["low","medium","high"])
-              if st.button("Add Clinical Rule"):
-                try:
-                 rule = ClinicalRule(name=rule_name, condition=condition, action=action, severity=severity)
-                 st.session_state.clinical_rules.add_rule(rule)
-                 st.success("Added Clinical Rule")
-                except Exception as e:
-                  st.error(f"Error in rule definition: {e}")
-                
-              st.subheader("Clinical KPI Definition")
-              kpi_name = st.text_input("Enter KPI name")
-              kpi_calculation = st.text_area("Enter KPI calculation (use 'df' for data frame), Example df['patient_count'].sum()")
-              threshold = st.text_input("Enter Threshold for KPI")
-              if st.button("Add Clinical KPI"):
-                  try:
-                    threshold_value = float(threshold) if threshold else None
-                    kpi = ClinicalKPI(name=kpi_name, calculation=kpi_calculation, threshold=threshold_value)
-                    st.session_state.kpi_monitoring.add_kpi(kpi)
-                    st.success(f"Added KPI {kpi_name}")
-                  except Exception as e:
-                      st.error(f"Error creating KPI: {e}")
-
-              if selected_data_key:
-                   data = st.session_state.data[selected_data_key]
-                   if st.button("Execute Clinical Rules"):
-                       with st.spinner("Executing Clinical Rules.."):
-                         result = st.session_state.clinical_rules.execute_rules(data)
-                         st.json(result)
-                   if st.button("Calculate Clinical KPIs"):
-                       with st.spinner("Calculating Clinical KPIs..."):
-                          result = st.session_state.kpi_monitoring.calculate_kpis(data)
-                          st.json(result)
-            with insights_tab:
-                if selected_data_key:
-                    data = st.session_state.data[selected_data_key]
-                    available_analysis = ["EDA", "temporal", "distribution", "hypothesis", "model"]
-                    selected_analysis = st.multiselect("Select Analysis", available_analysis)
-                    if st.button("Generate Automated Insights"):
-                        with st.spinner("Generating Insights"):
-                            results = st.session_state.automated_insights.generate_insights(data, analysis_names=selected_analysis)
-                            st.json(results)
-                    st.subheader("Diagnosis Support")
-                    target_col = st.selectbox("Select Target Variable for Diagnosis", data.columns.tolist())
-                    num_cols = data.select_dtypes(include=np.number).columns.tolist()
-                    selected_cols_diagnosis = st.multiselect("Select Feature Variables for Diagnosis", num_cols)
-                    if st.button("Generate Diagnosis"):
-                        if target_col
-                                                        if target_col and selected_cols_diagnosis:
-                                    with st.spinner("Generating Diagnosis"):
-                                        result = st.session_state.diagnosis_support.diagnose(data, target_col=target_col, columns=selected_cols_diagnosis, diagnosis_key="diagnosis_result")
-                                        st.json(result)
-
-                    st.subheader("Treatment Recommendation")
-                    condition_col = st.selectbox("Select Condition Column for Treatment Recommendation", data.columns.tolist())
-                    treatment_col = st.selectbox("Select Treatment Column for Treatment Recommendation", data.columns.tolist())
-                    if st.button("Generate Treatment Recommendation"):
-                        if condition_col and treatment_col:
-                            with st.spinner("Generating Treatment Recommendation"):
-                                result = st.session_state.treatment_recommendation.recommend(data, condition_col = condition_col, treatment_col = treatment_col, recommendation_key="treatment_recommendation")
-                                st.json(result)
-
-            with reports_tab:
-                 st.header("Reports")
-                 report_name = st.text_input("Report Name")
-                 report_def = st.text_area("Report definition")
-                 if st.button("Create Report Definition"):
-                   st.session_state.automated_reports.create_report_definition(report_name, report_def)
-                   st.success("Report definition created")
-                 if selected_data_key:
-                   data = st.session_state.data
-                   if st.button("Generate Report"):
-                     with st.spinner("Generating Report..."):
-                       report = st.session_state.automated_reports.generate_report(report_name, data)
-            with knowledge_tab:
-              st.header("Medical Knowledge")
-              query = st.text_input("Enter your medical question here:")
-              if st.button("Search"):
-                  with st.spinner("Searching..."):
-                      result = st.session_state.knowledge_base.search_medical_info(query, pub_email=st.session_state.pub_email)
-                      st.write(result)
+    if st.button("Add Clinical KPI"):
+        try:
+            threshold_value = float(threshold) if threshold else None
+            kpi = ClinicalKPI(
+                name=kpi_name, 
+                calculation=kpi_calculation, 
+                threshold=threshold_value
+            )
+            st.session_state.kpi_monitoring.add_kpi(kpi)
+            st.success(f"Added KPI '{kpi_name}' successfully.")
+        except ValueError:
+            st.error("Threshold must be a numeric value.")
+        except Exception as e:
+            st.error(f"Error creating KPI: {e}")
+
+    # Execute Clinical Rules and Calculate KPIs
+    selected_data_key = st.selectbox("Select Dataset for Clinical Logic", list(st.session_state.data.keys()))
+    if selected_data_key:
+        data = st.session_state.data[selected_data_key]
+        if st.button("Execute Clinical Rules"):
+            with st.spinner("Executing Clinical Rules..."):
+                result = st.session_state.clinical_rules.execute_rules(data)
+                st.json(result)
+        if st.button("Calculate Clinical KPIs"):
+            with st.spinner("Calculating Clinical KPIs..."):
+                result = st.session_state.kpi_monitoring.calculate_kpis(data)
+                st.json(result)
+
+def insights_section():
+    """Handles the Insights tab."""
+    st.header("Automated Insights")
+
+    selected_data_key = st.selectbox("Select Dataset for Insights", list(st.session_state.data.keys()))
+    if not selected_data_key:
+        st.warning("Please select a dataset to generate insights.")
+        return
+
+    data = st.session_state.data[selected_data_key]
+    available_analyses = ["EDA", "temporal", "distribution", "hypothesis", "model"]
+    selected_analyses = st.multiselect("Select Analyses for Insights", available_analyses)
+
+    if st.button("Generate Automated Insights"):
+        with st.spinner("Generating Insights..."):
+            results = st.session_state.automated_insights.generate_insights(
+                data, analysis_names=selected_analyses
+            )
+            st.json(results)
+
+    # Diagnosis Support
+    st.subheader("Diagnosis Support")
+    target_col = st.selectbox("Select Target Variable for Diagnosis", data.columns.tolist())
+    numeric_cols = data.select_dtypes(include=np.number).columns.tolist()
+    selected_feature_cols = st.multiselect("Select Feature Variables for Diagnosis", numeric_cols)
+
+    if st.button("Generate Diagnosis"):
+        if target_col and selected_feature_cols:
+            with st.spinner("Generating Diagnosis..."):
+                result = st.session_state.diagnosis_support.diagnose(
+                    data, target_col=target_col, columns=selected_feature_cols, diagnosis_key="diagnosis_result"
+                )
+                st.json(result)
+        else:
+            st.error("Please select both target and feature variables for diagnosis.")
+
+    # Treatment Recommendation
+    st.subheader("Treatment Recommendation")
+    condition_col = st.selectbox("Select Condition Column for Treatment Recommendation", data.columns.tolist())
+    treatment_col = st.selectbox("Select Treatment Column for Treatment Recommendation", data.columns.tolist())
+
+    if st.button("Generate Treatment Recommendation"):
+        if condition_col and treatment_col:
+            with st.spinner("Generating Treatment Recommendation..."):
+                result = st.session_state.treatment_recommendation.recommend(
+                    data, condition_col=condition_col, treatment_col=treatment_col, recommendation_key="treatment_recommendation"
+                )
+                st.json(result)
+        else:
+            st.error("Please select both condition and treatment columns.")
+
+def reports_section():
+    """Handles the Reports tab."""
+    st.header("Automated Reports")
+    
+    # Create Report Definition
+    st.subheader("Create Report Definition")
+    report_name = st.text_input("Report Name")
+    report_def = st.text_area("Report Definition")
+
+    if st.button("Create Report Definition"):
+        if report_name and report_def:
+            st.session_state.automated_reports.create_report_definition(report_name, report_def)
+            st.success("Report definition created successfully.")
+        else:
+            st.error("Please provide both report name and definition.")
+
+    # Generate Report
+    st.subheader("Generate Report")
+    report_name_to_generate = st.selectbox("Select Report to Generate", list(st.session_state.automated_reports.report_definitions.keys()))
+    
+    if st.button("Generate Report"):
+        if report_name_to_generate:
+            with st.spinner("Generating Report..."):
+                report = st.session_state.automated_reports.generate_report(report_name_to_generate, st.session_state.data)
+                if "error" not in report:
+                    st.header(f"Report: {report_name_to_generate}")
+                    st.write(f"**Definition:** {report['Report Definition']}")
+                    for df_name, df_content in report["Data"].items():
+                        st.subheader(f"Data: {df_name}")
+                        st.write(pd.DataFrame(df_content))
+                else:
+                    st.error(report["error"])
+        else:
+            st.error("Please select a report to generate.")
+
+def medical_knowledge_section():
+    """Handles the Medical Knowledge tab."""
+    st.header("Medical Knowledge")
+    query = st.text_input("Enter your medical question here:")
+
+    if st.button("Search"):
+        if query:
+            with st.spinner("Searching..."):
+                result = st.session_state.knowledge_base.search_medical_info(query, pub_email=st.session_state.pub_email)
+                st.markdown(result)
+        else:
+            st.error("Please enter a medical question to search.")
 
 if __name__ == "__main__":
-    main()
+    main()