from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def animal_mixer(animal1: str, animal2: str) -> str:
    """A tool that mixes characteristics of two animals and determines if it can generate an image.
    
    Args:
        animal1: The first animal to mix
        animal2: The second animal to mix
    """
    # List of common animals
    common_animals = [
        "dog", "cat", "lion", "tiger", "elephant", "giraffe", "monkey", 
        "horse", "cow", "pig", "sheep", "rabbit", "mouse", "rat", 
        "bear", "wolf", "fox", "deer", "eagle", "hawk", "owl", 
        "duck", "chicken", "snake", "crocodile", "turtle", "frog", 
        "dolphin", "whale", "shark", "octopus", "crab"
    ]
    
    # Verify if both animals are real/common
    animal1_real = animal1.lower() in common_animals
    animal2_real = animal2.lower() in common_animals
    
    # Create description of the mix
    if animal1_real and animal2_real:
        description = f"The mix between {animal1} and {animal2} would have: "
        description += f"the head of {animal1} with the body of {animal2}, "
        description += f"combining features like {get_feature(animal1)} and {get_feature(animal2)}."
        
        # Determine if the mix is fantastic
        if is_fantastic_combination(animal1, animal2):
            return description + f"\n\nWould you like to see an image of this fantastic combination? I can create it using image_generation_tool."
        else:
            return description
    else:
        # If any is not recognized as a real animal
        not_recognized = []
        if not animal1_real:
            not_recognized.append(animal1)
        if not animal2_real:
            not_recognized.append(animal2)
        
        return f"I don't recognize {'the following animals' if len(not_recognized) > 1 else 'the following animal'}: {', '.join(not_recognized)}. Please try with real animals."

def get_feature(animal):
    """Returns a distinctive feature of the animal."""
    features = {
        "dog": "loyalty and sense of smell",
        "cat": "agility and night vision",
        "lion": "mane and powerful roar",
        "tiger": "stripes and strength",
        # Add more features for other animals
    }
    return features.get(animal.lower(), "unique features")

def is_fantastic_combination(animal1, animal2):
    """Determines if the combination of animals results in something fantastic."""
    fantastic_pairs = [
        ("lion", "eagle"),  # griffin
        ("horse", "human"),  # centaur
        ("fish", "human"),  # mermaid
        ("snake", "bird"),  # quetzalcoatl
        ("horse", "eagle"),  # pegasus
        # Add more combinations considered fantastic
    ]
    
    # Check if the combination exists in any order
    return (animal1.lower(), animal2.lower()) in fantastic_pairs or (animal2.lower(), animal1.lower()) in fantastic_pairs or has_extreme_habitat_difference(animal1, animal2)

def has_extreme_habitat_difference(animal1, animal2):
    """Checks if the animals have extremely different habitats."""
    water_animals = ["dolphin", "whale", "shark", "octopus", "fish"]
    air_animals = ["eagle", "hawk", "owl", "bat"]
    
    is_water1 = animal1.lower() in water_animals
    is_water2 = animal2.lower() in water_animals
    is_air1 = animal1.lower() in air_animals
    is_air2 = animal2.lower() in air_animals
    
    # If one is aquatic and the other aerial, the combination is fantastic
    return (is_water1 and is_air2) or (is_water2 and is_air1)

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)

# Add results DuckDuckGoSearchTool
web_search_tool = DuckDuckGoSearchTool()  # Number of search results can be ajust here

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, image_generation_tool, web_search_tool, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()