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

"""
Single-shot diagnosis experiment
----------------------------------------
This single-shot experiment implements a direct diagnosis without iterative questioning. 
"""

import datetime
import json
import logging
import asyncio, yaml, time, copy, os
import numpy as np
from scipy.stats import entropy
from typing import List, Dict, Union, Tuple, Any, Optional
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_PROMPT1 = config['doctorai_agent']['differential_diagnosis']['main']
DOCTORAI_INTERACTIVE_DIAGNOSIS_PROMPT = config['doctorai_agent']['differential_diagnosis']['few_shot']
DOCTORAI_INQUIRIES_PROMPT1 = config['doctorai_agent']['inquiries']['main']
DOCTORAI_INQUIRIES_PROMPT = config['doctorai_agent']['inquiries']['multiple_questions']
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):
        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']
        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):
        self.logger = logger
        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
        )

    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, dlogger, 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.dlogger = dlogger

    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, logger, model: str = None):
        super().__init__(
            model=LLAMA_FREE if not model else model,
            system_prompt=EVALUATOR_DIAGNOSIS_PROMPT,
            output_type=Evaluation
        )
        self.logger = logger

    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 sort_result(diagnosis_result: List[Dict]) -> Dict[str, Any]:
    try:
        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
            
        return {
            **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 single_shot_experiment(top_k: int = None, 
                            incomplete: bool = False, 
                            data_source: str = 'nejm', 
                            description: str = None,
                            mask_type: str = None,
                            model: str = None,):
    global INPUT_TOKEN, OUTPUT_TOKEN, TOTAL_TOKEN, DOCTORAI_INTERACTIVE_DIAGNOSIS_PROMPT, DOCTORAI_INQUIRIES_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()
    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=5)
        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)
        temp = True #control whether to feed ddx list to the later predictions
        evaluator = EvaluatorAI(logger, 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 "No masking")
        logger.info(f"Doctor AI Model: {doctorai.__model_name__}")
        logger.info(f"Evaluator Model: {evaluator.__model_name__}")
        logger.info(f"LLM's Settings: {doctorai.setting}")
        if description:
            logger.info(f"Experiment Description: {description}")
    
        # Files and variables initializations (for metrics tracking)
        case_summary_file = initialize_log_file(logger, description, data_source, "case_summaries", log_path)
        total_cases = dataset.__len__()
        correct_diagnosis = 0
        correct_by_level = {'top1': 0, 'top2': 0, 'top3': 0, 'top4': 0, 'top5': 0}
        case_summaries = []

        print("Results can be found in the logs folder.")

        for i in range(total_cases):
            try:
                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} started *************************************')
                log_info(f"************************************************ New Case no. {i+1} (id:{task['metadata']['id']}) started *************************************************")
                logger.info(f"<Task>: {task['partial_patient_info'] if mask_type is not None else task['patient_info']}, <Ground Truth>: {task['ground_truth'] if 'ground_truth' in task else 'N/A'}")
                print(f"Simulating case {i+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 = sort_result(diagnosis_result)

                ### Evaluate correctness of diagnosis ###
                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
                        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                    
                                
                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}')

                # Create case summary
                case_summary = {
                    "type": "case_summary",
                    "case_id": i+1,
                    "case_index": agent_state['patient_case']['metadata']['id'],
                    "confidence_scores": [d['confidence'] for d in diagnosis_result],
                    "highest_confidence": diagnosis_result[0]['confidence'] if diagnosis_result else None,
                    "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)
                
            except Exception as e:
                logger.error(f"Error processing diagnosis for case {i+1}: {str(e)}")
                logger.error(f"Response was: {response}")

        # 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
        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'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'///////////////////////////////////////////////////////////////////////////')
        logger.info(f'************************************************ Experiment ended *********************************************************')
        print(f"Experiment completed.")

        # Prepare experiment results dictionary
        results = {
            "mask_type": mask_type if mask_type else "none",
            "correct_diagnoses": correct_diagnosis,
            "total_cases": total_cases,
            "execution_time": execution_time,
            "input_tokens": INPUT_TOKEN,
            "output_tokens": OUTPUT_TOKEN,
            "average_time_per_case": execution_time/total_cases if total_cases > 0 else 0,
            "top_k_correct_diagnoses": {level: correct_by_level[level] for level in ['top1', 'top2', 'top3', 'top4', 'top5']},
            "top_k_accuracies": top_k_accuracies,
            "experiment_details": {
                "experiment": description,
                "data_source": data_source,
                "model": model if model else "default",
                "top_k": top_k,
                "incomplete": incomplete
            },
            "case_summaries": case_summaries,
        }
        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'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(),
            "mask_type": results.get("mask_type", "none"),
            "total_cases": results.get("total_cases", 0),
            "correct_diagnoses": results.get("correct_diagnoses", 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}")