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

"""
Naive multi-turn baseline
----------------------------------------
This baseline implements a simple multi-turn interaction between a doctor agent and a patient agent. 
"""

import datetime
import json
import logging
import asyncio, yaml, time, os
import numpy as np
from scipy.stats import entropy
from typing import List, Dict, Union, Any
from utils import Logger, MedicalDataLoader, MedqaDataLoader, ModelInitializer, Model, MediqDataLoader, NejmDataLoader
from attributes.evaluation import Evaluation
from attributes.agent_state import AgentState, QA, DifferentialDiagnosis

# Ensure the working directory is set correctly
parent_dir = os.path.dirname(os.path.abspath(__file__))

# Load the configuration from YAML
with open(os.path.join(parent_dir, 'prompts.yaml'), 'rt') as f:
    config = yaml.safe_load(f.read())

# Models
GPT5_1 = "gpt-5.1-2025-11-13"
LLAMA_FREE = "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
LLAMA_8B = "meta-llama/Meta-Llama-3.1-8B-Instruct"
LLAMA_70B = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
DEEPSEEK_R1 = "deepseek-ai/DeepSeek-R1-0528"

# Workflow constants
MAX_NUM_QUESTIONS = 10
CONFIDENCE_THRESHOLD = 0.85
DOCTORAI_INTERACTIVE_DIAGNOSIS_PROMPT = config['doctorai_agent']['differential_diagnosis']['few_shot']
DOCTORAI_INQUIRIES_PROMPT1 = config['doctorai_agent']['inquiries']['main']
PATIENTAI_PROMPT = config['patientai_agent']['answer']
EVALUATOR_DIAGNOSIS_PROMPT = config['evaluator_agent']['evaluate']

# Token Usage Tracking
INPUT_TOKEN = 0
OUTPUT_TOKEN = 0
TOTAL_TOKEN = INPUT_TOKEN + OUTPUT_TOKEN

def log_info(message, logger_name="details_logger", print_to_std=False, mode="info"):
    logger = logging.getLogger(logger_name)
    if logger: 
        if mode == "error": logger.error(message)
        if mode == "warning": logger.warning(message)
        else: logger.info(message)
    if print_to_std: print(message + "\n")

def calculate_token_usage(response):
    """
    Calculate token usage from the response object.
    
    Args:
        response: The response object from the LLM API call.
        
    Returns:
        Tuple[int, int, int]: Input tokens, output tokens, and total tokens.
    """
    global INPUT_TOKEN, OUTPUT_TOKEN, TOTAL_TOKEN
    input_tokens = response.usage.input_tokens
    output_tokens = response.usage.output_tokens
    total_tokens = input_tokens + output_tokens
    
    INPUT_TOKEN += input_tokens
    OUTPUT_TOKEN += output_tokens
    TOTAL_TOKEN += total_tokens
    
    return input_tokens, output_tokens, total_tokens

class BaseAgent:
    def __init__(self, model: str, system_prompt: str, output_type: Any, setting: int = 5):
        """
        Initialize the agent with a specific model.
        
        Args:
            model (str): The model to use for the agent
            system_prompt (str): The system prompt to guide the agent's behavior
            output_type (Any): The expected output type from the model (e.g., str, dict, Pydantic model)
            setting (int): The setting number to determine model parameters (default is 5)
        """
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.output_type = output_type
        if self.__model_name__ == GPT5_1:
            self.setting = ModelInitializer.get_setting_attributes(setting_num=16)
        else:
            self.setting = ModelInitializer.get_setting_attributes(setting_num=setting)
        self.model = Model.create_model(model=self.__model_name__, system_prompt=self.system_prompt, setting=self.setting)
    
    def update_settings(self, **kwargs) -> None:
        """
        Update specific model settings without reinitializing the entire agent.
        
        Args:
            **kwargs: Key-value pairs of settings to update
            
        Examples:
            # Update temperature and top_p
            agent.update_settings(temperature=0.7, top_p=0.9)
            
            # Update system prompt
            agent.update_settings(system_prompt="New system prompt")
            
            # Change model
            agent.update_settings(model="meta-llama/Llama-3.3-70B-Instruct-Turbo")
        """
        # Track if we need to recreate the model
        recreate_model = False
        
        # Update system prompt if provided
        if "system_prompt" in kwargs:
            self.system_prompt = kwargs.pop("system_prompt")
            recreate_model = True
        
        # Update model if provided
        if "model" in kwargs:
            self.__model_name__ = kwargs.pop("model")
            recreate_model = True
        
        # Update output_type if provided
        if "output_type" in kwargs:
            self.output_type = kwargs.pop("output_type")
            recreate_model = True
            
        # Update setting if provided as an integer
        if "setting" in kwargs and isinstance(kwargs["setting"], int):
            self.setting = ModelInitializer.get_setting_attributes(setting_num=kwargs.pop("setting"))
            recreate_model = True
        
        # Update individual settings
        for key, value in kwargs.items():
            if key in self.setting:
                self.setting[key] = value
                recreate_model = True
        
        # Recreate model if necessary
        if recreate_model:
            self.model = Model.create_model(model=self.__model_name__, system_prompt=self.system_prompt, setting=self.setting)

class DoctorAI(BaseAgent):
    def __init__(self, dlogger, model: str = None, baseline: bool = False):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt=DOCTORAI_INTERACTIVE_DIAGNOSIS_PROMPT,
            output_type=DifferentialDiagnosis
        )
        self.dlogger = dlogger

    def __build_prompt(self, temp: bool, agent_state: AgentState, incomplete_information: bool = False, top_k: int = None):
        # Extract relevant information from the state
        patient_case = agent_state['patient_case']
        options = patient_case['options'] if 'options' in patient_case else None
        inquiries_history = agent_state['inquiries_history']
        differential_diagnosis = agent_state['differential_diagnosis']
        differential_diagnosis = [diagnosis['disease'] for diagnosis in differential_diagnosis['items']] if differential_diagnosis else None
        # Determine which patient info to use (masked or full)
        patient_info = patient_case['partial_patient_info'] if incomplete_information else patient_case['patient_info']

        # Build prompt components
        parts = []
        if top_k:
            parts.append(f"(Please give <k>={top_k} differential diagnosis)")
        # Add the basic task information
        parts.append(f"<Task>: {patient_info}")
        if inquiries_history != []:
            parts.append(f"<Inquiries History>: {inquiries_history}")
        if not temp and inquiries_history != [] and differential_diagnosis:
            parts.append(f"<Differential diagnosis>: {differential_diagnosis}")
        
        # Join all parts with newlines
        return "\n".join(parts)
    
    def diagnose(self, temp: bool, agent_state: AgentState, incomplete_information: bool = False, top_k: int = None) -> Dict:
        user_prompt = self.__build_prompt(
            temp=temp,
            agent_state=agent_state, 
            incomplete_information=incomplete_information, 
            top_k=top_k ###### Temporary set to 4 for testing purposes ######
        )
        response = self.model.run(user_prompt=user_prompt, output_type=self.output_type)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()

        if output:
            log_info(f"[DOCTOR]<prompt>: {self.system_prompt} \n\n[DOCTOR]<user_prompt>: {user_prompt} \n\n[DOCTOR]<output>: {output}")  
            self.dlogger.info(f"Differential Diagnosis --> {output['items']}")
        
        return output
    
    async def diagnose_async(self, temp: bool, agent_states: List[AgentState], incomplete_information: bool = False, top_k: int = None):
        user_prompts = []
        for agent_state in agent_states:
            user_prompt = self.__build_prompt(
                temp=temp,  ########## Temporally set to False for async diagnosis ###########
                agent_state=agent_state, 
                incomplete_information=incomplete_information, 
                top_k=top_k
            )
            user_prompts.append(user_prompt)
        
        async_tasks = []
        for i, prompt in enumerate(user_prompts):
            task = asyncio.create_task(self.model.run_async(user_prompt=prompt, output_type=self.output_type))
            async_tasks.append(task)
        responses = await asyncio.gather(*async_tasks)
        
        outputs = []
        for i, response in enumerate(responses):
            calculate_token_usage(response)
            output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()

            log_info(f"[DOCTOR]<prompt>: {self.system_prompt} \n\n[DOCTOR]<user_prompt>: {user_prompt} \n\n[DOCTOR]<output>: {output}")
            self.dlogger.info(f"Simulation {i+1}: {output['items']}...")  

            outputs.append(output)

        return outputs

class DoctorAIInquiries(BaseAgent):
    def __init__(self, logger, model: str = None, baseline: bool = False):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt=f"""You are a medical doctor trying to reason through a real-life clinical case. Based on your understanding of basic and clinical science, medical knowledge, and mechanisms underlying health, disease, and patient care, respond according to the task specified by the user. Base your response on the current and standard practices referenced in medical guidelines.\n\n{config['doctorai_agent']['inquiries']['main_deepseek']}""" if model == DEEPSEEK_R1 else \
                f"""You are a medical doctor trying to reason through a real-life clinical case. Based on your understanding of basic and clinical science, medical knowledge, and mechanisms underlying health, disease, and patient care, respond according to the task specified by the user. Base your response on the current and standard practices referenced in medical guidelines.\n\n{config['doctorai_agent']['inquiries']['main']}""",
            output_type=str
        )
        self.logger = logger

    def ask_question(self, agent_state: AgentState, questions=None) -> str:
        user_prompt = f"<Task>: {agent_state['patient_case']['partial_patient_info']} \n<Inquiries History>: {agent_state['inquiries_history']}"#{[dd['disease'] for dd in list(agent_state['differential_diagnosis']['items'])]}" #\n<Inquiries History>: {agent_state['inquiries_history']}"
        user_prompt += f"\n<Differential diagnosis>: {agent_state['differential_diagnosis']['items']}" if agent_state['differential_diagnosis'] else ""
        
        response = self.model.run(user_prompt=user_prompt, output_type=self.output_type)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()
        log_info(f"[INQUIRY]<prompt>: {self.system_prompt} \n\n[INQUIRY]<user_prompt>: {user_prompt} \n\n[INQUIRY]<output>: {output}")  

        return output
    
    async def ask_questions(self, agent_state: AgentState, questions=None) -> str:
        patient_case = agent_state['patient_case']['partial_patient_info']
        inquiries_history = agent_state['inquiries_history']
        diagnosis = sorted(agent_state['differential_diagnosis']['items'], key=lambda d: d.get('confidence', 0), reverse=True)
        user_prompts = []
        for gen_index in range(5):
            if gen_index == 0 or gen_index == 1 or gen_index == 2 or gen_index == 3:
                disease_a = diagnosis[0]['disease'] if diagnosis else "unknown"
                disease_b = diagnosis[gen_index+1]['disease'] if len(diagnosis) > gen_index+1 else "unknown"
                user_prompt = config['doctorai_agent']['inquiries']['question'].format(patient_case,
                                                                                    disease_a,
                                                                                    disease_b,
                                                                                    inquiries_history,
                                                                                    )
            else:
                current_diagnosis = [diag['disease'] for diag in diagnosis]
                user_prompt = config['doctorai_agent']['inquiries']['question_other'].format(patient_case,
                                                                                    current_diagnosis,
                                                                                    inquiries_history,
                                                                                    )
            user_prompts.append(user_prompt)

        async_tasks = []
        for i, prompt in enumerate(user_prompts):
            task = asyncio.create_task(self.model.run_async(user_prompt=prompt, output_type=self.output_type))
            async_tasks.append(task)
        responses = await asyncio.gather(*async_tasks)
        
        outputs = []
        for i, (prompt, response) in enumerate(zip(user_prompts, responses)):
            calculate_token_usage(response)
            output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()
            log_info(f"[INQUIRY]<prompt>: {self.system_prompt} \n\n[INQUIRY]<user_prompt>: {prompt} \n\n[INQUIRY]<output>: {output}")  
            try:
                output = output.split("Response:")[1].strip() if "Response:" in output else output
            except Exception as e:
                self.logger.error(f"Error extracting question: {e}")
                output = output
            outputs.append(output)

        return outputs

class PatientAI(BaseAgent):
    def __init__(self, model: str = None):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt=PATIENTAI_PROMPT,
            output_type=str
        )

    def answer_question(self, qa_num: int, agent_state: AgentState) -> str:
        """
        Generate an answer to a doctor's question based on patient information.
        
        Args:
            qa_num (int): Index of the question in the inquiries history
            agent_state (AgentState): Current state of the agent
            
        Returns:
            str: The patient's answer
        """
        user_prompt = f"<Task>: {agent_state['patient_case']['patient_info']} \n<Question>: {agent_state['inquiries_history'][qa_num]['question']}"
        response = self.model.run(user_prompt=user_prompt, output_type=self.output_type)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()
        log_info(f"[PATIENT]<prompt>: {self.system_prompt} \n\n[PATIENT]<user_prompt>: {user_prompt} \n\n[PATIENT]<output>: {output}")  

        return output

class DeciderAI(BaseAgent):
    def __init__(self, logger, model: str = None):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt="""You are a doctor. Consider a list of necessary features a doctor would need to make the right medical judgment; finally, consider whether all necessary information is given in the <Patient Case>. How confident are you to make a final diagnosis to the problem factually based on <Patient Case>? Choose between the following ratings:
"Very Confident" - The correct option is supported by all evidence, and there is enough evidence to eliminate the rest of the answers, so the option can be confirmed conclusively.
"Somewhat Confident" - I have reasonably enough information to tell that the correct option is more likely than other options, more information is helpful to make a conclusive decision.
"Neither Confident or Unconfident" - There are evident supporting the correct option, but further evidence is needed to be sure which one is the correct option.
"Somewhat Unconfident" - There are evidence supporting more than one options, therefore more questions are needed to further distinguish the options.
"Very Unconfident" - There are not enough evidence supporting any of the options, the likelihood of picking the correct option at this point is near random guessing.\n\nThink carefully step by step, respond with the chosen confidence rating ONLY and NOTHING ELSE.""",
            output_type=str
        )
        self.logger = logger

    def confidence_abstention(self, agent_state: AgentState, threshold: float = 4.0) -> bool:
        patient_case = agent_state['patient_case']['partial_patient_info']
        user_prompt = f"<Patient Case>: {patient_case}"
        
        response = self.model.run(user_prompt=user_prompt, output_type=self.output_type)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()
        scale = self.parse_likert_scale(output)
        log_info(f"[DECISION]<prompt>: {self.system_prompt} \n\n[DECISION]<user_prompt>: {user_prompt} \n\n[DECISION]<output>: {output} \n[DECISION]<conf>: {scale}")  

        try:
            self.dlogger.info(f"Confidence Abstention --> {output} \nCONF scale: {scale}")
        except (KeyError, TypeError) as e:
            self.dlogger.warning(f"Could not log reasoning or diagnosis: {str(e)}")
            
        if float(scale) >= threshold:
            return True
        else:
            return False
        
    def parse_likert_scale(self, response_text):
        temp_processed_response = response_text.lower().replace('.','').replace(',','').replace(';','').replace(':','')
        if "very confident" in temp_processed_response:
            conf_score = 5
        elif "somewhat confident" in temp_processed_response:
            conf_score = 4
        elif "neither confident nor unconfident" in temp_processed_response:
            conf_score = 3
        elif "neither confident or unconfident" in temp_processed_response:
            conf_score = 3
        elif "somewhat unconfident" in temp_processed_response:
            conf_score = 2
        elif "very unconfident" in temp_processed_response:
            conf_score = 1
        else:
            conf_score = 0
            self.dlogger.info("can't parse likert confidence score: {}".format(response_text), type="error")
        return conf_score

class EvaluatorAI(BaseAgent):
    def __init__(self, model: str = None):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt=EVALUATOR_DIAGNOSIS_PROMPT,
            output_type=Evaluation
        )

    def evaluate(self, ground_truth: str, diagnosis_result: str) -> bool:
        """
        Evaluate if a diagnosis is correct compared to the ground truth.
        
        Args:
            ground_truth (str): The correct diagnosis
            diagnosis_result (str): The model's diagnosis to evaluate
            
        Returns:
            bool: True if the diagnosis is correct, False otherwise
        """
        user_prompt = f"Ground truth: {ground_truth}\nContent to evaluate: {diagnosis_result}"
        response = self.model.run(user_prompt=user_prompt, output_type=self.output_type)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else response.output.to_dict()
        log_info(f"[EVALUATION]<prompt>: {self.system_prompt} \n\n[EVALUATION]<user_prompt>: {user_prompt} \n\n[EVALUATION]<output>: {output}")  

        return output.get('evaluation', None)
        
    def evaluate_all_results(self, ground_truth: str, diagnoses: Dict[str, str]) -> Dict[str, bool]:
        items_to_evaluate = "\n".join([f"{key}: {value}" for key, value in diagnoses.items()])
        user_prompt = f"Ground truth: {ground_truth}\n\nContents to evaluate:\n{items_to_evaluate}\n\n" + \
                     "Return your evaluation as JSON with keys matching the input identifiers and boolean values (true if there is one correct, false if all incorrect)."
                     
        # Using a string output type for this evaluation as we'll parse the JSON response
        self.update_settings(output_type=str)
        
        response = self.model.run(user_prompt=user_prompt, output_type=str)
        calculate_token_usage(response)
        output = response.output if isinstance(response.output, Union[dict, str]) else str(response.output)
        log_info(f"[BATCH_EVALUATION]<prompt>: {self.system_prompt} \n\n[BATCH_EVALUATION]<user_prompt>: {user_prompt} \n\n[BATCH_EVALUATION]<output>: {output}")
        
        # Reset to normal output type
        self.update_settings(output_type=Evaluation)
        
        try:
            import re
            import json
            
            # Try to find a JSON object in the response
            json_match = re.search(r'\{.*\}', output, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
                evaluation_results = json.loads(json_str)
            else:
                # Fallback to manual parsing if JSON not found
                evaluation_results = {}
                for key in diagnoses.keys():
                    evaluation_results[key] = 'true' in output.lower() and key.lower() in output.lower()
                    
            # Convert string 'true'/'false' to boolean values if needed
            for key, value in evaluation_results.items():
                if isinstance(value, str):
                    evaluation_results[key] = value.lower() == 'true'
                    
            return evaluation_results
            
        except Exception as e:
            self.logger.error(f"Error parsing batch evaluation results: {e}")
            # Fallback: perform individual evaluations
            results = {}
            for key, diagnosis in diagnoses.items():
                results[key] = self.evaluate(ground_truth, diagnosis)
            return results
    
def find_correct_result(diagnosis_result: List[Dict], threshold: float = CONFIDENCE_THRESHOLD) -> Dict[str, Any]:
    """
    Determine if there is a clear winner in the differential diagnosis and generate top-k results.
    
    Args:
        diagnosis_result (List[Dict]): List of diagnoses with confidence scores
        threshold (float): Confidence threshold for considering a diagnosis definitive
    
    Returns:
        Dict[str, Any]: Dictionary with early stopping result and top-k results
    """
    try:
        early_result = None
        sorted_results = sorted(diagnosis_result, key=lambda d: d.get('confidence', 0), reverse=True)

        # Generate all top-k results
        top_results = {}
        for k in range(1, 6):  # top-1 through top-5
            if k == 1 and sorted_results:
                top_results[f'top{k}'] = sorted_results[0]['disease']
            elif k <= len(sorted_results):
                top_results[f'top{k}'] = ", ".join(d["disease"] for d in sorted_results[:k])
            else:
                top_results[f'top{k}'] = ", ".join(d["disease"] for d in sorted_results)  # Use all available results
        
        # Check early stopping criteria
        if sorted_results and sorted_results[0]['confidence'] >= 0.97:
            early_result = sorted_results[0]['disease']
        elif len(sorted_results) > 1 and sorted_results[0]['confidence'] - sorted_results[1]['confidence'] >= threshold:
            early_result = sorted_results[0]['disease']
            
        return {
            'early_stop': early_result,
            **top_results,
            'sorted_results': sorted_results
        }
        
    except (IndexError, KeyError) as e:
        return {
            'early_stop': None,
            'top1': None, 'top2': None, 'top3': None, 'top4': None, 'top5': None,
            'sorted_results': []
        }
    
def calculate_entropy(likelihood: List[float]) -> float:
    pk = np.array(likelihood)
    H = entropy(pk, base=2)
    return H

def calculate_difference_sum(original_ddx: List[Dict], simulated_ddx: List[Dict]) -> float:
    """
    Calculate the sum of differences in confidence scores between original and simulated differential diagnoses.
    
    Args:
        original_ddx (List[Dict]): Original differential diagnosis with confidence scores
        simulated_ddx (List[Dict]): Simulated differential diagnoses with confidence scores
        
    Returns:
        float: Sum of differences in confidence scores
    """
    # Convert to dictionaries for easier lookup
    original_confidences = {str.lower(d['disease']): d['confidence'] for d in original_ddx}
    simulated_confidences = {str.lower(d['disease']): d['confidence'] for d in simulated_ddx}
    
    total_difference = 0.0
    for disease, original_confidence in original_confidences.items():
        simulated_confidence = simulated_confidences.get(disease, 0.0)
        total_difference += abs(original_confidence - simulated_confidence)

    # Add any diseases in simulated_ddx that are not in original_ddx
    for disease, simulated_confidence in simulated_confidences.items():
        if disease not in original_confidences:
            total_difference += simulated_confidence
    
    return total_difference

def analyze_entropy_metrics(case_summaries: List[Dict]) -> Dict:
    """
    Analyze entropy reduction and information gain metrics across all cases.
    
    Returns:
        Dict with average metrics and statistics
    """
    total_entropy_reductions = []
    avg_ig_per_query = []
    initial_entropies = []
    final_entropies = []
    all_iteration_ig = []  # All information gains from all iterations
    
    for case in case_summaries:
        if 'total_entropy_reduction' in case and case['total_entropy_reduction'] is not None:
            total_entropy_reductions.append(case['total_entropy_reduction'])
        if 'average_information_gain_per_query' in case and case['average_information_gain_per_query'] is not None:
            avg_ig_per_query.append(case['average_information_gain_per_query'])
        if 'initial_entropy' in case and case['initial_entropy'] is not None:
            initial_entropies.append(case['initial_entropy'])
        if 'final_entropy' in case and case['final_entropy'] is not None:
            final_entropies.append(case['final_entropy'])
        
        # Collect all iteration-level information gains
        if 'iteration_metrics' in case:
            for itr in case['iteration_metrics']:
                if itr.get('entropy_reduction') is not None:
                    all_iteration_ig.append(itr['entropy_reduction'])
    
    metrics = {
        "average_total_entropy_reduction": np.mean(total_entropy_reductions) if total_entropy_reductions else 0,
        "std_total_entropy_reduction": np.std(total_entropy_reductions) if total_entropy_reductions else 0,
        "average_information_gain_per_query": np.mean(avg_ig_per_query) if avg_ig_per_query else 0,
        "std_information_gain_per_query": np.std(avg_ig_per_query) if avg_ig_per_query else 0,
        "average_initial_entropy": np.mean(initial_entropies) if initial_entropies else 0,
        "average_final_entropy": np.mean(final_entropies) if final_entropies else 0,
        "entropy_reduction_percentage": (np.mean(total_entropy_reductions) / np.mean(initial_entropies) * 100) if initial_entropies and total_entropy_reductions else 0,
        "median_iteration_ig": np.median(all_iteration_ig) if all_iteration_ig else 0,
        "max_iteration_ig": np.max(all_iteration_ig) if all_iteration_ig else 0,
        "min_iteration_ig": np.min(all_iteration_ig) if all_iteration_ig else 0,
    }
    
    return metrics

async def baseline_experiment(top_k: int = None, 
             incomplete: bool = False, 
             max_num_questions: int = MAX_NUM_QUESTIONS, 
             data_source: str = 'nejm', 
             description: str = None,
             mask_type: str = None,
             model: str = None,):
    """   
    Naive multi-turn diagnosis workflow without directed question selection.
     
    Args:
        top_k (int, optional): Limit to top-k differential diagnoses
        incomplete (bool): Whether to use partial information
        max_num_questions (int): Maximum number of questions to ask
        data_source (str): Source of medical data
        description (str): Description of the experiment for logging purposes
        mask_type (str): Type of information masking ('Symptoms', 'Laboratory', etc.)
        model (str): Model name to use for all AIs
    
    Returns:
        Dict: Results of the experiment
    """
    global INPUT_TOKEN, OUTPUT_TOKEN, TOTAL_TOKEN, DOCTORAI_INTERACTIVE_DIAGNOSIS_PROMPT, CONFIDENCE_THRESHOLD

    # Initialize logs
    log_path = os.path.join(parent_dir, "logs", f"{description.replace(' ', '_')}_{data_source}_{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}")
    setup_logger('message_logger', log_path, f'{model.split("/")[-1]}_messages.log')
    setup_logger('details_logger', log_path, f'{model.split("/")[-1]}_details.log')
    logger = Logger(log_path).get_system_logger()
    dlogger = Logger(log_path).get_doctor_logger()
    plogger = Logger(log_path).get_patient_logger()
    alogger = Logger(log_path).get_audit_logger()

    # Reset token counters at the beginning of each workflow run
    INPUT_TOKEN = 0
    OUTPUT_TOKEN = 0
    TOTAL_TOKEN = 0
    
    logger.info(f"************************************* {description} started ************************************************")
    # Start timing
    start_time = time.time()
    
    try:
        # Initialize dataset and AI agents
        if data_source == 'nejm':
            dataset = NejmDataLoader().get_evaluation_set(num_cases=1)
        elif data_source == 'mediq':
            dataset = MediqDataLoader().get_evaluation_set(num_cases=129)
        elif data_source == 'medqa':
            dataset = MedqaDataLoader().get_evaluation_set(num_cases=170)
        elif data_source == 'structural_nejm' or data_source == 'structural_mediq' or data_source == 'structural_medqa':
            dataset = MedicalDataLoader(data_source=data_source, mask_type=mask_type, shuffle=False)

        doctorai = DoctorAI(dlogger, model=model, baseline=True)
        doctorai_inquiries = DoctorAIInquiries(logger, model=model, baseline=True)
        temp = True #control whether to feed ddx list to the later predictions
        patientai = PatientAI(model=model)
        evaluator = EvaluatorAI(model=LLAMA_70B)

        # Log experiment configuration
        logger.info(f"Loaded from {data_source} of total {len(dataset)} patient cases with partial information.")
        logger.info(f"Mask Type: {mask_type}" if mask_type else "All-features masking")
        logger.info(f"Doctor AI Model: {doctorai.__model_name__}")
        logger.info(f"Patient AI Model: {patientai.__model_name__}")
        logger.info(f"Evaluator Model: {evaluator.__model_name__}")
        logger.info(f"LLM's Settings: {doctorai.setting}")
        logger.info(f"Confidence Threshold: {CONFIDENCE_THRESHOLD}. Final Confidence: >= 0.97. Max number of questions: {max_num_questions}.")
        if description:
            logger.info(f"Experiment Description: {description}")

        # Files and variables initializations (for metrics tracking)
        iteration_log_file = initialize_log_file(logger, description, data_source, "iteration", log_path=log_path)
        case_summary_file = initialize_log_file(logger, description, data_source, "case_summaries", log_path=log_path)
        total_cases = dataset.__len__()
        total_questions_asked = 0
        correct_diagnosis = 0
        correct_by_level = {'top1': 0, 'top2': 0, 'top3': 0, 'top4': 0, 'top5': 0}
        all_iteration_logs = []
        case_summaries = []
            # --> For 5-iteration evaluation
        case_summary_iter_five_file = initialize_log_file(logger, description, data_source, "case_summaries_iter_five", log_path=log_path) ## For 5 iterations only ##
        correct_diagnosis_iter_five = 0 ## For 5 iterations only ##
        correct_by_level_iter_five = {'top1': 0, 'top2': 0, 'top3': 0, 'top4': 0, 'top5': 0} ## For 5 iterations only ##
        case_summaries_iter_five = []

        print("Results can be found in the logs folder.")
        
        # Process each case
        for i in range(total_cases):
            try:
                num_qa = 0
                task = dataset.get_single_case(i) if data_source == 'structural_nejm' or data_source == 'structural_mediq' or data_source == 'structural_medqa' else dataset[i]
                agent_state = AgentState.create(task)
                logger.info(f'********************** New Case no. {i+1} (id:{task["metadata"]["id"]}) started *************************************')
                log_info(f"************************************************ New Case no. {i+1} (id:{task['metadata']['id']}) started *************************************************")
                logger.info(f"<Task>: {task['partial_patient_info']}, <Ground Truth>: {task['ground_truth'] if 'ground_truth' in task else 'N/A'}")
                print(f"Simulating case {i+1}...")
                
                # Track iterations for this case
                case_iterations = []
                entropy_history = []
                total_information_gain = 0.0
                
                # Initial diagnosis
                try:
                    response = doctorai.diagnose(
                        temp=temp,
                        agent_state=agent_state, 
                        incomplete_information=incomplete, 
                        top_k=top_k,
                    )

                    # Update agent state with diagnosis
                    agent_state['differential_diagnosis'] = response
                    diagnosis_result = agent_state['differential_diagnosis']['items']
                    results_dict = find_correct_result(diagnosis_result)
                    final_result = results_dict['early_stop']
                    
                    # Calculate initial entropy
                    initial_confidences = [d['confidence'] for d in diagnosis_result]
                    initial_entropy = calculate_entropy(initial_confidences)
                    entropy_history.append(initial_entropy)
                    previous_entropy = initial_entropy

                    # Continue asking questions until diagnosis is certain or max questions reached
                    while len(diagnosis_result) > 1 and num_qa < max_num_questions and final_result is None:
                        logger.info(f"Not confident enough to make a final diagnosis, continue asking questions...")
                        
                        try:               
                            try:
                                response = doctorai_inquiries.ask_question(agent_state=agent_state)
                                current_question = response.split("Response:")[1].strip() if "Response:" in response else response
                            except Exception as e:
                                logger.error(f"Error extracting question: {e}")
                                current_question = response  # Fallback to using the entire response
                            dlogger.info(f"Question {num_qa+1} --> {current_question}")
                            
                            # Create Q&A record and add to history
                            qa_record = QA.create(question=current_question)
                            agent_state['inquiries_history'].append(qa_record)

                            # Get answer from the patient
                            response = patientai.answer_question(qa_num=num_qa, agent_state=agent_state) #num_qa if do not have integration step
                            
                            # Parse the answer
                            try:
                                answer = response.split("Response:")[1].strip() if "Response:" in response else response
                            except Exception as e:
                                logger.error(f"Error extracting answer: {e}")
                                answer = response  # Fallback to using the entire response

                            plogger.info(f"Answer {num_qa+1} --> {answer}")
                            agent_state['inquiries_history'][num_qa]['answer'] = answer #num_qa if do not have integration step
                            num_qa += 1

                            response = doctorai.diagnose(
                                temp=temp,
                                agent_state=agent_state, 
                                incomplete_information=incomplete, 
                                top_k=top_k,
                            )

                            agent_state['differential_diagnosis'] = response
                            diagnosis_result = agent_state['differential_diagnosis']['items']
                            results_dict = find_correct_result(diagnosis_result=diagnosis_result)
                            final_result = results_dict['early_stop']
                            
                            # Calculate entropy after this question
                            current_confidences = [d['confidence'] for d in diagnosis_result]
                            current_entropy = calculate_entropy(current_confidences)
                            entropy_history.append(current_entropy)
                            
                            # Calculate entropy reduction (information gain) for this query
                            information_gain = previous_entropy - current_entropy
                            total_information_gain += information_gain
                            previous_entropy = current_entropy

                            logger.info(f"Entropy: {current_entropy:.4f}, Entropy Reduction: {information_gain:.4f}")
                            
                            # Log this iteration's result
                            iteration_log = log_iteration_result(logger,
                                case_index=agent_state['patient_case']['metadata']['id'],
                                case_id=i+1,
                                iteration=num_qa,
                                question=current_question,
                                answer=answer,
                                agent_state=agent_state,
                                differential_diagnosis=diagnosis_result,
                                information_gain=information_gain,
                                current_entropy=current_entropy,
                                ground_truth=task['ground_truth'],
                                log_file_path=iteration_log_file
                            )
                            case_iterations.append(iteration_log)
                            all_iteration_logs.append(iteration_log)
                        
                            if num_qa == 5:
                                diagnoses_to_evaluate = {
                                    'top1': results_dict.get('top1'),
                                    'top2': results_dict.get('top2'),
                                    'top3': results_dict.get('top3'),
                                    'top4': results_dict.get('top4'),
                                    'top5': results_dict.get('top5'),
                                }
                                try:
                                    eval_results = {}
                                    if diagnoses_to_evaluate['top1'].lower() == task['ground_truth'].lower():
                                        eval_results['top1'] = eval_results['top2'] = eval_results['top3'] = eval_results['top4'] = eval_results['top5'] = True
                                    else:
                                        eval_results = evaluator.evaluate_all_results(
                                            ground_truth=task['ground_truth'],
                                            diagnoses=diagnoses_to_evaluate
                                        )
                                        # --> Correctness update
                                    if eval_results.get('top1', False):
                                        correct_diagnosis_iter_five += 1
                                    for level in ['top1', 'top2', 'top3', 'top4', 'top5']:
                                        if eval_results.get(level, False):
                                            correct_by_level_iter_five[level] += 1

                                except Exception as e:
                                    logger.error(f"Error in evaluation for 5 iterations: {e}")

                                # Create case summary
                                case_summary_iter_five = {
                                    "type": "case_summary",
                                    "case_id": i+1,
                                    "case_index": agent_state['patient_case']['metadata']['id'],
                                    "num_questions_asked": num_qa,
                                    "qa_history": agent_state['inquiries_history'],
                                    "confidence_scores": [d['confidence'] for d in diagnosis_result],
                                    "highest_confidence": diagnosis_result[0]['confidence'] if diagnosis_result else None,
                                    "entropy_history": entropy_history,
                                    "initial_entropy": initial_entropy if 'initial_entropy' in locals() else None,
                                    "final_entropy": current_entropy if 'current_entropy' in locals() else None,
                                    "total_entropy_reduction": initial_entropy - current_entropy if 'current_entropy' in locals() else 0.0,
                                    "average_information_gain_per_query": total_information_gain / num_qa if num_qa > 0 else 0.0,
                                    "ground_truth": task['ground_truth'],
                                    "final_diagnosis": {
                                        "top1": results_dict.get('top1'),
                                        "top2": results_dict.get('top2'),
                                        "top3": results_dict.get('top3'),
                                        "top4": results_dict.get('top4'),
                                        "top5": results_dict.get('top5'),
                                    },
                                    "evaluation": eval_results if 'eval_results' in locals() else None,
                                    "topk_correct_diagnoses": {level: correct_by_level_iter_five[level] for level in ['top1', 'top2', 'top3', 'top4', 'top5']},
                                    "topk_accuracies": {"top1": correct_by_level_iter_five['top1'] / (i+1),
                                                        "top2": correct_by_level_iter_five['top2'] / (i+1),
                                                        "top3": correct_by_level_iter_five['top3'] / (i+1),
                                                        "top4": correct_by_level_iter_five['top4'] / (i+1),
                                                        "top5": correct_by_level_iter_five['top5'] / (i+1)},
                                    "correct_diagnoses": correct_diagnosis_iter_five,
                                    "num_cases": i+1,
                                    "accuracy": correct_diagnosis_iter_five / (i+1),
                                }
                                case_summaries_iter_five.append(case_summary_iter_five)
                                
                                # Log case summary immediately
                                log_to_file(logger, case_summary_iter_five, case_summary_iter_five_file)
                        
                        except Exception as e:
                            logger.error(f"Error in Q&A iteration {num_qa+1}: {e}")
                            break  # Exit the Q&A loop if there's an error

                    ### End of q&a loop, making final summaries ###
                    total_questions_asked += num_qa
                    if final_result is None:
                        results_dict = find_correct_result(diagnosis_result)
                    diagnoses_to_evaluate = {
                        'top1': results_dict.get('top1'),
                        'top2': results_dict.get('top2'),
                        'top3': results_dict.get('top3'),
                        'top4': results_dict.get('top4'),
                        'top5': results_dict.get('top5'),
                    }
                    
                    alogger.info(f"Final Judgement <GROUND TRUTH: {task['ground_truth']}> vs. <DIAGNOSIS: {diagnoses_to_evaluate['top1']}>...")
                    try:
                        evaluation_results = {}
                        if diagnoses_to_evaluate['top1'].lower() == task['ground_truth'].lower():
                            evaluation_results['top1'] = evaluation_results['top2'] = evaluation_results['top3'] = evaluation_results['top4'] = evaluation_results['top5'] = True
                        else:
                            evaluation_results = evaluator.evaluate_all_results(
                                ground_truth=task['ground_truth'],
                                diagnoses=diagnoses_to_evaluate
                            )
                            # --> Correctness update
                        if evaluation_results.get('top1', False):
                            correct_diagnosis += 1
                            if num_qa < 5:
                                correct_diagnosis_iter_five += 1 ## For 5 iterations only ##
                            logger.info(f'Total correctness: {correct_diagnosis}/{i+1} (Correct!)')
                        else: 
                            logger.info(f'Total correctness: {correct_diagnosis}/{i+1} (Incorrect.)')
                        for level in ['top1', 'top2', 'top3', 'top4', 'top5']:
                            if evaluation_results.get(level, False):
                                correct_by_level[level] += 1
                                if num_qa < 5:
                                    correct_by_level_iter_five[level] += 1 ## For 5 iterations only ##

                    except Exception as e:
                        logger.error(f"Error in evaluation: {e}")

                    # Statistics announcement on the end of a case
                    accuracy = correct_diagnosis / (i+1)
                    logger.info(f'Accuracy: {accuracy:.4f}')
                    avg_information_gain = total_information_gain / num_qa if num_qa > 0 else 0.0
                    total_entropy_reduction = initial_entropy - current_entropy if 'current_entropy' in locals() else 0.0
                    logger.info(f"Total Entropy Reduction: {total_entropy_reduction:.4f}")
                    logger.info(f"Average Information Gain per Query: {avg_information_gain:.4f}")
                    
                    # Create case summary
                    case_summary = {
                        "type": "case_summary",
                        "case_id": i+1,
                        "case_index": agent_state['patient_case']['metadata']['id'],
                        "num_questions_asked": num_qa,
                        "qa_history": agent_state['inquiries_history'],
                        "confidence_scores": [d['confidence'] for d in diagnosis_result],
                        "highest_confidence": diagnosis_result[0]['confidence'] if diagnosis_result else None,
                        "entropy_history": entropy_history,
                        "initial_entropy": initial_entropy if 'initial_entropy' in locals() else None,
                        "final_entropy": current_entropy if 'current_entropy' in locals() else None,
                        "total_entropy_reduction": total_entropy_reduction,
                        "average_information_gain_per_query": avg_information_gain,
                        "ground_truth": task['ground_truth'],
                        "final_diagnosis": {
                            "top1": results_dict.get('top1'),
                            "top2": results_dict.get('top2'),
                            "top3": results_dict.get('top3'),
                            "top4": results_dict.get('top4'),
                            "top5": results_dict.get('top5'),
                        },
                        "evaluation": evaluation_results if 'evaluation_results' in locals() else None,
                        "topk_correct_diagnoses": {level: correct_by_level[level] for level in ['top1', 'top2', 'top3', 'top4', 'top5']},
                        "topk_accuracies": {"top1": correct_by_level['top1'] / (i+1),
                                            "top2": correct_by_level['top2'] / (i+1),
                                            "top3": correct_by_level['top3'] / (i+1),
                                            "top4": correct_by_level['top4'] / (i+1),
                                            "top5": correct_by_level['top5'] / (i+1)},
                        "correct_diagnoses": correct_diagnosis,
                        "num_cases": i+1,
                        "accuracy": correct_diagnosis / (i+1),
                    }
                    case_summaries.append(case_summary)
                    log_to_file(logger, case_summary, case_summary_file)

                    if num_qa < 5:
                        case_summary_iter_five = {
                            "type": "case_summary",
                            "case_id": i+1,
                            "case_index": agent_state['patient_case']['metadata']['id'],
                            "num_questions_asked": num_qa,
                            "reconstructed_case": agent_state['patient_case']['partial_patient_info'],
                            "qa_history": agent_state['questions_history'],
                            "confidence_scores": [d['confidence'] for d in diagnosis_result],
                            "highest_confidence": diagnosis_result[0]['confidence'] if diagnosis_result else None,
                            "entropy_history": entropy_history,
                            "initial_entropy": initial_entropy if 'initial_entropy' in locals() else None,
                            "final_entropy": current_entropy if 'current_entropy' in locals() else None,
                            "total_entropy_reduction": total_entropy_reduction,
                            "average_information_gain_per_query": avg_information_gain,
                            "ground_truth": task['ground_truth'],
                            "final_diagnosis": {
                                "top1": results_dict.get('top1'),
                                "top2": results_dict.get('top2'),
                                "top3": results_dict.get('top3'),
                                "top4": results_dict.get('top4'),
                                "top5": results_dict.get('top5'),
                            },
                            "evaluation": evaluation_results if 'evaluation_results' in locals() else None,
                            "topk_correct_diagnoses": {level: correct_by_level_iter_five[level] for level in ['top1', 'top2', 'top3', 'top4', 'top5']},
                            "topk_accuracies": {"top1": correct_by_level_iter_five['top1'] / (i+1),
                                                "top2": correct_by_level_iter_five['top2'] / (i+1),
                                                "top3": correct_by_level_iter_five['top3'] / (i+1),
                                                "top4": correct_by_level_iter_five['top4'] / (i+1),
                                                "top5": correct_by_level_iter_five['top5'] / (i+1)},
                            "correct_diagnoses": correct_diagnosis_iter_five,
                            "num_cases": i+1,
                            "accuracy": correct_diagnosis_iter_five / (i+1),
                        }
                        case_summaries_iter_five.append(case_summary_iter_five)
                        log_to_file(logger, case_summary_iter_five, case_summary_iter_five_file)
                
                except Exception as e:
                    logger.error(f"Error processing case {i+1}: {e}")
            except Exception as e:
                logger.error(f"Fatal error in case {i+1}, skipping: {e}")
                continue
                
        # End of all cases loop -> Final evaluation (time, accuracy, tokens, etc.)
        logger.info(f'////////////////////////// Experiment Summary ///////////////////////////////')
        end_time = time.time()
        execution_time = end_time - start_time
        avg_num_questions = total_questions_asked / total_cases if total_cases > 0 else 0.0
        top_k_accuracies = {"top1": correct_by_level['top1'] / total_cases,
                            "top2": correct_by_level['top2'] / total_cases,
                            "top3": correct_by_level['top3'] / total_cases,
                            "top4": correct_by_level['top4'] / total_cases,
                            "top5": correct_by_level['top5'] / total_cases}
        for level, acc in top_k_accuracies.items():
            logger.info(f'{level.capitalize()} accuracy: {acc:.4f} ({correct_by_level.get(level, 0)}/{total_cases})')
        logger.info(f'Average number of questions asked: {avg_num_questions:.1f}')
        logger.info(f'Total tokens usage: {TOTAL_TOKEN}, Input tokens: {INPUT_TOKEN}, Output tokens: {OUTPUT_TOKEN}')
        logger.info(f'Runtime: {execution_time:.2f} seconds ({execution_time/60:.2f} minutes)')
        logger.info(f'Average time per case: {execution_time/total_cases:.2f} seconds')
            # --> Calculate and log entropy metrics
        entropy_metrics = analyze_entropy_metrics(case_summaries)
        logger.info(f'Average Total Entropy Reduction: {entropy_metrics["average_total_entropy_reduction"]:.4f} ± {entropy_metrics["std_total_entropy_reduction"]:.4f}')
        logger.info(f'Average Information Gain per Query: {entropy_metrics["average_information_gain_per_query"]:.4f} ± {entropy_metrics["std_information_gain_per_query"]:.4f}')
        logger.info(f'Average Initial Entropy: {entropy_metrics["average_initial_entropy"]:.4f}')
        logger.info(f'Average Final Entropy: {entropy_metrics["average_final_entropy"]:.4f}')
        logger.info(f'Entropy Reduction Percentage: {entropy_metrics["entropy_reduction_percentage"]:.2f}%')
        logger.info(f'///////////////////////////////////////////////////////////////////////////')
        logger.info(f'************************************************ Experiment ended *********************************************************')
        print(f"Experiment completed.")
        
        # Prepare experiment results dictionary
        results = {
            "correct_diagnoses": correct_diagnosis,
            "total_cases": total_cases,
            "execution_time": execution_time,
            "input_tokens": INPUT_TOKEN,
            "output_tokens": OUTPUT_TOKEN,
            "average_num_questions": avg_num_questions,
            "average_time_per_case": execution_time/total_cases if total_cases > 0 else 0,
            "top_k_accuracies": top_k_accuracies,
            "entropy_metrics": entropy_metrics,
            "experiment_details": {
                "experiment": description,
                "data_source": data_source,
                "model": model if model else "default",
                "max_num_questions": max_num_questions,
                "top_k": top_k,
                "incomplete": incomplete
            },
            "case_summaries": case_summaries,
            "all_iterations": all_iteration_logs
        }
        finalize_log_file(logger, iteration_log_file, results)
        finalize_log_file(logger, case_summary_file, results)

        return results
        
    except Exception as e:
        logger.error(f"Fatal error in workflow: {e}")
        logger.info(f'Average number of questions asked: {avg_num_questions:.1f}')
        logger.info(f'Total tokens usage: {TOTAL_TOKEN}, Input tokens: {INPUT_TOKEN}, Output tokens: {OUTPUT_TOKEN}')
        logger.info(f'Runtime: {execution_time:.2f} seconds ({execution_time/60:.2f} minutes)')
        logger.info(f'Average time per case: {execution_time/total_cases:.2f} seconds')
        logger.info(f'************************************************ Experiment ended *********************************************************')
        raise

def setup_logger(name, log_path, filename):
    if not log_path: return None
    os.makedirs(log_path, exist_ok=True)
    logger = logging.getLogger(name)
    handler = logging.FileHandler(os.path.join(log_path, filename), mode='a')
    formatter = logging.Formatter('[%(asctime)s] [%(levelname)s] %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)
    return logger

def finalize_log_file(logger, log_file_path: str, results: dict) -> None:
    """
    Add experiment summary to a log file.
    """
    if not log_file_path or not os.path.exists(log_file_path):
        return
    
    try:
        experiment_summary = {
            "type": "experiment_summary",
            "end_time": datetime.datetime.now().isoformat(),
            "total_cases": results.get("total_cases", 0),
            "correct_diagnoses": results.get("correct_diagnoses", 0),
            "accuracy": results.get("accuracy", 0),
            "average_num_questions": results.get("average_num_questions", 0),
            "execution_time": results.get("execution_time", 0),
            "input_tokens": results.get("input_tokens", 0),
            "output_tokens": results.get("output_tokens", 0),
            "top_k_correct_diagnoses": results.get("top_k_correct_diagnoses", {}),
            "top_k_accuracies": results.get("top_k_accuracies", {}),
            "experiment_details": results.get("experiment_details", {})
        }
        
        with open(log_file_path, 'a', encoding='utf-8') as f:
            f.write(json.dumps(experiment_summary, ensure_ascii=False) + '\n')
            
    except Exception as e:
        logger.error(f"Error finalizing log file: {e}")

def initialize_log_file(logger, description: str, data_source: str, log_type: str, log_path: str) -> str:
    """
    Initialize a JSONL file for logging.
    """
    os.makedirs(log_path, exist_ok=True)
    
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    exp_description = (description or "experiment").replace(" ", "_").replace("/", "-").lower() + f"_{data_source}"
    filename = f"{exp_description}_{log_type}_{timestamp}.jsonl"
    file_path = os.path.join(log_path, filename)
    
    metadata = {
        "type": "metadata",
        "description": description + f" ({data_source})",
        "start_time": datetime.datetime.now().isoformat(),
        "timestamp": timestamp
    }
    
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(json.dumps(metadata, ensure_ascii=False) + '\n')
        logger.info(f"{log_type.capitalize()} log file initialized at: {file_path}")
        return file_path
    except Exception as e:
        logger.error(f"Error initializing {log_type} log file: {e}")
        return None

def log_to_file(logger, data: Dict, log_file_path: str) -> None:
    """
    Append data to a JSONL file.
    """
    if not log_file_path:
        return
    
    try:
        with open(log_file_path, 'a', encoding='utf-8') as f:
            f.write(json.dumps(data, ensure_ascii=False) + '\n')
    except Exception as e:
        logger.error(f"Error writing to file: {e}")

def log_iteration_result(logger, case_index: int, case_id: int, iteration: int, question: str, answer: str, 
                         agent_state: Dict, differential_diagnosis: List[Dict], information_gain: float = None,
                         current_entropy: float = None, selected_question_index: int = None,
                         all_questions: List[str] = None, all_information_gains: List[float] = None,
                         ground_truth: str = None, log_file_path: str = None) -> Dict:
    """
    Log the result of a single iteration.
    """
    itr = {
        "type": "iteration",
        "case_id": case_id,
        "case_index": case_index,
        "ground_truth": ground_truth,
        "timestamp": datetime.datetime.now().isoformat(),
        "num_of_iterations": iteration,
        "question": question,
        "answer": answer,
        "reconstructed_case": agent_state['patient_case']['partial_patient_info'],
        "complete_case": agent_state['patient_case']['patient_info'],
        "qa_history": agent_state['inquiries_history'],
        "differential_diagnosis": differential_diagnosis,
        "top_diagnosis": differential_diagnosis[0] if differential_diagnosis else None,
        "possible_diseases": [item['disease'] for item in differential_diagnosis] if differential_diagnosis else [],
        "conf_list": [item['confidence'] for item in differential_diagnosis] if differential_diagnosis else []
    }
    
    if information_gain is not None:
        itr["information_gain"] = float(information_gain)
    
    if current_entropy is not None:
        itr["current_entropy"] = float(current_entropy)
    
    if selected_question_index is not None:
        itr["selected_question_index"] = selected_question_index
        
    if all_questions:
        itr["candidate_questions"] = all_questions
        
    if all_information_gains:
        itr["candidate_information_gains"] = [float(ig) for ig in all_information_gains]
        if all_questions and len(all_questions) == len(all_information_gains):
            ranked = sorted(zip(all_questions, all_information_gains), 
                          key=lambda x: x[1], reverse=True)
            itr["ranked_questions"] = [
                {"question": q, "information_gain": float(ig)} for q, ig in ranked
            ]
    
    log_to_file(logger, itr, log_file_path)
    return itr