#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Main entry point
"""

import argparse
import asyncio
import os
import sys
from typing import Dict
from baseline import baseline_experiment
from single_shot import single_shot_experiment
from MedClarify import medclarify_experiment

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

class ExperimentRunner:
    """Run experiments from configuration file or CLI arguments."""
    
    def __init__(self):
        """Initialize the experiment runner."""
        pass
    
    def run_experiments(self, exp: Dict) -> Dict:
        """Run an experiment."""

        if exp.get('baseline', False):
            if exp.get('mask_type', None):
                print("-" * 80)
                print(f"Running naive multi-turn baseline (single-feature masking on {exp.get('mask_type', None)})...")
                print("-" * 80)
                asyncio.run(baseline_experiment(
                    top_k=exp.get('top_k'), 
                    incomplete=True,
                    max_num_questions=exp.get('max_num_questions'),
                    data_source=f"structural_{exp.get('dataset')}",
                    description=f"naive multi-turn baseline (single-feature masking on {exp.get('mask_type')})",
                    mask_type=exp.get('mask_type'),
                    model=exp.get('model')
                ))
            else:
                print("-" * 80)
                print(f"Running naive multi-turn baseline (all-features masking)...")
                print("-" * 80)
                asyncio.run(baseline_experiment(
                    top_k=exp.get('top_k'),
                    incomplete=True,
                    max_num_questions=exp.get('max_num_questions'),
                    data_source=exp.get('dataset'),
                    description='naive multi-turn baseline (all-features masking)',
                    model=exp.get('model')
                ))
        elif exp.get('single_shot_complete', False):
            print("-" * 80)
            print(f"Running single-shot diagnosis on complete cases...")
            print("-" * 80)
            single_shot_experiment(
                top_k=exp.get('top_k'),
                incomplete=False,
                data_source=f"structural_{exp.get('dataset')}",
                description='single shot diagnosis on complete cases',
                mask_type=exp.get('mask_type'),
                model=exp.get('model')
            )
        elif exp.get('single_shot_incomplete', False):
            print("-" * 80)
            print(f"Running single-shot diagnosis on incomplete cases...")
            print("-" * 80)
            single_shot_experiment(
                top_k=exp.get('top_k'),
                incomplete=True,
                data_source=f"structural_{exp.get('dataset')}",
                description='single shot diagnosis on incomplete cases',
                mask_type=exp.get('mask_type'),
                model=exp.get('model')
            )
        elif exp.get('medclarify', False):
            if exp.get('mask_type', None):
                print("-" * 80)
                print(f"Running MedClarify (single-feature masking on {exp.get('mask_type', None)})...")
                print("-" * 80)
                asyncio.run(medclarify_experiment(
                    top_k=exp.get('top_k'), 
                    incomplete=True,
                    max_num_questions=exp.get('max_num_questions'),
                    data_source=f"structural_{exp.get('dataset')}",
                    description=f"MedClarify (single-feature masking on {exp.get('mask_type')})",
                    mask_type=exp.get('mask_type'),
                    model=exp.get('model')
                ))
            else:
                print("-" * 80)
                print(f"Running MedClarify (all-features masking)...")
                print("-" * 80)
                asyncio.run(medclarify_experiment(
                    top_k=exp.get('top_k'),
                    incomplete=True,
                    max_num_questions=exp.get('max_num_questions'),
                    data_source=exp.get('dataset'),
                    description='MedClarify (all-features masking)',
                    model=exp.get('model')
                ))
            
        


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        description='Run medical diagnosis experiments',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run specific experiment
  python main.py --baseline --dataset nejm --max_num_questions 10

  # Run with specific model
  python main.py --baseline --model gpt-5.1-2025-11-13
  
  # Run single-shot diagnosis on complete cases
  python main.py --single_shot_complete 

  # Run single-shot diagnosis on incomplete cases with specific feature masking
  python main.py --single_shot_incomplete --mask_type symptoms
        """
    )
        
    parser.add_argument('--baseline', action='store_true', help='Run the naive multi-turn baseline experiment')
    parser.add_argument('--medclarify', action='store_true', help='Run the MedClarify experiment')
    parser.add_argument('--single_shot_complete', action='store_true', help='Run single-shot diagnosis on complete cases')
    parser.add_argument('--single_shot_incomplete', action='store_true', help='Run single-shot diagnosis on incomplete cases')
    parser.add_argument('--dataset', type=str, default='nejm', help='The dataset to use for experiments')
    parser.add_argument('--top_k', type=int, default=5, help='Top-k differential diagnoses to consider')
    parser.add_argument('--max_num_questions', type=int, default=10, help='Maximum number of questions to ask')
    parser.add_argument('--mask_type', type=str, default=None, help='Type of information masking (e.g., Symptoms, Laboratory, Imaging, Past_Medical_History, Physical_Examination, Social_History)')
    parser.add_argument('--model', type=str, default='meta-llama/Llama-3.3-70B-Instruct-Turbo', help='Model name to use for all AIs')

    args = parser.parse_args()
    runner = ExperimentRunner()
    exp = {
        'baseline': args.baseline,
        'medclarify': args.medclarify,
        'single_shot_complete': args.single_shot_complete,
        'single_shot_incomplete': args.single_shot_incomplete,
        'dataset': args.dataset, 
        'top_k': args.top_k,  
        'max_num_questions': args.max_num_questions,
        'mask_type': args.mask_type, 
        'model': args.model
    }

    if args.single_shot_incomplete and not args.mask_type:
        parser.error("--mask_type is required when using --single_shot_incomplete. \nChoose one of the medical information categories to mask (e.g., Symptoms, Laboratory, Imaging, Past_Medical_History, Physical_Examination, Social_History).")

    print(f"Running experiment with configuration: {exp}")
    try:
        runner.run_experiments(exp)
    except Exception as e:
        print(f"Error running experiments: {e}")
        sys.exit(1)



if __name__ == '__main__':
    main()