import requests
from bs4 import BeautifulSoup
from attributes.agent_state import PatientCase
import os, yaml, time, random, asyncio, aiohttp, nest_asyncio, functools, inspect, json, copy, tiktoken, logging.config, re, urllib3
from typing import Dict, List, Any, Optional, Union
from together import AsyncTogether, Together
from openai import OpenAI, AsyncOpenAI
from pydantic import BaseModel, ValidationError
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider
import ollama, anthropic, openai
from google import genai
from google.genai import types
from datetime import datetime
from dotenv import load_dotenv
from similarity_mapping import similarity_scores

load_dotenv()
nest_asyncio.apply()
DIR_PATH = os.path.dirname(os.path.abspath(__file__))

class Usage:
    def __init__(self, input_tokens: int, output_tokens: int):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens

class Response:
    """
    Response class to handle model responses.
    
    This class provides a unified interface for handling responses from different
    language models, including methods for parsing and formatting the response.
    
    Attributes:
        content (str): The content of the response
        model (str): The model that generated the response
        usage (Dict[str, int]): Token usage statistics
    """
    
    def __init__(self, output: str, model: str, usage: Usage, body):
        self.output = output
        self.model = model
        self.usage = usage
        self.body = body

class ModelInitializer:
    """
    A unified interface for initializing and running different language models.
    
    This class provides a consistent API for working with various LLM providers including:
    - OpenAI models (GPT-4o)
    - Anthropic models (Claude)
    - Together.ai models (Llama, DeepSeek, Qwen)
    - Locally hosted models via Ollama
    - Google Gemini models
    
    Attributes:
        model (str): The identifier for the model to use
        system_prompt (Optional[str]): Default system prompt to use with the model
        gemini_models (List[str]): List of supported Google Gemini model identifiers
        openai_models (List[str]): List of supported OpenAI model identifiers
        anthropic_models (List[str]): List of supported Anthropic model identifiers
        together_ai_models (List[str]): List of supported Together.ai model identifiers
        ollama_local_models (List[str]): List of supported local Ollama model identifiers
        client: The initialized client for the selected model provider
    """

    def __init__(self, model:str, system_prompt:Optional[str] = None):
        """
        Initialize the ModelInitializer with a specific model.
        
        Args:
            model (str): The identifier for the model to use
            system_prompt (Optional[str]): Default system prompt to use with the model
        """
        self.model = model
        self.system_prompt = system_prompt
        self.chutes_models = ["deepseek-ai/DeepSeek-V3-0324", "deepseek-ai/DeepSeek-R1"]
        self.gemini_models = ["gemini-2.5-pro-exp-03-25", "gemini-2.5-flash-preview-04-17"]
        self.openai_models = ["gpt-4o-2024-08-06"]
        self.anthropic_models = ["claude-3-7-sonnet-latest"]
        self.together_ai_models = ["meta-llama/Llama-3.3-70B-Instruct-Turbo-Free", "meta-llama/Llama-3.3-70B-Instruct-Turbo", "deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free", "Qwen/QwQ-32B", "Qwen/Qwen2.5-72B-Instruct-Turbo"]
        self.ollama_local_models = ["phi4", "qwq:32b"]
        self.client = self.init_client()
        self.token_tracker = TokenTracker()


    def init_client(self):
        """
        Initialize the appropriate client based on the selected model.
        
        Returns:
            The initialized client object for the selected model provider
        
        Raises:
            KeyError: If the required API key environment variable is not set
        """
        if self.model in self.together_ai_models: 
            return Together(api_key=os.environ["TOGETHER_API_KEY"])
        elif self.model in self.ollama_local_models:
            return ollama.Client()
        elif self.model in self.anthropic_models:
            return anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
        elif self.model in self.openai_models:
            return openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
        elif self.model in self.gemini_models:
            return genai.Client(api_key=os.environ["GEMINI_API_KEY"])
        elif self.model in self.chutes_models:
            headers = {
                "Authorization": "Bearer " + os.environ["CHUTES_API_KEY"],
                "Content-Type": "application/json"
            }
            return headers
    
    @classmethod
    def find_setting(cls, setting_num: int):
        if setting_num == 1:
            return {'temperature': 1.97, 'min_p': 0.05, 'seed': 42}
        if setting_num == 2:
            return {'temperature': 0.3, 'min_p': 0.1, 'seed': 42} # deepseek suggested, low temperature for precision, low min_p for reducing hallucinations, high top_p for diversity
        elif setting_num == 3:
            return {'temperature': 1.97, 'min_p': 0.05, 'top_p': 0.9, 'seed': 42} # temperature, min_p, top_p, seed
        elif setting_num == 4:
            return {'temperature': 0.3, 'min_p': 0.1, 'top_p': 0.9, 'seed': 42} 
        elif setting_num == 5:
            # best settings for llama 3.3
            return {'temperature': 0.3, 'min_p': 0.1, 'top_p': 0.9, 'repetition_penalty': 0.9, 'seed': 42} # temperature, min_p, top_p, repetition_penalty, seed
        elif setting_num == 6:
            return {'temperature': 1.97, 'min_p': 0.05, 'top_p': 0.9, 'repetition_penalty': 0.9, 'seed': 42} 
        elif setting_num == 7:
            return {'temperature': 1.97, 'min_p': 0.05, 'top_p': 0.6, 'repetition_penalty': 0.9, 'seed': 42}
        elif setting_num == 8:
            return {'temperature': 0.3,'min_p': 0.05, 'top_p': 0.9,'repetition_penalty': 1.2,'seed': 42}
        elif setting_num == 9:
            return {'temperature': 0.3,'min_p': 0.05, 'top_p': 0.85, 'repetition_penalty': 1, 'seed': 42}
        elif setting_num == 10:
            return {'temperature': 0.3,'min_p': 0.1, 'top_p': 0.85,'repetition_penalty': 0.95,'seed': 42}
        elif setting_num == 11:
            return {'temperature': 0.5,'seed': 42}
        elif setting_num == 12:
            return {'temperature': 0.5, 'top_p': 0.9, 'repetition_penalty': 0.9}
        elif setting_num == 13:
            return {'temperature': 0.8,'seed': 42}
        elif setting_num == 14:
            return {
                'temperature': 0.3,  # Keeps original creativity level
                'top_p': 0.8,       # More restrictive than 0.9 to mimic min_p filtering
                'frequency_penalty': 0.2,  # Reduces repetition (like repetition_penalty=0.9)
                'presence_penalty': 0.1,   # Mild penalty for new tokens already in text
                'seed': 42          # Keep deterministic if possible
            }
        elif setting_num == 15:
            return {'temperature': 0.5, 'min_p': 0.08, 'top_p': 0.9, 'repetition_penalty': 1.0, 'seed': 42}
        elif setting_num == 16:
            return {'temperature': 0.3, 'top_p': 0.9, 'seed': 42}
    
    @classmethod
    def get_setting_attributes(cls, setting_num: int):
        return cls.find_setting(setting_num)

    def run_llm(self, user_prompt : str, system_prompt : Optional[str] = None, setting: int = None):
        """
        Run inference with the language model using the provided prompt.
        
        This method handles the differences between various LLM APIs and provides
        a unified interface for text generation. It automatically tracks token usage.
        
        Args:
            user_prompt (str): The user prompt to send to the model
            system_prompt (Optional[str]): System prompt to override the default
            setting (int, optional): Setting number for model configurations
            
        Returns:
            str: The generated text response from the model
            
        Note:
            Default parameters include temperature=0.7 and max_tokens=700
            Token usage is automatically counted and logged
        """
        setting = self.find_setting(setting)
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        
        messages.append({"role": "user", "content": user_prompt})
        print(messages)
        if self.model in self.ollama_local_models:
            response = self.client.chat(
                model=self.model,
                messages=messages,
            )
            # Track token usage
            self.token_tracker.track(messages, response.message.content)
            return response.message.content
        elif self.model in self.together_ai_models:
            # Use dictionary-based settings
            if setting:
                # Create a dictionary with default values
                params = {
                    'messages': messages,
                    'model': self.model,
                    'max_tokens': 500,
                }
                
                # Add available parameters from the setting dictionary
                if 'temperature' in setting:
                    params['temperature'] = setting['temperature']
                if 'min_p' in setting:
                    params['min_p'] = setting['min_p']
                if 'top_p' in setting:
                    params['top_p'] = setting['top_p']
                if 'repetition_penalty' in setting:
                    params['repetition_penalty'] = setting['repetition_penalty']
                if 'seed' in setting:
                    params['seed'] = setting['seed']
                
                # Create the completion with all available parameters
                response = self.client.chat.completions.create(**params)
            else:
                # Default case with no settings
                response = self.client.chat.completions.create(
                    messages=messages,
                    model=self.model,
                    max_tokens=700,
                )
            output = response.choices[0].message.content
            response = Response(
                output=output,
                model=self.model,
                usage=Usage(input_tokens=self.calculate_tokens(json.dumps(messages)),
                            output_tokens=self.calculate_tokens(output))
            )

            return response
        elif self.model in self.anthropic_models:
            response = self.client.messages.create(
                max_tokens=700,
                temperature=0.7,
                model=self.model,
                messages=messages,
            )

            final_response = next(
                (block.text for block in response.content if hasattr(block, "text")),
                None,
            )
            # Track token usage
            self.token_tracker.track(messages, final_response)
            return final_response
        elif self.model in self.openai_models:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                max_tokens=700,
            )
            # Track token usage
            self.token_tracker.track(messages, response.choices[0].message.content)
            return response.choices[0].message.content
        elif self.model in self.gemini_models:
            response = self.client.models.generate_content(
                model=self.model,
                config=types.GenerateContentConfig(
                    system_instruction=system_prompt,
                    temperature=0.7,
                    max_output_tokens=700,
                    safety_settings=[
                        types.SafetySetting(
                            category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                            threshold=types.HarmBlockThreshold.BLOCK_NONE,
                        ),
                        types.SafetySetting(
                            category=types.HarmCategory.HARM_CATEGORY_HARASSMENT,
                            threshold=types.HarmBlockThreshold.BLOCK_NONE,
                        ),
                        types.SafetySetting(
                            category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
                            threshold=types.HarmBlockThreshold.BLOCK_NONE,
                        ),
                        types.SafetySetting(
                            category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
                            threshold=types.HarmBlockThreshold.BLOCK_NONE,
                        )
                    ]
                ),
                contents=user_prompt,
            )
            # Track token usage
            self.token_tracker.track(messages, response.text)
            return response.text
        elif self.model in self.chutes_models:
            async def invoke_chute(messages):
                params =     {
                    "model": self.model,  #"deepseek-ai/DeepSeek-R1"
                    "messages": messages,
                    "stream": False,
                    "max_tokens": 10000,
                    #"temperature": 0.5
                }

                # Add available parameters from the setting dictionary
                if 'temperature' in setting:
                    params['temperature'] = setting['temperature']
                if 'min_p' in setting:
                    params['min_p'] = setting['min_p']
                if 'top_p' in setting:
                    params['top_p'] = setting['top_p']
                if 'repetition_penalty' in setting:
                    params['repetition_penalty'] = setting['repetition_penalty']
                if 'seed' in setting:
                    params['seed'] = setting['seed']
                
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        "https://llm.chutes.ai/v1/chat/completions", 
                        headers=self.client,
                        json=params
                    ) as response:
                        # Since stream is set to False, we can just get the full response at once
                        response_json = await response.json()
                        #print(json.dumps(response_json, indent=2))  # Pretty print the response
                        # Track token usage
                        self.token_tracker.track(messages, response_json['choices'][0]['message']['content'])
                        return response_json['choices'][0]['message']['content']

            return asyncio.run(invoke_chute(messages))

            

    def JSON_llm(self, user_prompt : str, schema : BaseModel, system_prompt : Optional[str] = None, setting: int = None):
        """
        Run inference with the language model and parse the response as structured JSON.
        
        This method is specifically designed for cases where a structured response
        conforming to a Pydantic schema is required.
        
        Args:
            user_prompt (str): The user prompt to send to the model
            schema (BaseModel): Pydantic model defining the expected response structure
            system_prompt (Optional[str]): System prompt to override the default
            setting (int, optional): Setting number for model configurations
            
        Returns:
            Dict or BaseModel instance: The structured response from the model
            
        Raises:
            ValueError: If the model response cannot be parsed according to the schema
            
        Token usage is tracked and can be retrieved using the count_tokens method
        """        
        try:
            setting = self.find_setting(setting)
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            
            messages.append({"role": "user", "content": user_prompt})
            if self.model in self.ollama_local_models:
                extract = self.client.chat(
                    model=self.model,
                    messages=messages,
                    format=schema.model_json_schema(),
                )

                response = json.loads(extract.message.content)
                # Track token usage
                self.token_tracker.track(messages, extract.message.content)
                return response
            elif self.model in self.together_ai_models:
                # Use dictionary-based settings
                if setting:
                    # Create a dictionary with default values
                    params = {
                        'messages': messages,
                        'model': self.model,
                        'max_tokens': 5000, # 500 for preliminary experiment, 5000 for main experiment
                        'response_format': {
                            "type": "json_object",
                            "schema": schema.model_json_schema(),
                        }
                    }
                    
                    # Add available parameters from the setting dictionary
                    if 'temperature' in setting:
                        params['temperature'] = setting['temperature']
                    if 'min_p' in setting:
                        params['min_p'] = setting['min_p']
                    if 'top_p' in setting:
                        params['top_p'] = setting['top_p']
                    if 'repetition_penalty' in setting:
                        params['repetition_penalty'] = setting['repetition_penalty']
                    if 'seed' in setting:
                        params['seed'] = setting['seed']
                    
                    # Create the completion with all available parameters
                    response = self.client.chat.completions.create(**params)
                else:
                    # Default case with no settings
                    response = self.client.chat.completions.create(
                        messages=messages,
                        model=self.model,
                        response_format={
                            "type": "json_object",
                            "schema": schema.model_json_schema(),
                        },
                        max_tokens=500,
                    )
                output = response.choices[0].message.content
                response = Response(
                    output=json.loads(output),
                    model=self.model,
                    usage=Usage(input_tokens=self.calculate_tokens(json.dumps(messages)),
                                output_tokens=self.calculate_tokens(output))
                )

                return response
            elif self.model in self.openai_models:
                response = self.client.responses.parse(
                    input=messages,
                    model=self.model,
                    text_format=schema,
                )
                # Track token usage
                self.token_tracker.track(messages, response.output_text)
                return response.output_text
            elif self.model in self.chutes_models:
                async def invoke_chute(messages):
                    # Add schema instructions to the system message
                    schema_instructions = f"You must respond with a valid JSON object that follows this schema: {json.dumps(schema.model_json_schema())}"
                    
                    # Add or update system message with schema instructions
                    has_system = False
                    for msg in messages:
                        if msg["role"] == "system":
                            msg["content"] = msg["content"] + "\n\n" + schema_instructions
                            has_system = True
                            break
                    
                    if not has_system:
                        messages.insert(0, {"role": "system", "content": schema_instructions})
                    
                    # Add explicit instruction to user message
                    for i, msg in enumerate(messages):
                        if msg["role"] == "user":
                            messages[i]["content"] = msg["content"] + "\n\nYou must respond with a valid JSON object only, no other text."
                    
                    body = {
                        "model": self.model,
                        "messages": messages,
                        "stream": False,
                        "max_tokens": 10000,
                        "temperature": 0.5
                    }
                    
                    async with aiohttp.ClientSession() as session:
                        async with session.post(
                            "https://llm.chutes.ai/v1/chat/completions", 
                            headers=self.client,
                            json=body
                        ) as response:
                            response_json = await response.json()
                            json_content = response_json['choices'][0]['message']['content']
                            # Track token usage
                            self.token_tracker.track(messages, json_content)
                            
                            # Clean up the response to ensure it's valid JSON
                            # Remove any markdown code block markers
                            json_content = json_content.replace("```json", "").replace("```", "").strip()
                            
                            # Try to find JSON object in the response
                            import re
                            json_match = re.search(r'(\{.*\})', json_content, re.DOTALL)
                            if json_match:
                                json_content = json_match.group(1)
                            
                            # Parse the JSON string to ensure it's valid and matches schema
                            try:
                                parsed_json = json.loads(json_content)
                                return parsed_json
                            except json.JSONDecodeError:
                                # If parsing fails, try to fix common JSON issues
                                try:
                                    # Replace single quotes with double quotes
                                    fixed_content = json_content.replace("'", "\"")
                                    parsed_json = json.loads(fixed_content)
                                    return parsed_json
                                except:
                                    raise ValueError(f"Model response is not valid JSON: {json_content}")
                
                return asyncio.run(invoke_chute(messages))
            
            
        except ValidationError as e:
            raise ValueError(f"Schema validation failed: {str(e)}")
        except Exception as e:
            # Catch-all for any other exceptions
            raise ValueError(f"Error in JSON_llm: {str(e)}")
        
    def get_tiktoken_stats(self):
        """
        Get the token usage statistics.
        
        Returns:
            Dict: A dictionary containing total input and output tokens
        """
        return self.token_tracker.get_stats()
    
    def calculate_tokens(self, string: str) -> int:
        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(string))
    
# Global state to track cooldowns across all decorator instances
_key_cooldowns: Dict[str, float] = {}
# Add global retry counter for exponential backoff
_retry_counts: Dict[str, int] = {}

def rotate_api_keys(cooldown_period=10, debug=False):
    """
    Simple decorator for API key rotation and error handling.
    Works with both sync and async functions.
    No retry limit - will keep trying until success.
    """
    def decorator(func):
        is_async = inspect.iscoroutinefunction(func)
        
        # Get API keys from environment
        api_keys = []
        # Check if we're using Together API keys or OpenAI API keys based on the function's class
        if 'TogetherModel' in func.__qualname__ or 'PydanticModel' in func.__qualname__:
            for key in os.environ:
                if "TOGETHER_API_KEY" in key and os.environ[key]:
                    api_keys.append(os.environ[key])
        elif 'ChutesModel' in func.__qualname__:
            for key in os.environ:
                if "CHUTES_API_KEY" in key and os.environ[key]:
                    api_keys.append(os.environ[key])
        elif 'DeepinfraModel' in func.__qualname__:
            for key in os.environ:
                if "DEEPINFRA_API_KEY" in key and os.environ[key]:
                    api_keys.append(os.environ[key])
        elif 'OpenaiAPIModel' in func.__qualname__:
            for key in os.environ:
                if "OPENAI_API_KEY" in key and os.environ[key]:
                    api_keys.append(os.environ[key])
        
        # eliminate repeated keys
        api_keys = list(set(api_keys))
        if not api_keys:
            raise ValueError(f"No API keys found in environment variables for {func.__qualname__}")
        
        def log(message):
            if debug:
                print(f"[API Key Rotation] {message}")

        def create_client(api_key: str, func_qualname: str, is_async: bool):
            """Create appropriate client based on model type"""
            if 'TogetherModel' in func_qualname:
                if is_async:
                    return AsyncTogether(api_key=api_key,
                                        timeout=55,)
                else:
                    return Together(api_key=api_key,
                                    timeout=55,)
            elif 'DeepinfraModel' in func_qualname:
                pass
                # if is_async:
                #     return AsyncOpenAI(api_key=api_key, base_url="https://api.deepinfra.com/v1/openai")
                # else:
                #     return OpenAI(api_key=api_key, base_url="https://api.deepinfra.com/v1/openai")
            elif 'ChutesModel' in func_qualname:
                return {
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                }
            elif 'PydanticModel' in func_qualname:
                # For PydanticModel, we'll pass the api_key and let the method handle client creation
                return api_key
            elif 'OpenaiAPIModel' in func_qualname:
                if is_async:
                    return AsyncOpenAI(api_key=api_key)
                else:
                    return OpenAI(api_key=api_key)
            else:
                raise ValueError(f"Unknown model type: {func_qualname}")
        
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            while True:  # No retry limit, continue until success
                # Get available keys not on cooldown
                now = time.time()
                available_keys = [key for key in api_keys 
                                if key not in _key_cooldowns 
                                or (now - _key_cooldowns[key]) >= cooldown_period]
                
                task = kwargs.get('user_prompt', 'No task provided')

                # If no keys available, wait for the first one
                if not available_keys:
                    earliest = min([_key_cooldowns[key] for key in _key_cooldowns]) + cooldown_period
                    wait_time = max(0, earliest - now)
                    log(f"[Task: {task[:20]}] All keys on cooldown. Waiting {wait_time:.1f}s")
                    await asyncio.sleep(wait_time)
                    continue
                
                # Try request with a random available key
                key = random.choice(available_keys)
                try:
                    client = create_client(key, func.__qualname__, is_async=True)
                    kwargs['client'] = client
                    kwargs['api_key'] = key

                    # Reset retry count on successful attempt
                    key_id = id(key)
                    if key_id in _retry_counts:
                        del _retry_counts[key_id]

                    return await func(*args, **kwargs)
                except Exception as e:
                    error_msg = str(e).lower()

                    # Add retry counter for exponential backoff
                    key_id = id(key)
                    _retry_counts[key_id] = _retry_counts.get(key_id, 0) + 1

                    if "rate limit" in error_msg or "429" in error_msg:
                        _key_cooldowns[key] = time.time()
                        log(f"Rate limit for key. Cooling down for {cooldown_period}s")
                    elif "500" in error_msg or "502" in error_msg or "503" in error_msg or "internal server error" in error_msg:
                        backoff_time = min(30, 2 ** _retry_counts[key_id])
                        log(f"Server error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        await asyncio.sleep(backoff_time)
                    elif "error communicating with together" in error_msg:
                        backoff_time = min(20, 2 ** _retry_counts[key_id])
                        log(f"API communication error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        await asyncio.sleep(backoff_time)
                    elif "connection" in error_msg or "timeout" in error_msg or "request timed out" in error_msg or "read timeout" in error_msg:
                        backoff_time = min(25, 3 ** _retry_counts[key_id])
                        log(f"Connection/Timeout error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        await asyncio.sleep(backoff_time)
                    elif e.__class__.__name__ == "JSONDecodeError":
                        log(f"{e.__class__.__name__} error detected. Retry immediately.")
                    elif "parsed output is not a valid json object" in error_msg or "failed to parse valid json object" in error_msg or "parsed output <differential diagnosis> does not contain required 'items' field" in error_msg:
                        log(f"Invalid JSON output detected. Retry immediately.")
                    elif "401" in error_msg or "invalid api key" in error_msg:
                        log(f"Invalid API key detected. Removing key from rotation.")
                        if key in api_keys:
                            api_keys.remove(key)
                        else:
                            log(f"Key already removed from rotation.")
                    else:
                        log(f"Unexpected error: {e}")
                        raise  # Only raise for truly unexpected errors
        
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            while True:  # No retry limit, continue until success
                # Get available keys not on cooldown
                now = time.time()
                available_keys = [key for key in api_keys 
                                if key not in _key_cooldowns 
                                or (now - _key_cooldowns[key]) >= cooldown_period]
                
                task = kwargs.get('user_prompt', 'No task provided')

                # If no keys available, wait for the first one
                if not available_keys:
                    earliest = min([_key_cooldowns[key] for key in _key_cooldowns]) + cooldown_period
                    wait_time = max(0, earliest - now)
                    log(f"[Task: {task[:20]}] All keys on cooldown. Waiting {wait_time:.1f}s")
                    time.sleep(wait_time)
                    continue
                
                # Try request with a random available key
                key = random.choice(available_keys)
                try:
                    client = create_client(key, func.__qualname__, is_async=False)
                    kwargs['client'] = client
                    kwargs['api_key'] = key

                     # Reset retry count on successful attempt
                    key_id = id(key)
                    if key_id in _retry_counts:
                        del _retry_counts[key_id]
                    
                    return func(*args, **kwargs)
                except Exception as e:
                    error_msg = str(e).lower()

                     # Add retry counter for exponential backoff
                    key_id = id(key)
                    _retry_counts[key_id] = _retry_counts.get(key_id, 0) + 1

                    if "rate limit" in error_msg or "429" in error_msg:
                        _key_cooldowns[key] = time.time()
                        log(f"Rate limit for key. Cooling down for {cooldown_period}s")
                    elif "500" in error_msg or "502" in error_msg or "503" in error_msg or "internal server error" in error_msg:
                        backoff_time = min(30, 2 ** _retry_counts[key_id])
                        log(f"Server error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        time.sleep(backoff_time)
                    elif "error communicating with together" in error_msg:
                        backoff_time = min(20, 2 ** _retry_counts[key_id])
                        log(f"API communication error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        time.sleep(backoff_time)
                    elif "connection" in error_msg or "timeout" in error_msg or "request timed out" in error_msg or "read timeout" in error_msg:
                        backoff_time = min(25, 3 ** _retry_counts[key_id])
                        log(f"Connection/Timeout error, backing off for {backoff_time}s (attempt {_retry_counts[key_id]}): {e}")
                        time.sleep(backoff_time)
                    elif e.__class__.__name__ == "JSONDecodeError":
                        log(f"{e.__class__.__name__} error detected. Retry immediately.")
                    elif "parsed output is not a valid json object" in error_msg or "failed to parse valid json object" in error_msg or "parsed output <differential diagnosis> does not contain required 'items' field" in error_msg:
                        log(f"Invalid JSON output detected. Retry immediately.")
                    elif "401" in error_msg or "invalid api key" in error_msg:
                        log(f"Invalid API key detected. Removing key from rotation.")
                        if key in api_keys:
                            api_keys.remove(key)
                        else:
                            log(f"Key already removed from rotation.")
                    else:
                        log(f"Unexpected error: {e}")
                        raise  # Only raise for truly unexpected errors
        
        return async_wrapper if is_async else sync_wrapper
    return decorator

def log_info(message, logger_name="message_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")

class TogetherModel():
    """
    Model class for Together.ai models.
    
    This class provides methods to run inference with Together.ai models and
    handle token usage tracking.
    
    Attributes:
        model (str): The identifier for the model to use
        system_prompt (Optional[str]): Default system prompt to use with the model
        client: The initialized Together client
        token_tracker: TokenTracker instance for tracking token usage
    """
    
    def __init__(self, model: str, system_prompt: Optional[str] = None, setting: dict = None):
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.setting = setting
    
    def _build_messages(self, user_prompt: Union[str, List[str]]) -> Union[List[dict], List[List[dict]]]:
        def generate_messages(user_prompt: str) -> List[dict]:
            messages = []
            if self.system_prompt:
                messages.append({"role": "system", "content": self.system_prompt})
            messages.append({"role": "user", "content": user_prompt})
            return messages
        
        if isinstance(user_prompt, list):
            # Return a list of message lists (one for each prompt)
            return [generate_messages(prompt) for prompt in user_prompt]
        else:
            return generate_messages(user_prompt)
    
    def _build_params(self, messages: List[dict], output_type: Any, **kwargs) -> dict:
        base_params = {
            "messages": messages,
            "model": self.__model_name__,
            "max_tokens": 2000,
            **self.setting,
            **kwargs
        }
        if output_type != str:
            base_params["response_format"] = {
                "type": "json_object",
                "schema": output_type.model_json_schema(),
            }
        
        return base_params
    
    def _create_response(self, output: Any, output_type: Any) -> Response:
        content = output.choices[0].message.content
        parsed_output = content if output_type == str else json.loads(content)
        
        return Response(
            output=parsed_output,
            model=self.__model_name__,
            usage=Usage(
                input_tokens=output.usage.prompt_tokens,
                output_tokens=output.usage.completion_tokens
            )
        )
    
    @rotate_api_keys(debug=True)
    def run(self, client: Together, api_key: str, user_prompt: str, output_type: Any = str, **kwargs) -> Response:
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        output = client.chat.completions.create(**params)
        log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
        return self._create_response(output, output_type)

    @rotate_api_keys(debug=True)
    async def run_async(self, client: AsyncTogether, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        output = await client.chat.completions.create(**params)
        log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
        return self._create_response(output, output_type)
    
class PydanticModel():
    def __init__(self,model: str, system_prompt: Optional[str] = None, setting: Dict = None,):
        self.base_url = "https://api.together.xyz/v1"
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.setting = setting

    def _build_agent(self, api_key: str, output_type: Any, params: dict) -> Agent:
        # Create the model instance
        self.model = OpenAIModel(
            self.__model_name__,
            provider=OpenAIProvider(base_url=self.base_url, api_key=api_key)
        )

        return Agent(
            self.model,
            system_prompt=self.system_prompt,
            model_settings=params,
            output_type=output_type,
        )   
    
    def _build_params(self, **kwargs) -> dict:
        base_params = {
            "max_tokens": 2000,
            **self.setting,
            **kwargs
        }
        
        return base_params
    
    def _create_response(self, output: Any, output_type: Any) -> Response:
        content = output.choices[0].message.content
        parsed_output = content if output_type == str else json.loads(content)
        
        return Response(
            output=parsed_output,
            model=self.__model_name__,
            usage=Usage(
                input_tokens=output.usage.prompt_tokens,
                output_tokens=output.usage.completion_tokens
            )
        )
    
    @rotate_api_keys(debug=True)
    def run(self, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        params = self._build_params(**kwargs)
        agent = self._build_agent(api_key=api_key, output_type=output_type, params=params)

        output = agent.run_sync(user_prompt=user_prompt)
        return output

    @rotate_api_keys(debug=True)
    async def run_async(self, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        params = self._build_params(**kwargs)
        agent = self._build_agent(api_key=api_key, output_type=output_type, params=params)

        output = await agent.run(user_prompt=user_prompt)
        return output

class DeepinfraModel():
    def __init__(self, model: str, system_prompt: Optional[str] = None, setting: dict = None):
        self.base_url = "https://api.deepinfra.com/v1/inference/"
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.setting = setting
    
    def _build_messages(self, user_prompt: Union[str, List[str]], output_type) -> Union[List[dict], List[List[dict]]]:
        def generate_messages(user_prompt: str) -> List[dict]:
            messages = ""
            deepseek_think_restriction = f"""You must use <think> tags.
BUT: Keep your thinking under 50 words.
Format:
<think>[Brief thinking here, max 50 words]</think>
[JSON or TEXT answer here]"""
            end_length_restriction = "\nRemember to STRICTLY keep your thinking short under 50 words."
            if output_type != str:
                pydantic_schema = output_type.model_json_schema()

                schema_instructions = f"""You must respond with a valid JSON object that STRCTLY follows this schema: {json.dumps(pydantic_schema)} \n\nExample output:
{{
"items": [
    {{"disease": "disease", "confidence": 1.0}},
    {{"disease": "disease", "confidence": 1.0}}
    ...
]
}}""" if output_type.__name__ == "DifferentialDiagnosis" and self.__model_name__ == "deepseek-ai/DeepSeek-R1-0528" else f"""You must respond with a valid JSON object that STRCTLY follows this schema: {json.dumps(pydantic_schema)}"""                
                system_prompt = copy.deepcopy(self.system_prompt) if self.system_prompt else ""
                system_prompt += f"\n\n{schema_instructions}"
                user_prompt += "\n\nYou must respond with a valid JSON object only, no other text."

            if self.system_prompt:
                if self.__model_name__ == "deepseek-ai/DeepSeek-R1-0528":
                    messages += f"<｜begin▁of▁sentence｜>{deepseek_think_restriction + (self.system_prompt if output_type == str else system_prompt)}<｜User｜>{user_prompt + end_length_restriction}<｜Assistant｜>"
                else:
                    messages += f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{self.system_prompt if output_type == str else system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            else:
                if self.__model_name__ == "deepseek-ai/DeepSeek-R1-0528":
                    messages += f"<｜begin▁of▁sentence｜><｜User｜>{deepseek_think_restriction + user_prompt + end_length_restriction}<｜Assistant｜>"
                else:
                    messages += f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            return messages

        if isinstance(user_prompt, list):
            return [generate_messages(prompt) for prompt in user_prompt]
        else:
            return generate_messages(user_prompt)
    
    def _build_params(self, messages: str, output_type: Any, **kwargs) -> dict:
        base_params = {
            "input": messages,
            "max_new_tokens": 2500 if self.__model_name__ == "deepseek-ai/DeepSeek-R1-0528" else 2000,
            **self.setting,
        }
        if output_type != str:
            base_params["response_format"] = {
                "type": "json_object",
            }
        
        return base_params
    
    def _create_response(self, output: Any, output_type: Any) -> Response:
        content = output["results"][0]["generated_text"]
        if self.__model_name__ == "deepseek-ai/DeepSeek-R1-0528":
            content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
        if output_type != str:
            parsed_output = json.loads(content)
            if not isinstance(parsed_output, dict):
                raise ValueError("Parsed output is not a valid JSON object (expected dict)")
            if output_type.__name__ == "DifferentialDiagnosis" and parsed_output.get("items", None) is None:
                raise ValueError("Parsed output <Differential Diagnosis> does not contain required 'items' field") 
        else:
            parsed_output = content  

        return Response(
            output=parsed_output,
            model=self.__model_name__,
            usage=Usage(
                input_tokens=output['inference_status']['tokens_input'],
                output_tokens=output['inference_status']['tokens_generated']
            ),
            body=output
        )
    
    @rotate_api_keys(debug=True)
    def run(self, client: OpenAI, api_key: str, user_prompt: str, output_type: Any = str, **kwargs) -> Response:
        messages = self._build_messages(user_prompt=user_prompt, output_type=output_type)
        params = self._build_params(messages, output_type, **kwargs)
        # log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        # output = client.chat.completions.create(**params)
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        try:
            response = requests.post(
                self.base_url + self.__model_name__,
                headers=headers,
                json=params,
            )
            response.raise_for_status()  # Raise exception for HTTP errors
            output = response.json()
            log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
            return self._create_response(output, output_type)
        except requests.exceptions.RequestException as e:
            raise ValueError(f"error communicating with deepinfra: {e}")
        # log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
        # return self._create_response(output, output_type)

    @rotate_api_keys(debug=True)
    async def run_async(self, client: AsyncOpenAI, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        messages = self._build_messages(user_prompt=user_prompt, output_type=output_type)
        params = self._build_params(messages, output_type, **kwargs)
        # log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        # output = await client.chat.completions.create(**params)
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(
                    self.base_url + self.__model_name__,
                    headers=headers,
                    json=params,
                ) as response:
                    if response.status >= 400:
                        error_text = await response.text()
                        raise ValueError(f"deepinfra error ({response.status}): {error_text}")

                    output = await response.json()
                    log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
                    return self._create_response(output, output_type)
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                raise ValueError(f"error communicating with deepinfra: {e}")
            
class OpenaiAPIModel():
    def __init__(self, model: str, system_prompt: Optional[str] = None, setting: dict = None):
        self.base_url = "https://api.openai.com/v1/chat/completions"
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.setting = setting
    
    def _build_messages(self, user_prompt: Union[str, List[str]]) -> Union[List[dict], List[List[dict]]]:
        def generate_messages(user_prompt: str) -> List[dict]:
            messages = []
            if self.system_prompt:
                messages.append({"role": "system", "content": self.system_prompt})
            messages.append({"role": "user", "content": user_prompt})
            return messages
        
        if isinstance(user_prompt, list):
            return [generate_messages(prompt) for prompt in user_prompt]
        else:
            return generate_messages(user_prompt)
    
    def _build_params(self, messages: str, output_type: Any, **kwargs) -> dict:
        base_params = {
            "messages": messages,
            "model": self.__model_name__,
            "max_completion_tokens": 2000,
            **self.setting,
            **kwargs
        }
        if output_type != str:
            base_params["response_format"] = output_type
        
        return base_params
    
    def _create_response(self, output: Any, output_type: Any) -> Response:
        content = output.choices[0].message.content
        parsed_output = content if output_type == str else json.loads(content)
        
        return Response( 
            output=parsed_output,
            model=self.__model_name__,
            usage=Usage(
                input_tokens=output.usage.prompt_tokens,
                output_tokens=output.usage.completion_tokens
            ),
            body=output
        )
    
    @rotate_api_keys(debug=True)
    def run(self, client: OpenAI, api_key: str, user_prompt: str, output_type: Any = str, **kwargs) -> Response:
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        if output_type != str:
            output = client.chat.completions.parse(**params)
        else:
            output = client.chat.completions.create(**params)
        log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
        return self._create_response(output, output_type)

    @rotate_api_keys(debug=True)
    async def run_async(self, client: AsyncOpenAI, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        log_info(f"[{self.__model_name__}][INPUT]: {messages}")

        if output_type != str:
            output = await client.chat.completions.parse(**params)
        else:
            output = await client.chat.completions.create(**params)
        log_info(f"[{self.__model_name__}][OUTPUT]: {output}")
        return self._create_response(output, output_type)
    
class ChutesModel():
    def __init__(self, model: str, system_prompt: Optional[str] = None, setting: dict = None):
        self.__model_name__ = model
        self.system_prompt = system_prompt
        self.setting = setting
    
    def _build_messages(self, user_prompt: Union[str, List[str]]) -> Union[List[dict], List[List[dict]]]:
        def generate_messages(user_prompt: str) -> List[dict]:
            messages = []
            if self.system_prompt:
                messages.append({"role": "system", "content": self.system_prompt})
            messages.append({"role": "user", "content": user_prompt})
            return messages
        
        if isinstance(user_prompt, list):
            # Return a list of message lists (one for each prompt)
            return [generate_messages(prompt) for prompt in user_prompt]
        else:
            return generate_messages(user_prompt)
    
    def _build_params(self, messages: List[dict], output_type: Any, **kwargs) -> dict:
        base_params = {
            "messages": messages,
            "model": self.__model_name__,
            "max_tokens": 2000,
            "stream": False,
            **self.setting,
            **kwargs
        }
        if output_type != str:
            # Format the schema for Chutes.ai
            schema_name = output_type.__name__.lower()
            pydantic_schema = output_type.model_json_schema()
            
            # Add schema instructions to system message instead
            for i, msg in enumerate(base_params["messages"]):
                if msg["role"] == "system":
                    schema_instructions = f"You must respond with a valid JSON object that follows this schema: {json.dumps(pydantic_schema)}"
                    base_params["messages"][i]["content"] += f"\n\n{schema_instructions}"
                    break
            else:
                # If no system message, add one
                schema_instructions = f"You must respond with a valid JSON object that follows this schema: {json.dumps(pydantic_schema)}"
                base_params["messages"].insert(0, {"role": "system", "content": schema_instructions})
            
            # Add instruction to user message
            for i, msg in enumerate(base_params["messages"]):
                if msg["role"] == "user":
                    base_params["messages"][i]["content"] += "\n\nYou must respond with a valid JSON object only, no other text."
                    break
        
        return base_params
    
    def _create_response(self, output: Any, output_type: Any) -> Response:
        content = output.get('choices', [{}])[0].get('message', {}).get('content', '')
        
        if output_type != str:
            try:
                # Clean up the response to ensure it's valid JSON
                json_content = content.replace("```json", "").replace("```", "").strip()
                
                # Try to find JSON object in the response if it's embedded in other text
                import re
                json_match = re.search(r'(\{.*\})', json_content, re.DOTALL)
                if json_match:
                    json_content = json_match.group(1)
                
                # Parse the JSON string to ensure it's valid
                try:
                    parsed_json = json.loads(json_content)
                except json.JSONDecodeError:
                    # If parsing fails, try to fix common JSON issues
                    fixed_content = json_content.replace("'", "\"")
                    parsed_json = json.loads(fixed_content)
                
                parsed_output = parsed_json
            except Exception as e:
                raise ValueError(f"Failed to parse JSON response: {e}\nResponse: {content}")
        else:
            parsed_output = content
        
        # Extract token usage if available
        usage = output.get('usage', {})
        
        return Response(
            output=parsed_output,
            model=self.__model_name__,
            usage=Usage(
                input_tokens=usage.get('prompt_tokens', 0),
                output_tokens=usage.get('completion_tokens', 0)
            )
        )
    
    @rotate_api_keys(debug=True)
    def run(self, api_key: str, user_prompt: str, output_type: Any = str, **kwargs) -> Response:
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        try:
            response = requests.post(
                "https://llm.chutes.ai/v1/chat/completions",
                headers=headers,
                json=params,
                timeout=60
            )
            response.raise_for_status()  # Raise exception for HTTP errors
            output = response.json()
            return self._create_response(output, output_type)
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Error communicating with Chutes.ai: {e}")

    @rotate_api_keys(debug=True)
    async def run_async(self, api_key: str, user_prompt: str, output_type: Any = str, **kwargs):
        messages = self._build_messages(user_prompt)
        params = self._build_params(messages, output_type, **kwargs)
        
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(
                    "https://llm.chutes.ai/v1/chat/completions",
                    headers=headers,
                    json=params,
                    timeout=aiohttp.ClientTimeout(total=60)
                ) as response:
                    if response.status >= 400:
                        error_text = await response.text()
                        raise ValueError(f"Chutes.ai API error ({response.status}): {error_text}")
                    
                    output = await response.json()
                    return self._create_response(output, output_type)
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                raise ValueError(f"Error communicating with Chutes.ai: {e}")
    
class Model:
    """Enhanced factory class with model registry pattern"""
    
    _model_registry = {
        "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": TogetherModel,
        "meta-llama/Llama-3.3-70B-Instruct-Turbo": DeepinfraModel,
        "deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free": TogetherModel,
        "Qwen/QwQ-32B": TogetherModel,
        "Qwen/Qwen2.5-72B-Instruct-Turbo": TogetherModel,
        "meta-llama/Llama-Vision-Free": TogetherModel,
        "meta-llama/Meta-Llama-3.1-8B-Instruct": DeepinfraModel,
        # "meta-llama/llama-3.3-8b-instruct:free": OpenaiModel,
        "chutesai/Mistral-Small-3.1-24B-Instruct-2503": ChutesModel,
        "deepseek-ai/DeepSeek-R1": ChutesModel,
        "deepseek-ai/DeepSeek-V3-0324": ChutesModel,
        "gpt-5.1-2025-11-13": OpenaiAPIModel,
        "deepseek-ai/DeepSeek-R1-0528": DeepinfraModel,
    }
    
    @classmethod
    def create_model(
        cls,
        model: str,
        system_prompt: Optional[str] = None,
        setting: Optional[dict] = None,
    ):
        if model not in cls._model_registry:
            raise ValueError(f"Unsupported model: {model}")
        
        return cls._model_registry[model](
            model=model,
            system_prompt=system_prompt,
            setting=setting,
        )

class TokenTracker:
    def __init__(self):
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        
    def track(self, prompt: List[dict], completion: str):
        # Using tiktoken for counting
        enc = tiktoken.get_encoding("cl100k_base")
        for msg in prompt:
            message = msg.get("content", "")
            self.total_input_tokens += len(enc.encode(message))
        self.total_output_tokens += len(enc.encode(completion))
        
    def get_stats(self):
        return {
            "total_tokens": self.total_input_tokens + self.total_output_tokens,
            "input_tokens": self.total_input_tokens,
            "output_tokens": self.total_output_tokens
        }
    
    def combine_stats(self, stats_list: List[Dict]) -> Dict:
        """
        Combine multiple token statistics dictionaries into a single dictionary.
        
        Args:
            stats_list (List[Dict]): List of dictionaries containing token statistics
            
        Returns:
            Dict: Combined token statistics
        """
        combined_stats = {
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0
        }
        
        for stats in stats_list:
            combined_stats["total_tokens"] += stats.get("total_tokens", 0)
            combined_stats["input_tokens"] += stats.get("input_tokens", 0)
            combined_stats["output_tokens"] += stats.get("output_tokens", 0)
        
        return combined_stats
    
class MedqaDataLoader:
    def __init__(self, seed: int = 123):
        self.data_path = os.path.join(DIR_PATH, 'dataset', 'medqa.jsonl')
        with open(self.data_path, 'r', encoding='utf-8') as f:
            self.data = [json.loads(line.strip()) for line in f if line.strip()]

        random.seed(seed)  # Set a seed for reproducibility
        random.shuffle(self.data)

    def get_original_data(self) -> List[dict]:
        return self.data
    
    def get_evaluation_set(self, num_cases: int = 170) -> List[PatientCase]:
        dataset = []
        for line in self.data[:num_cases]:
            index = line.get('id', None)
            facts = line.get('facts', None)
            options = line.get('options', None)
            answer = line.get('answer', None)
            patient_case = PatientCase.create(patient_info=facts, partial_patient_info=facts[0], options=options, ground_truth=answer, metadata={"id": index})
            dataset.append(patient_case)
        return dataset
    
class MediqDataLoader:
    def __init__(self, seed: int = 123):
        self.data_path = os.path.join(DIR_PATH, 'dataset', 'mediq.jsonl')
        with open(self.data_path, 'r', encoding='utf-8') as f:
            self.data = [json.loads(line.strip()) for line in f if line.strip()]

        random.seed(seed)  # Set a seed for reproducibility
        random.shuffle(self.data)

    def get_original_data(self) -> List[dict]:
        return self.data
    
    def get_evaluation_set(self, num_cases: int = 129) -> List[PatientCase]:
        dataset = []
        for line in self.data[:num_cases]:
            index = line.get('id', None)
            context = line.get('context', None)
            facts = line.get('facts', None)
            options = line.get('options', None)
            answer = line.get('answer', None)
            patient_case = PatientCase.create(patient_info=facts, partial_patient_info=context[0], options=options, ground_truth=answer, metadata={"id": index})
            dataset.append(patient_case)
        return dataset

class NejmDataLoader:
    def __init__(self, seed: int = 123, data = None):
        self.data_path = os.path.join(DIR_PATH, 'dataset', 'nejm.jsonl')
        if data is None:
            with open(self.data_path, 'r', encoding='utf-8') as f:
                self.data = [json.loads(line.strip()) for line in f if line.strip()]
        else:
            self.data = data

        random.seed(seed)  # Set a seed for reproducibility
        random.shuffle(self.data)

    def get_original_data(self) -> List[dict]:
        return self.data
    
    def get_evaluation_set(self, num_cases: int = 170) -> List[PatientCase]:
        dataset = []
        for line in self.data[:num_cases]:
            index = line.get('id', None)
            facts = line.get('facts', None)
            answer = line.get('Correct_Answer', None)
            patient_case = PatientCase.create(patient_info=facts, partial_patient_info=facts[0], options=[], ground_truth=answer, metadata={"id": index})
            dataset.append(patient_case)
        return dataset
    
class NEJMDataLoader:
    def __init__(self):
        self.data_path = os.path.join(DIR_PATH, 'dataset', 'structural_nejm.jsonl')
        with open(self.data_path, 'r', encoding='utf-8') as f:
            self.data = [json.loads(line.strip()) for line in f if line.strip()]
        
        random.seed(22)  # Set a seed for reproducibility
        random.shuffle(self.data)

    def get_original_data(self) -> List[dict]:
        return self.data
    
    def get_evaluation_set(self, num_cases: int = 170, seed: str = None, num_removal: int = 1) -> List[PatientCase]:
        """
        Use random generator to generate number 0-8, respectively representing 'Demographics', 
        'Past_Medical_History', 'Social_History', 'Review_of_Systems', 'Symptoms', 'Laboratory', 'Imaging',
        'History', 'Past_Medical_History'. For each number, create a PatientCase with missing information.

        """
        missing_cases = []
        categories = ['Past_Medical_History', 'Social_History', 'Review_of_Systems', 
                'Physical_Examination', 'Laboratory', 'Imaging', 'History']
        weights = [0.4, 0.8, 0.9, 0.4, 0.9, 0.9, 0.3]  # Adjust weights as needed

        if seed is not None:
            random.seed(seed)

        def weighted_sample(items, weights, k):
            weighted_pool = random.choices(items, weights=weights, k=10*k)  # Oversample
            unique_items = list(dict.fromkeys(weighted_pool))  # Remove duplicates
            return unique_items[:k]  # Return up to k items

        counts = [0, 0, 0, 0, 0, 0, 0]  # Counts for each category
        for i in range(num_cases):
            missing_info = weighted_sample(categories, weights, len(categories))
            #k = min(6, len(categories))
            #missing_info = random.sample(categories, k=k)

            # remove every empty ""/[]/{}/None from Patient_Case
            def clean_empty(obj):
                if isinstance(obj, dict):
                    return {k: clean_empty(v) for k, v in obj.items() if v not in ("", [], {}, None)}
                elif isinstance(obj, list):
                    return [clean_empty(v) for v in obj if v not in ("", [], {}, None)]
                else:
                    return obj
                
            

            case_to_process = copy.deepcopy(self.data[i]['Patient_Case'])
            case_to_process = clean_empty(case_to_process)

            for category in missing_info:
                if category in case_to_process['Patient_Factor']:
                    if category == 'History':
                        case_to_process['Patient_Factor'].pop('Symptoms', None)
                    case_to_process['Patient_Factor'].pop(category, None)
                    counts[categories.index(category)] += 1
                    print(f"Case {i} removed {category}.")
                    break
                elif category in case_to_process['Test_Results']:
                    case_to_process['Test_Results'].pop(category, None)
                    counts[categories.index(category)] += 1
                    print(f"Case {i} removed {category}.")
                    break
                elif category in case_to_process:
                    case_to_process.pop(category, None)
                    counts[categories.index(category)] += 1
                    print(f"Case {i} removed {category}.")
                    break
                else:
                    '''if len(missing_info) == 1 or (len(missing_info) > 1 and category == missing_info[-1]):
                        print(f"Category {category} not found in Patient_Case for case {i}.")'''

            patient_case = {
                "patient_info": self.data[i]['Patient_Case'],
                "partial_patient_info": clean_empty(case_to_process),
                "ground_truth": self.data[i]['Correct_Answer'],
            }
            
            missing_cases.append(patient_case)
        
        '''for index, item in enumerate(categories):
            print(f"Category {item} was removed {counts[index]} times.")'''
        return missing_cases
        

class MedicalDataLoader:
    """
    DataLoader for medical diagnosis datasets.
    
    This class handles loading patient data from files or directories,
    with support for various data formats, batching, shuffling, and masking
    specific types of information.
    
    Attributes:
        data_path (str): Path to the dataset file or directory
        mask_type (str): Type of information to mask in the patient data
        shuffle (bool): Whether to shuffle the dataset
        data (List): The loaded patient data
    """
    
    def __init__(self, 
                 data_source: str = "nejm",
                 is_sentence: bool = False,
                 year: int = None, 
                 mask_type: str | List[str] = None,
                 shuffle: bool = False):
        """
        Initialize the DataLoader.
        
        Args:
            data_source (str): Data source (e.g., "nejm", "medqa")
            mask_type (str, optional): Type of information to mask (e.g., "age", "main_symptoms")
            shuffle (bool, default=False): Whether to shuffle the dataset
        """
        self.year = year
        self.mask_type = mask_type
        self.shuffle = shuffle
        self.data_source = data_source
        self.is_sentence = is_sentence
        self.data_path = self.get_file_path()
        self.data = []
        self._load_data()

    def get_file_path(self):
        """
        Get the file path for the dataset.

        Returns:
            str: File path of the dataset
        """
        '''if self.mask_type: 
            if self.data_source == "nejm":
                return os.path.join(DIR_PATH, 'dataset', f'nejm_no_{self.mask_type}.jsonl')'''
        
        if self.data_source == "structural_nejm":
            return os.path.join(DIR_PATH, 'dataset', 'structural_nejm.jsonl')
        elif self.data_source == "agentclinic_nejm":
            return os.path.join(DIR_PATH, 'dataset', 'agentclinic_nejm.jsonl')
        elif self.data_source == "structural_medqa":
            return os.path.join(DIR_PATH, 'dataset', 'structural_medqa.jsonl')
        elif self.data_source == "nejm_test":
            return os.path.join(DIR_PATH, 'dataset', 'tmp', f'subset_test_Laboratory.jsonl')
        elif self.data_source == "structural_mediq":
            return os.path.join(DIR_PATH, 'dataset', 'structural_mediq.jsonl')

    def _load_data(self):
        """
        Load data from file or directory.
        
        Supports loading from JSONL files and handles different data formats.
        Sets self.data with the loaded patient records.
        
        Raises:
            FileNotFoundError: If the data path does not exist
            ValueError: If the data format is not supported
        """
        if os.path.isfile(self.data_path):
            self._load_from_file(self.data_path)
        
        if self.year:
            self.data = self.get_subset_by_year(self.year)

        if self.shuffle:
            random.shuffle(self.data)

    
    def _load_from_file(self, filepath: str):
        """Load data from a single file"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                if filepath.endswith('.jsonl'):
                    # Load JSONL file (one JSON object per line)
                    for line in f:
                        line = line.strip()
                        if line:  # Skip empty lines
                            try:
                                record = json.loads(line)
                                self.data.append(record)
                            except json.JSONDecodeError:
                                print(f"Warning: Could not parse line as JSON: {line}")
                else:
                    # Assume each line is a separate record
                    for line in f:
                        try:
                            record = json.loads(line.strip())
                            self.data.append(record)
                        except json.JSONDecodeError:
                            print(f"Warning: Could not parse line as JSON: {line}")

        except Exception as e:
            print(f"Error loading data from {filepath}: {e}")
    
    def __len__(self) -> int:
        """
        Get the number of patient records in the dataset.
        
        Returns:
            int: Number of patient records
        """
        if (self.data_source == 'structural_nejm' or self.data_source == 'structural_mediq') and self.mask_type:
            return len(self.get_subset())
        else:
            return len(self.data)
    
    def dict_to_formatted_text(self, data: dict, indent_level: int = 0) -> str:
        """
        Convert a dictionary to a formatted text representation.
        
        Args:
            data (dict): The dictionary to format
            indent_level (int): Current indentation level (used for recursion)
            
        Returns:
            str: Formatted text representation of the dictionary
        """
        result = []
        indent = "  " * indent_level
        
        if self.data_source == "structural_nejm" or self.data_source == "nejm_test" or self.data_source == "structural_mediq":
            for key, value in data.items():
                # Convert key from snake_case to Title Case
                formatted_key = key.replace('_', ' ').title()
                if value != '' and value != [] and value != {}:
                    if isinstance(value, dict):
                        # Handle nested dictionaries
                        result.append(f"{indent}## {formatted_key}")
                        result.append(self.dict_to_formatted_text(value, indent_level + 1))
                    elif isinstance(value, list):
                        # Handle list values
                        result.append(f"{indent}### {formatted_key}")
                        for item in value:
                            result.append(f"{indent}- {item}")
                    else:
                        # Handle simple values
                        if indent_level == 0:
                            result.append(f"{indent}# {formatted_key}")
                            result.append(f"{indent}{value}")
                        else:
                            result.append(f"{indent}- {formatted_key}: {value}")
            
            
        elif self.data_source == "structural_medqa":
            for key, value in data.items():
                # Convert key from snake_case to Title Case
                formatted_key = key.replace('_', ' ').title()
                
                if isinstance(value, dict):
                    # Handle nested dictionary
                    result.append(f"{indent}## {formatted_key}")
                    result.append(self.dict_to_formatted_text(value, indent_level + 1))
                elif isinstance(value, list):
                    # Handle list values
                    result.append(f"{indent}### {formatted_key}")
                    for item in value:
                        result.append(f"{indent}- {item}")
                else:
                    # Handle simple values
                    if indent_level == 0:
                        result.append(f"{indent}# {formatted_key}")
                        result.append(f"{indent}{value}")
                    else:
                        result.append(f"{indent}- {formatted_key}: {value}")
        
        return "\n".join(result)

    def dict_to_sentence_text(self, data: dict) -> str:
        """
        Convert a dictionary to a natural sentence text representation.
        
        Args:
            data (dict): The dictionary to format
            
        Returns:
            str: Natural sentence text representation of the dictionary
        """
        result = []
        
        if self.data_source == "structural_nejm" or self.data_source == "nejm_test":
            for key, value in data.items():
                # Convert key from snake_case to readable format
                formatted_key = key.replace('_', ' ').title()
                
                if value != '' and value != [] and value != {}:
                    if isinstance(value, dict):
                        # Handle nested dictionaries
                        result.append(f"Regarding {formatted_key.lower()}, {self.dict_to_sentence_text(value)}")
                    elif isinstance(value, list):
                        # Handle list values
                        if len(value) == 1:
                            result.append(f"the {formatted_key.lower()} shows {value[0]}.")
                        else:
                            items = ", ".join(value[:-1]) + " and " + value[-1]
                            result.append(f"The {formatted_key.lower()} shows {items}.")
                    else:
                        # Handle simple values
                        result.append(f"The {formatted_key.lower()} is {value}.")
            
        elif self.data_source == "structural_medqa":
            for key, value in data.items():
                # Convert key from snake_case to readable format
                formatted_key = key.replace('_', ' ').title()
                
                if isinstance(value, dict):
                    # Handle nested dictionary
                    result.append(f"Regarding {formatted_key.lower()}, {self.dict_to_sentence_text(value)}")
                elif isinstance(value, list):
                    # Handle list values
                    if len(value) == 1:
                        result.append(f"the {formatted_key.lower()} shows {value[0]}.")
                    else:
                        items = ", ".join(value[:-1]) + " and " + value[-1]
                        result.append(f"the {formatted_key.lower()} shows {items}.")
                else:
                    # Handle simple values
                    result.append(f"The {formatted_key.lower()} is {value}.")
        
        return " ".join(result)

    def _convert_to_patient_case(self, item: Dict[str, Any]) -> PatientCase:
        """Convert a data item to PatientCase"""
        if self.data_source == "structural_nejm" or self.data_source == "nejm_test" or self.data_source == "structural_mediq":
            masked_case = None
            if self.mask_type is not None and isinstance(self.mask_type, list):
                try:
                    masked_case = dict(copy.deepcopy(item['Patient_Case']))
                    for category in self.mask_type:
                        if masked_case['Patient_Factor'].__contains__(category) and category in ['Demographics', 'Past_Medical_History', 
                                                'Social_History', 'Review_of_Systems']:
                            masked_case['Patient_Factor'].pop(category)
                        elif masked_case['Patient_Factor'].__contains__(category) and category in ['Symptoms', 'History']:
                            masked_case['Patient_Factor'].pop('Symptoms')
                            masked_case['Patient_Factor'].pop('History')
                        elif masked_case['Test_Results'].__contains__(category) and category in ['Laboratory', 'Imaging']:
                            masked_case['Test_Results'].pop(category)
                            
                            masked_case['Test_Results'] = {k: v for k, v in masked_case['Test_Results'].items() if v != []}
                            if not masked_case['Test_Results']:
                                masked_case.pop('Test_Results')
      
                        elif masked_case.__contains__(category):
                            masked_case.pop(category)
                except:
                    print(f"Key {self.mask_type} not found in the dictionary.")
                    masked_case = None

            elif self.mask_type is not None and isinstance(self.mask_type, str):
                try:
                    masked_case = dict(copy.deepcopy(item['Patient_Case']))
                    if not masked_case['Test_Results']['Imaging'] and not masked_case['Test_Results']['Laboratory']:
                        masked_case.pop('Test_Results')
                    if self.mask_type in ['Demographics', 'Past_Medical_History', 
                                                    'Social_History', 'Review_of_Systems']:
                                masked_case['Patient_Factor'].pop(self.mask_type)
                    elif self.mask_type in ['Symptoms', 'History']:
                        masked_case['Patient_Factor'].pop('Symptoms')
                        masked_case['Patient_Factor'].pop('History')
                    elif self.mask_type in ['Laboratory', 'Imaging']:
                        masked_case['Test_Results'].pop(self.mask_type)

                        masked_case['Test_Results'] = {k: v for k, v in masked_case['Test_Results'].items() if v != []}
                        if not masked_case['Test_Results']:
                            masked_case.pop('Test_Results')
                    else:
                        masked_case.pop(self.mask_type)
                except:
                    print(f"Key {self.mask_type} not found in the dictionary.")

            if self.is_sentence:
                partial_patient_info = self.dict_to_sentence_text(masked_case) if masked_case is not None else None
            else:
                partial_patient_info = self.dict_to_formatted_text(masked_case) if masked_case is not None else None
            patient_case: PatientCase = {
                "patient_info": self.dict_to_formatted_text(item['Patient_Case']) if not self.is_sentence else self.dict_to_sentence_text(item['Patient_Case']),
                "partial_patient_info": None if masked_case is None else partial_patient_info,
                "ground_truth": item['Correct_Answer'],
                "metadata": {"id": item['id']} if 'id' in item else {}
            }
    
        elif self.data_source == "agentclinic_nejm":            
            patient_case: PatientCase = {
                "patient_info": item['Patient_Info'],
                "partial_patient_info": None,
                "ground_truth": item['Correct_Answer'],
                "metadata": {"id": item['id']} if 'id' in item else {}
            }
        elif self.data_source == "structural_medqa":
            body = item['OSCE_Examination']
            body.pop('Objective_for_Doctor')
            ground_truth = body.pop('Correct_Diagnosis')
              
            if self.mask_type is not None:
                masked_case = copy.deepcopy(body)
                try:
                    if isinstance(self.mask_type, list):
                        for category in self.mask_type:
                            if category in ['Demographics', 'Past_Medical_History', 
                                                    'Social_History', 'Review_of_Systems']:
                                masked_case['Patient_Actor'].pop(category)
                            elif category in ['Symptoms', 'History']:
                                masked_case['Patient_Actor'].pop('Symptoms')
                                masked_case['Patient_Actor'].pop('History')
                            elif category in ['Physical_Examination']:
                                masked_case.pop('Physical_Examination_Findings')
                            elif category in ['Laboratory', 'Imaging']:
                                if 'Test_Results' in masked_case:
                                    masked_case['Test_Results'].pop(category)
                            else:
                                masked_case.pop(category)
                    else:
                        if self.mask_type in ['Demographics', 'Past_Medical_History', 
                                                'Social_History', 'Review_of_Systems']:
                            masked_case['Patient_Actor'].pop(self.mask_type)
                        elif self.mask_type in ['Symptoms', 'History']:
                            masked_case['Patient_Actor'].pop('Symptoms')
                            masked_case['Patient_Actor'].pop('History')
                        elif self.mask_type in ['Physical_Examination']:
                            masked_case.pop('Physical_Examination_Findings')
                        elif self.mask_type in ['Laboratory', 'Imaging']:
                            if 'Test_Results' in masked_case:
                                masked_case.pop('Test_Results')
                        else:
                            masked_case.pop(self.mask_type)
                except:
                    print(f"Key {self.mask_type} not found in the dictionary.")
                patient_case: PatientCase = {
                    "patient_info": self.dict_to_formatted_text(body),
                    "partial_patient_info": self.dict_to_formatted_text(masked_case),
                    "ground_truth": ground_truth,
                    "metadata": {"id": item['id']} if 'id' in item else {}
                }
            else:
                patient_case: PatientCase = {
                    "patient_info": self.dict_to_formatted_text(body),
                    "partial_patient_info": None,
                    "ground_truth": ground_truth,
                    "metadata": {"id": item['id']} if 'id' in item else {}
                }
        
        return patient_case
    
    def get_subset_by_year(self, year: int) -> List[dict]:
        """
        Get a subset of patient cases based on a specified year.
        Args:
            year (int): Year to filter patient cases
        Returns:
            List[dist]: Subset of dictionaries containing patient cases
        """
        subset = []
        for item in self.data:
            item['Date'] = datetime.strptime(item['Date'], '%B %d, %Y')
            if item['Date'].year == year:
                subset.append(item)
            
        return subset

    def get_subset(self) -> List[dict]:
        """
        Get a subset of patient cases based on a specified mask type.

        Args:
            mask_type (str): Type of information to mask in the patient data
        Returns:
            dict: Subset of patient cases        
        """
        subset = []
        for item in self.data:
            if isinstance(self.mask_type, str):
                if self.mask_type == 'Past_Medical_History' and item['Patient_Case']['Patient_Factor']['Past_Medical_History'] != '':
                    subset.append(item)
                elif self.mask_type == 'Demographics' and item['Patient_Case']['Patient_Factor']['Demographics']!= '':
                    subset.append(item)
                elif self.mask_type == 'Symptoms' and (item['Patient_Case']['Patient_Factor']['Symptoms']['Primary_Symptom'] != '' or item['Patient_Case']['Patient_Factor']['History']!= ''):
                    subset.append(item)
                elif self.mask_type == 'Social_History' and item['Patient_Case']['Patient_Factor']['Social_History'] != '':
                    subset.append(item)
                elif self.mask_type == 'Review_of_Systems' and item['Patient_Case']['Patient_Factor']['Review_of_Systems'] != '':
                    subset.append(item)
                elif self.mask_type == 'Physical_Examination' and item['Patient_Case']['Physical_Examination'] != []:
                    subset.append(item)
                elif self.mask_type == 'Laboratory' and item['Patient_Case']['Test_Results']['Laboratory'] != []:
                    subset.append(item)
                elif self.mask_type == 'Imaging' and (item['Patient_Case']['Test_Results']['Imaging'] != [] and item['Patient_Case']['Test_Results']['Imaging'] != ''):
                    subset.append(item)
            else:
                subset.append(item)
                
        return subset

    def get_single_case(self, index: int, original_data: bool = False) -> PatientCase:
        """
        Get a single patient case by index.
        
        Args:
            index (int): Index of the patient case to retrieve
            
        Returns:
            PatientCase: The patient case at the specified index
            
        Raises:
            IndexError: If the index is out of range
        """
        if (self.data_source == 'structural_nejm' or self.data_source == 'structural_mediq') and self.mask_type:
            subset = self.get_subset()

            if original_data:
                if 0 <= index < len(subset):
                    return subset[index]
                raise IndexError(f"Index {index} out of range for dataset with {len(subset)} items")
            else:
                if 0 <= index < len(subset):
                    return self._convert_to_patient_case(subset[index])
                raise IndexError(f"Index {index} out of range for dataset with {len(subset)} items")
        else:
            data = copy.deepcopy(self.data[index])

            if 0 <= index < len(self.data):
                return self._convert_to_patient_case(data)
            raise IndexError(f"Index {index} out of range for dataset with {len(self.data)} items")

# Example usage:
# data_loader = MedicalDataLoader(
#     data_path="path/to/your/data.jsonl",
# )
# 
# # Process all batches
# for batch in data_loader:
#     for patient_state in batch:
#         result = compiled_graph.invoke(input=patient_state, config={"recursion_limit": 35})
#         print(f"Diagnosis: {result['diagnosis']}, Ground Truth: {result['ground_truth']}")
#
# # Or get a single item
# single_case = data_loader.get_single_item(0)
# result = compiled_graph.invoke(input=single_case, config={"recursion_limit": 35})



class TokenCounter:
    def __init__(self):
        """
        Initialize the TokenCounter.

        Args:
            encoding_name (str): Name of the encoding to use
        """
        self.total_tokens = 0

    @classmethod
    def count_tokens(cls, string: str) -> int:
        """
        Count the number of tokens in a string.
        Args:
            string (str): The string to count tokens for
        Returns:
            int: The number of tokens in the string
        """
        encoding = tiktoken.get_encoding("cl100k_base")
        num_tokens = len(encoding.encode(string))

        return num_tokens

    def update_token_usage(self, string: str) -> int:
        """
        Count the number of tokens in a string.
        Args:
            string (str): The string to count tokens for
        Returns:
            int: The number of tokens in the string
        """
        num_tokens = self.count_tokens(string)
        self.total_tokens += num_tokens


    def get_total_tokens(self) -> int:
        """
        Get the total number of tokens counted.
        Returns:
            int: The total number of tokens counted
        """
        return self.total_tokens


    def reset(self):
        """
        Reset the total number of tokens counted.
        """
        self.total_tokens = 0




class Logger:
    """
    Unified logging interface for the medical diagnosis system.
    
    This class initializes and provides access to multiple logger instances
    for different components of the system, configured from a YAML file.
    
    Attributes:
        system_logger: Logger for system-level messages
        patient_logger: Logger for patient-related messages
        doctor_logger: Logger for doctor agent messages
        audit_logger: Logger for audit trail messages
        tiktoken_logger: Logger for token counting messages
    """

    def __init__(self, log_path, config_path=None):
        """
        Initialize the Logger class.
        
        Args:
            log_path (str): Path to the directory where log files should be stored
            config_path (str): Path to the logging configuration YAML file
            
        Raises:
            FileNotFoundError: If the logging config file is not found
        """
        if config_path is None:
            config_path = os.path.join(DIR_PATH, 'logs', 'logging_config.yaml')

        # Check if config file exists
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Logging config file not found at {config_path}")
            
        # Load the config file
        with open(config_path, 'rt') as f:
            config = yaml.safe_load(f.read())

        # change log file path
        config['handlers']['file']['filename'] = os.path.join(log_path, config['handlers']['file']['filename'])
            
        # Configure the logging module with the config file
        logging.config.dictConfig(config)
        
        # Initialize logger instances
        self.system_logger = logging.getLogger('system')
        self.patient_logger = logging.getLogger('patient')
        self.doctor_logger = logging.getLogger('doctor')
        self.audit_logger = logging.getLogger('audit')
    
    def get_system_logger(self):
        """Get the system logger instance"""
        return self.system_logger

    def get_patient_logger(self):
        """Get the patient logger instance"""
        return self.patient_logger
        
    def get_doctor_logger(self):
        """Get the doctor logger instance"""
        return self.doctor_logger
        
    def get_audit_logger(self):
        """Get the audit logger instance"""
        return self.audit_logger


class Inquiries_Generator:
    def __init__(self):
        self.FACTOR_RANK = [
            'Laboratory',
            'Symptoms',
            'Social History',
            'Physical Examination',
            'Past Medical History',
            'Review of Systems',
            'Demographics'
        ]
    
    def generate_questions(self, patient_info: str) -> str:
        """
        Generate a question based on the patient information.
        Args:
            patient_info (str): The patient's information
        Returns:
            str: A question
        """
        pass

def search_children(parent_uri: str) -> set:
    # Disable SSL warnings - Only do this for testing/development
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    # fetch content and parse html
    def parse_html(url, client_id, client_secret):
        # Token endpoint and credentials
        token_endpoint = 'https://icdaccessmanagement.who.int/connect/token'
        client_id = client_id
        client_secret = client_secret
        scope = 'icdapi_access'
        grant_type = 'client_credentials'

        # Get the OAUTH2 token
        payload = {
            'client_id': client_id,
            'client_secret': client_secret,
            'scope': scope,
            'grant_type': grant_type
        }

        r = requests.post(token_endpoint, data=payload, verify=False).json()
        token = r['access_token']

        # HTTP header fields to set
        headers = {
            'Authorization': 'Bearer ' + token,
            'Accept': 'application/json',
            'Accept-Language': 'en',
            'API-Version': 'v2'
        }
        response = requests.get(url, headers=headers, verify=False)
        return response.json()
    
    body = parse_html(parent_uri, os.environ.get('ICD_CLIENT_ID'), os.environ.get('ICD_CLIENT_SECRET'))
    child = body.get("child", [])
    list = []
    for c in child:
        content = parse_html(c, os.environ.get('ICD_CLIENT_ID'), os.environ.get('ICD_CLIENT_SECRET'))
        list.append(content.get("title").get("@value"))
    return set(list)

def search_icd(search_term: str):
    # Disable SSL warnings - Only do this for testing/development
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    def parse_html_string(html_string):
        """Parse HTML string and extract clean text without HTML tags"""
        if not html_string:
            return ""
        
        # Use BeautifulSoup to parse and extract text
        soup = BeautifulSoup(html_string, 'html.parser')
        clean_text = soup.get_text()
        return clean_text.strip()
    # Token endpoint and credentials
    token_endpoint = 'https://icdaccessmanagement.who.int/connect/token'
    client_id = os.environ.get('ICD_CLIENT_ID')
    client_secret = os.environ.get('ICD_CLIENT_SECRET')
    scope = 'icdapi_access'
    grant_type = 'client_credentials'

    # Get the OAUTH2 token
    payload = {
        'client_id': client_id,
        'client_secret': client_secret,
        'scope': scope,
        'grant_type': grant_type
    }

    print(f"{client_id} {client_secret}")
    r = requests.post(token_endpoint, data=payload, verify=False).json()
    token = r['access_token']

    search_uri = 'https://id.who.int/icd/release/11/2022-02/mms/search'

    # HTTP header fields to set
    headers = {
        'Authorization': 'Bearer ' + token,
        'Accept': 'application/json',
        'Accept-Language': 'en',
        'API-Version': 'v2'
    }

    # Make search request
    search_response = requests.get(search_uri + f"?q={search_term}", headers=headers, verify=False)
    if search_response.status_code == 200:
        search_data = search_response.json()
        # Implement flexible search if no results found
        if 'destinationEntities' not in search_data or not search_data['destinationEntities'] or len(search_data['destinationEntities']) == 0:
            search_response = requests.get(search_uri + f"?q={search_term}&useFlexisearch=true", headers=headers, verify=False)
            if search_response.status_code == 200:
                search_data = search_response.json()
            else:
                print(f"Search failed with status code: {search_response.status_code}")
                return {}
        
        info_mapping = {}
        # sort the search_data by score
        search_data['destinationEntities'] = sorted(search_data['destinationEntities'], key=lambda x: x['score'], reverse=True)
        if len(search_data['destinationEntities']) > 0:
            result = search_data['destinationEntities'][0] 
            chapter = result.get('chapter', 'Unknown')
            code = result.get('theCode', 'Unknown')
            result['matchingPVs'] = sorted(result.get('matchingPVs', []), key=lambda x: x['score'], reverse=True)

            if len(result['matchingPVs']) > 0:
                result = result['matchingPVs'][0]
                # Get details about searches
                title = parse_html_string(result.get('label', 'Unknown'))
                entity_uri = result['foundationUri']
            elif len(result['matchingPVs']) == 0:
                title = parse_html_string(result.get('title', 'Unknown'))
                entity_uri = result['id']
                
            score = result.get('score', 0)
            
            # print("\n" + "=" * 80)
            # print(f"**Entity Details**")
            # print(f"Title: {title}")
            # print(f"URI: {entity_uri}")
            # print(f"Chapter: {chapter}")
            # print(f"Code: {code}")
            # print(f"Score: {score}")
            # print("=" * 80 + "\n")

            # Request for details
            entity_response = requests.get(entity_uri, headers=headers, verify=False)
            if entity_response.status_code == 200:
                hierarchy = []
                inclusions = []
                exclusions = []
                entity_data = entity_response.json()
                ent_title = entity_data.get('title', '').get('@value', '') if 'title' in entity_data else None
                definition = entity_data['definition'].get('@value', '') if 'definition' in entity_data else None
                long_definition = entity_data.get('longDefinition', '').get('@value', '') if 'longDefinition' in entity_data else None
                exclusion = entity_data.get('exclusion', {}) if 'exclusion' in entity_data else None
                exclusion_disease = [disease.get('label', '').get('@value', '') for disease in exclusion] if exclusion else []
                inclusion = entity_data.get('inclusion', {}) if 'inclusion' in entity_data else None
                inclusion_disease = [disease.get('label', '').get('@value', '') for disease in inclusion] if inclusion else []

                if ent_title:
                    hierarchy.append((f"**{title}**"))
                    if definition:
                        hierarchy[-1] += f"\n{definition}"
                    if long_definition:
                        hierarchy[-1] += f"\n{long_definition}"
                    if exclusion_disease:
                        exclusions.extend(exclusion_disease)
                    if inclusion_disease:
                        inclusions.extend(inclusion_disease)

                while (ent_title != 'ICD Category' and ent_title != 'ICD-11 for Mortality and Morbidity Statistics'):
                    parent = entity_data.get('parent', [])[0]
                    entity_response = requests.get(parent, headers=headers, verify=False)
                    if entity_response.status_code == 200:
                        entity_data = entity_response.json()
                        ent_title = entity_data.get('title', '').get('@value', '') if 'title' in entity_data else None
                        ent_definition = entity_data['definition'].get('@value', '') if 'definition' in entity_data else None
                        ent_long_definition = entity_data.get('longDefinition', '').get('@value', '') if 'longDefinition' in entity_data else None
                        exclusion = entity_data.get('exclusion', {}) if 'exclusion' in entity_data else None
                        exclusion_disease = [disease.get('label', '').get('@value', '') for disease in exclusion] if exclusion else []
                        inclusion = entity_data.get('inclusion', {}) if 'inclusion' in entity_data else None
                        inclusion_disease = [disease.get('label', '').get('@value', '') for disease in inclusion] if inclusion else []

                        if ent_title and ent_title != 'ICD Category' and ent_title != 'ICD-11 for Mortality and Morbidity Statistics':
                            hierarchy.append((f"**{ent_title}**"))
                            if ent_definition:
                                hierarchy[-1] += f"\n{ent_definition}"
                            if ent_long_definition:
                                hierarchy[-1] += f"\n{ent_long_definition}"
                            if exclusion_disease:
                                exclusions.extend(exclusion_disease)
                            if inclusion_disease:
                                inclusions.extend(inclusion_disease)

                chapter_desc = hierarchy[-1].split("**")[-2] if hierarchy else "Unknown"
                # Print all definitions in reverse (top-down hierarchy)
                # print("\n" + "=" * 80)
                # print("**Hierarchical Definitions**")
                # for defn in reversed(hierarchy):
                #     print("\n" + "-" * 40)
                #     print(defn)
                # print("=" * 80 + "\n")
                # print("\n" + "=" * 80)
                # print("**Exclusion Diseases**")
                # print(" | ".join(exclusions))
                # print("=" * 80 + "\n")
                # print("\n" + "=" * 80)
                # print("**Inclusion Diseases**")
                # print(" | ".join(inclusions))
                # print("=" * 80 + "\n")

                info_mapping[title] = {
                    'id': entity_uri,
                    'chapter': chapter.lstrip('0'),
                    'chapter_description': chapter_desc,
                    'code': code,
                    'definition': definition,
                    'long_definition': long_definition,
                    'hierarchy': hierarchy,
                    'exclusion_diseases': exclusions,
                    'inclusion_diseases': inclusions,
                    'score': [score],
                    'term': [search_term.lower()]
                }
                return info_mapping
            else:
                print(f"Error getting entity details: {entity_response.status_code}")
                return {}

    else:
        print(f"Search failed with status code: {search_response.status_code}")
        return {}
    
def get_chapter_from_diagnosis_direct(diagnosis: Dict) -> str:
    """Extract chapter directly from diagnosis item."""
    return diagnosis.get('chapter', '21')  # Default to chapter 21 if not found

def calculate_icd_contrast_chapter_similarity(yes_diagnoses: List[Dict], no_diagnoses: List[Dict]) -> float:
    """
    Calculate ICD contrast using chapter-based similarity scores.
    Uses chapter information directly from diagnosis items.
    
    Args:
        yes_diagnoses: Diagnosis list after "yes" answer
        no_diagnoses: Diagnosis list after "no" answer
        similarity_scores: Chapter similarity mapping from similarity_mapping.py
        
    Returns:
        float: ICD contrast score (average similarity between chapters)
    """
    if not yes_diagnoses or not no_diagnoses:
        return 0.0
    
    # Get chapters directly from diagnosis items
    yes_chapters = [get_chapter_from_diagnosis_direct(diag) for diag in yes_diagnoses]
    no_chapters = [get_chapter_from_diagnosis_direct(diag) for diag in no_diagnoses]
    
    print(f"YES Chapters: {yes_chapters} \nNO Chapters: {no_chapters}")
    
    if not yes_chapters or not no_chapters:
        return 0.0
    
    # Calculate pairwise similarities between yes and no chapters
    total_similarity = 0.0
    comparison_count = 0
    
    for yes_chapter in yes_chapters:
        for no_chapter in no_chapters:
            # Get similarity score from the mapping
            similarity = similarity_scores.get(yes_chapter, {}).get(no_chapter, 0.05)  # Default low similarity
            total_similarity += similarity
            comparison_count += 1
    
    if comparison_count == 0:
        return 0.0
    
    # Calculate average similarity
    avg_similarity = total_similarity / comparison_count
    return avg_similarity

def calculate_weighted_icd_contrast_chapter(yes_diagnoses: List[Dict], no_diagnoses: List[Dict]) -> float:
    """
    Calculate weighted ICD contrast using chapter similarity scores, considering confidence scores.
    Uses information directly from diagnosis items.
    
    Args:
        yes_diagnoses: Diagnosis list after "yes" answer
        no_diagnoses: Diagnosis list after "no" answer
        similarity_scores: Chapter similarity mapping
        
    Returns:
        float: Weighted ICD contrast score
    """
    if not yes_diagnoses or not no_diagnoses:
        return 0.0
    
    # Calculate weighted similarities
    total_weighted_similarity = 0.0
    total_weight = 0.0
    
    for yes_diag in yes_diagnoses:
        yes_chapter = get_chapter_from_diagnosis_direct(yes_diag)
        yes_confidence = yes_diag.get('confidence', 0.5)
        
        for no_diag in no_diagnoses:
            no_chapter = get_chapter_from_diagnosis_direct(no_diag)
            no_confidence = no_diag.get('confidence', 0.5)
            
            # Weight is the product of confidences
            weight = yes_confidence * no_confidence
            
            # Get similarity score
            similarity = similarity_scores.get(yes_chapter, {}).get(no_chapter, 0.05)
            
            total_weighted_similarity += similarity * weight
            total_weight += weight
    
    if total_weight == 0:
        return 0.0
    
    # Calculate weighted average similarity
    weighted_avg_similarity = total_weighted_similarity / total_weight
    return weighted_avg_similarity