app.py

import streamlit as st
import pandas as pd
from typing import Dict, List, Optional, Any
from pydantic import BaseModel, Field
import base64
import io
import matplotlib.pyplot as plt
import seaborn as sns
from abc import ABC, abstractmethod
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import adfuller
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

# Initialize Groq Client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# ---------------------- Base Classes and Schemas ---------------------------
class ResearchInput(BaseModel):
    """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"""
    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"""
    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"""
    target_col: str = Field(..., description="Name of target column")

class DataAnalyzer(ABC):
    """Abstract base class for data analysis modules"""
    @abstractmethod
    def invoke(self, **kwargs) -> Dict[str, Any]:
        pass

# ---------------------- Concrete Analyzer Implementations ---------------------------
class AdvancedEDA(DataAnalyzer):
    """Comprehensive Exploratory Data Analysis"""
    def invoke(self, data: pd.DataFrame, **kwargs) -> Dict[str, Any]:
        try:
            analysis = {
                "dimensionality": {
                    "rows": len(data),
                    "columns": list(data.columns),
                    "memory_usage": f"{data.memory_usage().sum() / 1e6:.2f} MB"
                },
                "statistical_profile": data.describe(percentiles=[.25, .5, .75]).to_dict(),
                "temporal_analysis": {
                    "date_ranges": {
                        col: {
                            "min": data[col].min(),
                            "max": data[col].max()
                        } for col in data.select_dtypes(include='datetime').columns
                    }
                },
                "data_quality": {
                    "missing_values": data.isnull().sum().to_dict(),
                    "duplicates": data.duplicated().sum(),
                    "cardinality": {
                        col: data[col].nunique() for col in data.columns
                    }
                }
            }
            return analysis
        except Exception as e:
            return {"error": f"EDA Failed: {str(e)}"}

class DistributionVisualizer(DataAnalyzer):
    """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)}"

class TemporalAnalyzer(DataAnalyzer):
    """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],
                    "seasonality_strength": max(decomposition.seasonal)
                },
                "visualization": plot_data
            }
        except Exception as e:
            return {"error": f"Temporal Analysis Failed: {str(e)}"}

class HypothesisTester(DataAnalyzer):
    """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:
            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"

class LogisticRegressionTrainer(DataAnalyzer):
    """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)}"}
# ---------------------- Business Logic Layer ---------------------------

class ClinicalRule(BaseModel):
    """Defines a clinical rule"""
    name: str
    condition: str
    action: str
    severity: str # low, medium or high

class ClinicalRulesEngine():
    """Executes rules against patient data."""
    def __init__(self):
       self.rules: Dict[str, ClinicalRule] = {}
    
    def add_rule(self, rule: ClinicalRule):
        self.rules[rule.name] = rule

    def execute_rules(self, data: pd.DataFrame):
        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}
            except Exception as e:
                 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

class ClinicalKPIMonitoring():
    """Calculates KPIs based on data"""
    def __init__(self):
      self.kpis : Dict[str, ClinicalKPI] = {}

    def add_kpi(self, kpi:ClinicalKPI):
      self.kpis[kpi.name] = kpi

    def calculate_kpis(self, data: pd.DataFrame):
        results = {}
        for kpi_name, kpi in self.kpis.items():
            try:
                results[kpi_name] = eval(kpi.calculation, {}, {"df": data})
            except Exception as e:
                results[kpi_name] = {"error": str(e)}
        return results

class DiagnosisSupport(ABC):
  """Abstract class for implementing clinical diagnoses."""
  @abstractmethod
  def diagnose(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
         pass

class SimpleDiagnosis(DiagnosisSupport):
    """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], **kwargs) -> pd.DataFrame:
        try:
             result = self.model.invoke(data, target_col=target_col, columns = columns)
             if "accuracy" in result:
               return pd.DataFrame({"diagnosis": [f"Accuracy {result['accuracy']}"],
                                   "model": result["model_type"]})
             else:
                return pd.DataFrame({"diagnosis": [f"Diagnosis failed: {result}"]})

        except Exception as e:
                return pd.DataFrame({"diagnosis":[f"Error during diagnosis {e}"]})


class TreatmentRecommendation(ABC):
  """Abstract class for treatment recommendations"""
  @abstractmethod
  def recommend(self, data: pd.DataFrame, **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, **kwargs) -> pd.DataFrame:
       if condition_col not in data.columns or treatment_col not in data.columns:
          return pd.DataFrame({"recommendation": ["Condition or Treatment columns not found!"]})
       treatment = data[data[condition_col] == "High"][treatment_col].to_list()
       if len(treatment)>0:
         return pd.DataFrame({"recommendation": [f"Treatment recommended for High risk patients: {treatment}"]})
       else:
           return pd.DataFrame({"recommendation": [f"No treatment recommendation found!"]})


class MedicalKnowledgeBase():
  """Abstract class for Medical Knowledge"""
  @abstractmethod
  def search_medical_info(self, query: str) -> str:
    pass

class SimpleMedicalKnowledge(MedicalKnowledgeBase):
    """Simple Medical Knowledge Class"""
    def search_medical_info(self, query: str) -> str:
          if "diabetes treatment" in query.lower():
            return "The recommended treatment for diabetes includes lifestyle changes, medication, and monitoring"
          elif "heart disease risk factors" in query.lower():
               return "Risk factors for heart disease include high blood pressure, high cholesterol, and smoking"
          else:
            return "No specific information is available"


class ForecastingEngine(ABC):
    @abstractmethod
    def predict(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
         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"]})

# ---------------------- Insights and Reporting Layer ---------------------------
class AutomatedInsights():
    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]
             results[name] = analyzer.invoke(data=data, **kwargs)
           else:
               results[name] = {"error": "Analysis not found"}
       return results

class Dashboard():
    def __init__(self):
        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)

# ---------------------- 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


class CSVDataSource(DataSource):
    """Data source for CSV files."""
    def __init__(self, file_path: str):
        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
    
class DatabaseSource(DataSource):
   def __init__(self, connection_string: str, database_type: str):
        self.connection_string = connection_string
        self.database_type = database_type
        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:
        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}"]})
     

class DataIngestion:
    def __init__(self):
      self.sources : Dict[str, DataSource] = {}
    
    def add_source(self, source_name: str, source: DataSource):
         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)
       
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():
    def __init__(self):
      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]
# ---------------------- Main Streamlit Application ---------------------------
def main():
    st.set_page_config(page_title="AI Clinical Intelligence Hub", layout="wide")
    st.title("🏥 AI-Powered Clinical Intelligence Hub")

    # Session State
    if 'data' not in st.session_state:
        st.session_state.data = {}  # store pd.DataFrame under a name
    if 'data_ingestion' not in st.session_state:
      st.session_state.data_ingestion = DataIngestion()
    if 'data_modelling' not in st.session_state:
      st.session_state.data_modelling = DataModelling()
    if 'clinical_rules' not in st.session_state:
      st.session_state.clinical_rules = ClinicalRulesEngine()
    if 'kpi_monitoring' not in st.session_state:
       st.session_state.kpi_monitoring = ClinicalKPIMonitoring()
    if 'forecasting_engine' not in st.session_state:
       st.session_state.forecasting_engine = SimpleForecasting()
    if 'automated_insights' not in st.session_state:
      st.session_state.automated_insights = AutomatedInsights()
    if 'dashboard' not in st.session_state:
       st.session_state.dashboard = Dashboard()
    if 'automated_reports' not in st.session_state:
      st.session_state.automated_reports = AutomatedReports()
    if 'diagnosis_support' not in st.session_state:
       st.session_state.diagnosis_support = SimpleDiagnosis()
    if 'treatment_recommendation' not in st.session_state:
       st.session_state.treatment_recommendation = BasicTreatmentRecommendation()
    if 'knowledge_base' not in st.session_state:
      st.session_state.knowledge_base = SimpleMedicalKnowledge()
  

    # 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:
                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}")
                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")

    if st.session_state.data:
        col1, col2 = st.columns([1, 3])

        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 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)
                               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)
                              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)
                      st.write(result)

if __name__ == "__main__":
    main()