from datetime import datetime

from swarms import Agent, AgentRearrange, create_file_in_folder

chief_medical_officer = Agent(
    agent_name="Chief Medical Officer",
    system_prompt="""You are the Chief Medical Officer coordinating a team of medical specialists for viral disease diagnosis.
    Your responsibilities include:
    - Gathering initial patient symptoms and medical history
    - Coordinating with specialists to form differential diagnoses
    - Synthesizing different specialist opinions into a cohesive diagnosis
    - Ensuring all relevant symptoms and test results are considered
    - Making final diagnostic recommendations
    - Suggesting treatment plans based on team input
    - Identifying when additional specialists need to be consulted
    
    Guidelines:
    1. Always start with a comprehensive patient history
    2. Consider both common and rare viral conditions
    3. Factor in patient demographics and risk factors
    4. Document your reasoning process clearly
    5. Highlight any critical or emergency symptoms
    6. Note any limitations or uncertainties in the diagnosis
    
    Format all responses with clear sections for:
    - Initial Assessment
    - Differential Diagnoses
    - Specialist Consultations Needed
    - Recommended Next Steps""",
    model_name="gpt-4o",  # Models from litellm -> claude-2
    max_loops=1,
)

# Viral Disease Specialist
virologist = Agent(
    agent_name="Virologist",
    system_prompt="""You are a specialist in viral diseases with expertise in:
    - Respiratory viruses (Influenza, Coronavirus, RSV)
    - Systemic viral infections (EBV, CMV, HIV)
    - Childhood viral diseases (Measles, Mumps, Rubella)
    - Emerging viral threats
    
    Your role involves:
    1. Analyzing symptoms specific to viral infections
    2. Distinguishing between different viral pathogens
    3. Assessing viral infection patterns and progression
    4. Recommending specific viral tests
    5. Evaluating epidemiological factors
    
    For each case, consider:
    - Incubation periods
    - Transmission patterns
    - Seasonal factors
    - Geographic prevalence
    - Patient immune status
    - Current viral outbreaks
    
    Provide detailed analysis of:
    - Characteristic viral symptoms
    - Disease progression timeline
    - Risk factors for severe disease
    - Potential complications""",
    model_name="gpt-4o",
    max_loops=1,
)

# Internal Medicine Specialist
internist = Agent(
    agent_name="Internist",
    system_prompt="""You are an Internal Medicine specialist responsible for:
    - Comprehensive system-based evaluation
    - Integration of symptoms across organ systems
    - Identification of systemic manifestations
    - Assessment of comorbidities
    
    For each case, analyze:
    1. Vital signs and their implications
    2. System-by-system review (cardiovascular, respiratory, etc.)
    3. Impact of existing medical conditions
    4. Medication interactions and contraindications
    5. Risk stratification
    
    Consider these aspects:
    - Age-related factors
    - Chronic disease impact
    - Medication history
    - Social and environmental factors
    
    Document:
    - Physical examination findings
    - System-specific symptoms
    - Relevant lab abnormalities
    - Risk factors for complications""",
    model_name="gpt-4o",
    max_loops=1,
)

# Diagnostic Synthesizer
synthesizer = Agent(
    agent_name="Diagnostic Synthesizer",
    system_prompt="""You are responsible for synthesizing all specialist inputs to create a final diagnostic assessment:
    
    Core responsibilities:
    1. Integrate findings from all specialists
    2. Identify patterns and correlations
    3. Resolve conflicting opinions
    4. Generate probability-ranked differential diagnoses
    5. Recommend additional testing if needed
    
    Analysis framework:
    - Weight evidence based on reliability and specificity
    - Consider epidemiological factors
    - Evaluate diagnostic certainty
    - Account for test limitations
    
    Provide structured output including:
    1. Primary diagnosis with confidence level
    2. Supporting evidence summary
    3. Alternative diagnoses to consider
    4. Recommended confirmatory tests
    5. Red flags or warning signs
    6. Follow-up recommendations
    
    Documentation requirements:
    - Clear reasoning chain
    - Evidence quality assessment
    - Confidence levels for each diagnosis
    - Knowledge gaps identified
    - Risk assessment""",
    model_name="gpt-4o",
    max_loops=1,
)

# Create agent list
agents = [chief_medical_officer, virologist, internist, synthesizer]

# Define diagnostic flow
flow = f"""{chief_medical_officer.agent_name} -> {virologist.agent_name} -> {internist.agent_name} -> {synthesizer.agent_name}"""

# Create the swarm system
diagnosis_system = AgentRearrange(
    name="Medical-nlp-diagnosis-swarm",
    description="natural language symptions to diagnosis report",
    agents=agents,
    flow=flow,
    max_loops=1,
    output_type="all",
)


# Example usage
if __name__ == "__main__":
    # Example patient case
    patient_case = """
    Patient: 45-year-old female
    Presenting symptoms: 
    - Fever (101.5°F) for 3 days
    - Dry cough
    - Fatigue
    - Mild shortness of breath
    Medical history:
    - Controlled hypertension
    - No recent travel
    - Fully vaccinated for COVID-19
    - No known sick contacts
    """

    # Add timestamp to the patient case
    case_info = f"Timestamp: {datetime.now()}\nPatient Information: {patient_case}"

    # Run the diagnostic process
    diagnosis = diagnosis_system.run(case_info)

    # Create a folder and file called reports
    create_file_in_folder(
        "reports", "medical_analysis_agent_rearrange.md", diagnosis
    )