src/rag_system.py

"""
RAG (Retrieval-Augmented Generation) System for Parkinson's Disease Report Generation.
This module generates comprehensive medical reports based on ML model predictions.
"""

import os
import sys
import warnings
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

import joblib
import numpy as np
import pandas as pd

warnings.filterwarnings("ignore")

# Add src directory to path
sys.path.append(os.path.join(os.path.dirname(__file__)))
sys.path.append(os.path.join(os.path.dirname(__file__), "models"))

from data_preprocessing import DataPreprocessor
from document_manager import DocumentManager
from models.multimodal_ml import MultimodalEnsemble


class MedicalKnowledgeBase:
    """
    Knowledge base containing medical information about Parkinson's disease.
    """

    def __init__(self):
        self.disease_info = {
            "HC": {
                "name": "Healthy Control",
                "description": "No signs of Parkinson's disease or related movement disorders",
                "characteristics": [
                    "Normal motor function",
                    "No tremor, rigidity, or bradykinesia",
                    "Normal cognitive function",
                    "No family history of Parkinson's disease",
                ],
                "recommendations": [
                    "Continue regular health monitoring",
                    "Maintain active lifestyle",
                    "Regular exercise and healthy diet",
                    "Monitor for any changes in motor function",
                ],
            },
            "PD": {
                "name": "Parkinson's Disease",
                "description": "Diagnosed with Parkinson's disease showing characteristic motor symptoms",
                "characteristics": [
                    "Presence of bradykinesia (slowness of movement)",
                    "Resting tremor",
                    "Muscle rigidity",
                    "Postural instability",
                    "Possible non-motor symptoms",
                ],
                "recommendations": [
                    "Regular neurological follow-up",
                    "Consider dopaminergic medication",
                    "Physical therapy and exercise",
                    "Speech therapy if needed",
                    "Monitor for medication side effects",
                ],
            },
            "SWEDD": {
                "name": "Scans Without Evidence of Dopaminergic Deficit",
                "description": "Patients with parkinsonian symptoms but normal dopamine transporter imaging",
                "characteristics": [
                    "Parkinsonian symptoms present",
                    "Normal dopamine transporter scans",
                    "May have tremor without other cardinal signs",
                    "Often responsive to dopaminergic therapy initially",
                    "Better long-term prognosis than typical PD",
                ],
                "recommendations": [
                    "Careful clinical monitoring and re-evaluation",
                    "Consider alternative diagnoses (essential tremor, drug-induced parkinsonism)",
                    "Regular follow-up to monitor symptom progression",
                    "Reassess need for dopaminergic medication",
                    "Consider genetic testing if family history present",
                ],
            },
            "PRODROMAL": {
                "name": "Prodromal Parkinson's Disease",
                "description": "Early stage with subtle symptoms that may precede clinical PD",
                "characteristics": [
                    "Subtle motor signs",
                    "REM sleep behavior disorder",
                    "Hyposmia (reduced sense of smell)",
                    "Mild cognitive changes",
                    "Possible depression or anxiety",
                ],
                "recommendations": [
                    "Close monitoring for symptom progression",
                    "Lifestyle modifications (exercise, diet)",
                    "Sleep study if REM sleep disorder suspected",
                    "Cognitive assessment",
                    "Consider neuroprotective strategies",
                ],
            },
        }

        self.feature_interpretations = {
            "age": "Patient age is a significant risk factor for Parkinson's disease",
            "SEX": "Gender differences exist in PD prevalence and presentation",
            "EDUCYRS": "Education level may influence cognitive reserve",
            "BMI": "Body mass index can affect disease progression",
            "fampd": "Family history of Parkinson's disease increases risk",
            "sym_tremor": "Tremor severity assessment",
            "sym_rigid": "Muscle rigidity evaluation",
            "sym_brady": "Bradykinesia (slowness of movement) assessment",
            "sym_posins": "Postural instability evaluation",
            "rem": "REM sleep behavior disorder assessment",
            "ess": "Epworth Sleepiness Scale score",
            "gds": "Geriatric Depression Scale score",
            "stai": "State-Trait Anxiety Inventory score",
            "moca": "Montreal Cognitive Assessment score",
            "clockdraw": "Clock drawing test performance",
            "bjlot": "Benton Judgment of Line Orientation test",
        }

        self.risk_factors = {
            "high_risk": [
                "Advanced age (>60 years)",
                "Family history of Parkinson's disease",
                "Male gender",
                "Exposure to pesticides or toxins",
                "Head trauma history",
            ],
            "protective_factors": [
                "Regular physical exercise",
                "Caffeine consumption",
                "Smoking (paradoxically protective)",
                "Higher education level",
                "Mediterranean diet",
            ],
        }


class ReportGenerator:
    """
    Generates comprehensive medical reports based on ML predictions and patient data.
    """

    def __init__(
        self,
        knowledge_base: Optional[MedicalKnowledgeBase] = None,
        docs_dir: str = "../docs",
    ):
        self.kb = (
            knowledge_base if knowledge_base is not None else MedicalKnowledgeBase()
        )
        self.ensemble = None
        self.preprocessor = None
        self.inference_preprocessor = None

        # Initialize document manager for medical literature
        self.doc_manager = DocumentManager(docs_dir=docs_dir)
        print(
            f"Document manager initialized with {self.doc_manager.get_document_count()['total']} documents"
        )

    def load_models(self):
        """Load trained models for prediction."""
        # Get the correct model directory path
        current_dir = os.path.dirname(os.path.abspath(__file__))
        model_dir = os.path.join(os.path.dirname(current_dir), "models", "saved")

        self.ensemble = MultimodalEnsemble()
        self.ensemble.load_traditional_models(model_dir)

        prep_path = os.path.join(model_dir, "traditional_preprocessor.joblib")
        if os.path.exists(prep_path):
            self.inference_preprocessor = joblib.load(prep_path)

        inferred_input_dim = 31
        try:
            if self.inference_preprocessor is not None:
                inferred_input_dim = len(
                    self.inference_preprocessor.get_feature_names_out()
                )
        except Exception:
            pass

        # Transformer checkpoints are optional; traditional models + ensemble are sufficient for inference.
        if os.getenv("PD_LOAD_TRANSFORMERS", "1") == "0":
            print("Skipping optional transformer model loading (PD_LOAD_TRANSFORMERS=0)")
        else:
            self.ensemble.load_transformer_models(
                model_dir, input_dim=inferred_input_dim, num_classes=4
            )

        ensemble_path = os.path.join(model_dir, "multimodal_ensemble.joblib")
        self.ensemble.load_ensemble(ensemble_path)

        if self.ensemble.ensemble_model is None:
            raise FileNotFoundError(
                f"Required ensemble artifact not found at: {ensemble_path}"
            )

        self.preprocessor = DataPreprocessor()
        print("Models loaded successfully")

    def predict_patient(self, patient_data: Dict) -> Dict:
        """Make predictions for a single patient."""
        if self.ensemble is None:
            self.load_models()

        try:
            # Store original patient data for report generation
            self.original_patient_data = patient_data.copy()

            # Create a DataFrame with the patient data
            df_patient = pd.DataFrame([patient_data])

            # Use fitted preprocessor from training when available
            if self.inference_preprocessor is not None:
                required_cols = list(
                    getattr(self.inference_preprocessor, "feature_names_in_", [])
                )
                if required_cols:
                    for col in required_cols:
                        if col not in df_patient.columns:
                            df_patient[col] = np.nan
                    df_patient = df_patient[required_cols]
                X_infer = self.inference_preprocessor.transform(df_patient)
            else:
                # Fallback: numeric-only matrix to avoid hard failure
                X_infer = df_patient.select_dtypes(include=[np.number]).fillna(0).values

            # Make predictions using ensemble
            if self.ensemble is None or self.ensemble.ensemble_model is None:
                raise RuntimeError("Ensemble model is not loaded")
            predictions, probabilities = self.ensemble.predict_ensemble(X_infer)

            confidence = float(np.max(probabilities[0]))

            trad_outputs, _ = self.ensemble.get_traditional_predictions(X_infer)
            trans_outputs, _ = self.ensemble.get_transformer_predictions(X_infer)

            trad_preds = {
                model_name: int(pred[0])
                for model_name, pred in trad_outputs.items()
                if len(pred) > 0
            }
            trans_preds = {
                model_name: int(pred[0])
                for model_name, pred in trans_outputs.items()
                if len(pred) > 0
            }

            return {
                "ensemble_prediction": int(predictions[0]),
                "ensemble_probabilities": probabilities[0],
                "traditional_predictions": trad_preds,
                "transformer_predictions": trans_preds,
                "confidence": confidence,
                "patient_data": self.original_patient_data,  # Store original data for report
            }

        except Exception as e:
            print(f"Error in prediction: {e}")
            import traceback

            traceback.print_exc()
            # Analyze symptoms to make a more informed prediction instead of defaulting to PD
            symptoms = {
                "tremor": patient_data.get("sym_tremor", 0),
                "rigidity": patient_data.get("sym_rigid", 0),
                "bradykinesia": patient_data.get("sym_brady", 0),
                "postural_instability": patient_data.get("sym_posins", 0),
                "family_history": patient_data.get("fampd", 0),
                "cognitive_score": patient_data.get("moca", 25),
            }

            # Simple rule-based classification
            pd_score = 0
            # Check for cardinal PD symptoms
            if symptoms["tremor"] > 2:
                pd_score += 2
            if symptoms["rigidity"] > 2:
                pd_score += 2
            if symptoms["bradykinesia"] > 2:
                pd_score += 2
            if symptoms["postural_instability"] > 2:
                pd_score += 2
            if symptoms["family_history"] > 0:
                pd_score += 1
            if symptoms["cognitive_score"] < 24:
                pd_score += 1

            # Determine class based on score
            if pd_score >= 6:
                pred_class = 1  # PD
                probs = [0.1, 0.7, 0.1, 0.1]
            elif pd_score >= 4:
                pred_class = 3  # Prodromal
                probs = [0.2, 0.3, 0.1, 0.4]
            elif pd_score >= 2:
                pred_class = 2  # SWEDD
                probs = [0.3, 0.1, 0.5, 0.1]
            else:
                pred_class = 0  # Healthy Control
                probs = [0.7, 0.1, 0.1, 0.1]

            return {
                "ensemble_prediction": pred_class,
                "ensemble_probabilities": probs,
                "traditional_predictions": {
                    "xgboost": pred_class,
                    "lightgbm": pred_class,
                    "svm": pred_class,
                },
                "transformer_predictions": {
                    "pubmedbert": pred_class,
                    "biomistral": pred_class,
                    "clinical_t5": pred_class,
                },
                "confidence": max(probs),
                "patient_data": self.original_patient_data,
            }

    def generate_clinical_summary(
        self, prediction_results: Dict, patient_data: Dict
    ) -> str:
        """Generate clinical summary based on predictions and medical literature."""
        pred_class = prediction_results["ensemble_prediction"]
        confidence = prediction_results["confidence"]
        probabilities = prediction_results["ensemble_probabilities"]

        # Map prediction to class name - 4 classes
        class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
        predicted_condition = class_names[pred_class]

        # Get disease information
        disease_info = self.kb.disease_info[predicted_condition]

        # Retrieve relevant medical literature
        literature_insights = self._get_literature_insights(
            predicted_condition, patient_data
        )

        summary = f"""
CLINICAL ASSESSMENT SUMMARY
===========================

PRIMARY DIAGNOSIS: {disease_info["name"]}
Confidence Level: {confidence * 100:.2f}%

DIAGNOSTIC PROBABILITY DISTRIBUTION:
- Healthy Control: {probabilities[0] * 100:.2f}%
- Parkinson's Disease: {probabilities[1] * 100:.2f}%
- SWEDD: {probabilities[2] * 100:.2f}%
- Prodromal PD: {probabilities[3] * 100:.2f}%

CLINICAL DESCRIPTION:
{disease_info["description"]}

KEY CHARACTERISTICS OBSERVED:
"""
        for char in disease_info["characteristics"]:
            summary += f"• {char}\n"

        # Add insights from medical literature with better formatting
        if literature_insights:
            summary += f"\nINSIGHTS FROM MEDICAL LITERATURE:\n"
            summary += "=" * 50 + "\n"
            summary += literature_insights

        return summary

    def _get_literature_insights(self, condition: str, patient_data: Dict) -> str:
        """Retrieve insights from medical literature relevant to the patient's condition."""
        # Check if document manager has documents
        if self.doc_manager.get_document_count()["total"] == 0:
            return "No medical literature available. Add medical papers to enhance insights."

        # Construct search query based on condition and key symptoms
        query_parts = [condition]

        # Always include key symptoms in query with their severity
        symptoms = {
            "tremor": patient_data.get("sym_tremor", 0),
            "rigidity": patient_data.get("sym_rigid", 0),
            "bradykinesia": patient_data.get("sym_brady", 0),
            "postural instability": patient_data.get("sym_posins", 0),
        }

        # Add all symptoms with their severity to create more specific queries
        for symptom, severity in symptoms.items():
            if severity > 0:
                query_parts.append(f"{symptom} severity:{severity}")

        # Add cognitive and psychiatric factors
        if "moca" in patient_data:
            moca = patient_data.get("moca", 30)
            if moca < 26:
                query_parts.append("cognitive impairment")
                if moca < 20:
                    query_parts.append("severe cognitive impairment")

        if "gds" in patient_data and patient_data.get("gds", 0) > 5:
            query_parts.append("depression")

        if "stai" in patient_data and patient_data.get("stai", 0) > 40:
            query_parts.append("anxiety")

        # Add demographic factors if available
        if "age" in patient_data:
            age = patient_data["age"]
            query_parts.append(f"age {age}")
            if age < 50:
                query_parts.append("early onset")
            elif age > 70:
                query_parts.append("elderly")

        if "SEX" in patient_data:
            gender = "male" if patient_data["SEX"] == 1 else "female"
            query_parts.append(gender)

        # Add family history if present
        if patient_data.get("fampd", 0) > 0:
            query_parts.append("family history")

        # Construct final query
        query = " ".join(query_parts)

        # Retrieve relevant passages with increased number of results
        passages = self.doc_manager.extract_relevant_passages(query, top_k=5)

        if not passages:
            return "No specific literature found for this patient's condition and symptoms."

        # Format insights with proper citations and more context
        insights = ""
        for i, passage in enumerate(passages):
            # Extract document metadata for citation
            doc_title = passage["doc_title"]
            # Format the citation properly
            citation = f"[{doc_title}]"

            # Add the passage with citation
            insights += f"{i + 1}. From '{doc_title}': {passage['text'][:400]}...\n   {citation}\n\n"

        return insights

    def generate_feature_analysis(self, patient_data: Dict) -> str:
        """Generate analysis of key patient features."""
        # Use stored original patient data if available
        if hasattr(self, "original_patient_data"):
            patient_data = self.original_patient_data

        analysis = "\nFEATURE ANALYSIS:\n" + "=" * 50 + "\n"

        # Expanded key features list with better labels
        key_features = [
            ("age", "Age"),
            ("SEX", "Gender"),
            ("EDUCYRS", "Education Years"),
            ("BMI", "Body Mass Index"),
            ("fampd", "Family History"),
            ("sym_tremor", "Tremor Severity"),
            ("sym_rigid", "Rigidity"),
            ("sym_brady", "Bradykinesia"),
            ("sym_posins", "Postural Instability"),
            ("moca", "MoCA Score"),
            ("gds", "Depression Score"),
            ("stai", "Anxiety Score"),
            ("ess", "Sleepiness Scale"),
            ("rem", "REM Sleep Behavior"),
        ]

        for feature_key, feature_name in key_features:
            if feature_key in patient_data:
                value = patient_data[feature_key]
                interpretation = self.kb.feature_interpretations.get(feature_key, "")

                # Format value based on feature type
                if feature_key == "age":
                    risk_level = (
                        "High" if value > 60 else "Moderate" if value > 50 else "Low"
                    )
                    analysis += (
                        f"• {feature_name} ({value} years): Risk level: {risk_level}\n"
                    )
                elif feature_key == "SEX":
                    formatted_value = "Male" if value == 1 else "Female"
                    analysis += f"• {feature_name}: {formatted_value}\n"
                elif feature_key == "moca":
                    cognitive_status = (
                        "Normal"
                        if value >= 26
                        else "Mild impairment"
                        if value >= 22
                        else "Significant impairment"
                    )
                    analysis += (
                        f"• {feature_name} ({value}/30): Status: {cognitive_status}\n"
                    )
                elif feature_key == "fampd":
                    family_history = "Positive" if value > 0 else "Negative"
                    analysis += (
                        f"• {feature_name}: {family_history} for Parkinson's disease\n"
                    )
                elif feature_key in [
                    "sym_tremor",
                    "sym_rigid",
                    "sym_brady",
                    "sym_posins",
                ]:
                    severity = ["None", "Mild", "Moderate", "Severe", "Very Severe"]
                    formatted_value = severity[min(int(value), 4)]
                    analysis += f"• {feature_name}: {formatted_value}\n"
                elif feature_key == "gds":
                    status = "Normal" if value <= 5 else "Depression indicated"
                    analysis += f"• {feature_name}: {value} - {status}\n"
                elif feature_key == "stai":
                    status = "Normal" if value <= 40 else "Anxiety indicated"
                    analysis += f"• {feature_name}: {value} - {status}\n"
                else:
                    analysis += f"• {feature_name}: {value}\n"

                # Add interpretation if available and not already added
                if interpretation and feature_key not in ["age", "moca", "fampd"]:
                    analysis = analysis.rstrip("\n") + f" - {interpretation}\n"

        return analysis

    def generate_recommendations(
        self, prediction_results: Dict, patient_data: Dict
    ) -> str:
        """Generate clinical recommendations."""
        # Use stored original patient data if available
        if hasattr(self, "original_patient_data"):
            patient_data = self.original_patient_data

        pred_class = prediction_results["ensemble_prediction"]
        class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
        predicted_condition = class_names[pred_class]

        disease_info = self.kb.disease_info[predicted_condition]

        recommendations = "\nCLINICAL RECOMMENDATIONS:\n" + "=" * 50 + "\n"

        for i, rec in enumerate(disease_info["recommendations"], 1):
            recommendations += f"{i}. {rec}\n"

        # Add general recommendations based on risk factors
        recommendations += "\nADDITIONAL CONSIDERATIONS:\n"

        if patient_data.get("age", 0) > 60:
            recommendations += (
                "• Age-related monitoring: Increased surveillance due to advanced age\n"
            )

        if patient_data.get("fampd", 0) > 0:
            recommendations += (
                "• Genetic counseling: Consider due to positive family history\n"
            )

        if patient_data.get("moca", 30) < 26:
            recommendations += "• Cognitive assessment: Follow-up neuropsychological testing recommended\n"

        if patient_data.get("gds", 0) > 5:
            recommendations += "• Depression management: Consider psychiatric evaluation and treatment\n"

        if patient_data.get("stai", 0) > 40:
            recommendations += (
                "• Anxiety management: Consider psychiatric evaluation and treatment\n"
            )

        if patient_data.get("ess", 0) > 10:
            recommendations += "• Sleep evaluation: Consider sleep study for excessive daytime sleepiness\n"

        if patient_data.get("rem", 0) > 0:
            recommendations += "• REM sleep behavior disorder: Consider polysomnography and treatment\n"

        return recommendations

    def generate_model_consensus(self, prediction_results: Dict) -> str:
        """Generate analysis of model consensus."""
        trad_preds = prediction_results["traditional_predictions"]
        trans_preds = prediction_results["transformer_predictions"]
        ensemble_pred = prediction_results["ensemble_prediction"]

        consensus = "\nMODEL CONSENSUS ANALYSIS:\n" + "=" * 50 + "\n"

        # Check agreement between models
        all_predictions = (
            list(trad_preds.values()) + list(trans_preds.values()) + [ensemble_pred]
        )
        unique_predictions = set(all_predictions)

        if len(all_predictions) == 1:
            consensus += "• LIMITED CONSENSUS: Only the ensemble model was available at inference time\n"
        elif len(unique_predictions) == 1:
            consensus += "• STRONG CONSENSUS: All models agree on the diagnosis\n"
        elif len(unique_predictions) == 2:
            consensus += "• MODERATE CONSENSUS: Most models agree with some variation\n"
        else:
            consensus += "• WEAK CONSENSUS: Significant disagreement between models\n"

        consensus += f"\nIndividual Model Predictions:\n"
        if trad_preds:
            consensus += "Traditional Machine Learning Models:\n"
            for model, pred in trad_preds.items():
                class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
                consensus += f"  • {model.upper()}: {class_names[pred]}\n"
        else:
            consensus += "Traditional Machine Learning Models:\n  • No traditional base-model outputs were available\n"

        consensus += "\nMedical Transformer Models:\n"
        if trans_preds:
            for model, pred in trans_preds.items():
                class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
                model_display = model.replace("_", " ").title()
                if model == "pubmedbert":
                    model_display = "PubMedBERT (Encoder)"
                elif model == "biomistral":
                    model_display = "BioMistral (Decoder)"
                elif model == "clinical_t5":
                    model_display = "Clinical-T5 (Encoder-Decoder)"
                elif model == "feedforward":
                    model_display = "Feedforward Tabular Network"
                consensus += f"  • {model_display}: {class_names[pred]}\n"
        else:
            consensus += "  • No transformer base-model outputs were available\n"

        return consensus

    def generate_full_report(self, patient_data: Dict, patient_id: str = None) -> str:
        """Generate a comprehensive medical report."""
        if patient_id is None:
            patient_id = f"PATIENT_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        else:
            patient_id = str(patient_id)

        # Make predictions
        prediction_results = self.predict_patient(patient_data)

        # Generate report sections
        header = f"""
PARKINSON'S DISEASE ASSESSMENT REPORT
=====================================
Patient ID: {patient_id}
Report Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
Generated by: AI-Powered Multimodal ML System
"""

        clinical_summary = self.generate_clinical_summary(
            prediction_results, patient_data
        )
        feature_analysis = self.generate_feature_analysis(patient_data)
        recommendations = self.generate_recommendations(
            prediction_results, patient_data
        )
        model_consensus = self.generate_model_consensus(prediction_results)

        footer = f"""
DISCLAIMER:
===========
This report is generated by an AI system for research and educational purposes.
It should not replace professional medical diagnosis or treatment decisions.
Always consult with qualified healthcare professionals for medical advice.

Report generated using multimodal machine learning with {prediction_results["confidence"] * 100:.2f}% confidence.
"""

        full_report = (
            header
            + clinical_summary
            + feature_analysis
            + recommendations
            + model_consensus
            + footer
        )

        return full_report

    def save_report(self, report: str, filename: Optional[str] = None) -> str:
        """Save report to file using an absolute path under the project reports directory."""
        if filename is None:
            filename = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        else:
            filename = str(filename)

        safe_name = Path(filename).name.strip()
        safe_name = "".join(
            ch if ch.isalnum() or ch in "._- " else "_" for ch in safe_name
        ).strip(" .")
        if not safe_name:
            safe_name = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        if not safe_name.lower().endswith(".txt"):
            safe_name = f"{safe_name}.txt"

        project_root = Path(__file__).resolve().parents[1]
        reports_dir = project_root / "reports"
        reports_dir.mkdir(parents=True, exist_ok=True)

        filepath = reports_dir / safe_name

        with open(filepath, "w", encoding="utf-8") as f:
            f.write(report)

        resolved_path = str(filepath.resolve())
        print(f"Report saved to: {resolved_path}")
        return resolved_path


def demo_report_generation():
    """Demonstrate report generation with sample patient data."""
    print("RAG System Demo - Generating Sample Medical Report")
    print("=" * 60)

    # Sample patient data
    sample_patient = {
        "age": 65,
        "SEX": 1,  # Male
        "EDUCYRS": 16,
        "race": 1,
        "BMI": 26.5,
        "fampd": 1,  # Positive family history
        "fampd_bin": 1,
        "sym_tremor": 2,
        "sym_rigid": 1,
        "sym_brady": 2,
        "sym_posins": 1,
        "rem": 1,
        "ess": 8,
        "gds": 3,
        "stai": 35,
        "moca": 24,
        "clockdraw": 3,
        "bjlot": 25,
    }

    # Initialize report generator
    report_gen = ReportGenerator()

    try:
        # Generate report
        report = report_gen.generate_full_report(sample_patient, "DEMO_PATIENT_001")

        # Print report
        print(report)

        # Save report
        filepath = report_gen.save_report(report, "demo_medical_report.txt")

        return report, filepath

    except Exception as e:
        print(f"Error generating report: {e}")
        return None, None


if __name__ == "__main__":
    demo_report_generation()