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

from Gradio_UI import GradioUI


@tool
def maths(operation: str, a: float, b: float | None = None) -> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that can perform various mathematical operations and provide an answer to the query. The operations should be simple and should contain at most two operands.
    Args:
        a: Float operand a.
        b: Float operand b. b should be None if operation is 'not', 'is_prime' or 'fibonacci'.
        operation: A string telling which mathematical operation to perform. In some operations, the argument b should not be given.
        Available operations are 'addition', 'subtraction', 'division', 'multiplication', 'modulo', 'power': a raised to b, 'equal': checks if a is equal to b,
        'gcd' / 'hcf', 'lcm', 'greater': checks if a > b, 'less_than': checks if a < b, bitwise operators 'or', 'and', 'xor', 'left_shift': a << b, 'right_shift': a >> b,
        'log': log a to the base b, 'root': a raised to the power 1/b
    """
    if b is not None:
        if operation == 'addition':
            return f'Addition of {a} and {b} is {a+b}'
        elif operation == 'subtraction':
            return f"Subtraction of {a} and {b} is {a-b}"
        elif operation == 'division':
            if b == 0:
                return "Mathematical error. Invalid operation - cannot divide by zero."
            else:
                return f'{a} divide by {b} is {a/b}'
        elif operation == 'multiplication':
            return f'Multiplication of {a} and {b} is {a*b}'
        elif operation == 'modulo':
            return f'{a} modulo {b} is {a%b}'
        elif operation == 'power':
            return f'{a} raised to {b} is {pow(a,b)}'
        elif operation in ('gcd', 'lcm', 'hcf'):
            x, y = a, b
            while b:
                a, b = b, a % b
            if operation in ('gcd', 'hcf'):
                return f'Greatest Common Divisor or Highest Common Factor of {x} and {y} is {a}'
            else:
                return f'Lowest Common Multiple of {x} and {y} is {(x*y)/a}'
        elif operation == 'equal':
            if a == b:
                return f"{a} and {b} are equal"
            else:
                return f'{a} and {b} are not equal'
        elif operation == 'greater':
            if a > b:
                return f'{a} is greater than {b}'
            else:
                return f'{a} is not greater than {b}'
        elif operation == 'less_than':
            if a < b:
                return f'{a} is less than {b}'
            else:
                return f'{a} is not less than {b}'
        elif operation == 'or':
            return f'Bitwise {a} or {b} is {a|b}'
        elif operation == 'and':
            return f'Bitwise {a} and {b} is {a&b}'
        elif operation == 'xor':
            return f'Bitwise {a} xor {b} is {a^b}'
        elif operation == 'left_shift':
            return f'Left shift {a} by {b} bits is {a << b}'
        elif operation == 'right_shift':
            return f'Right shift {a} by {b} bits is {a >> b}'
        elif operation == 'log':
            return f'Log {a} to the base {b} is {log(a,b)}'
        elif operation == 'root':
            if b == 2:
                return f'Square root of {a} is {sqrt(a)}'
            elif b == 3:
                return f'Cube root of {a} is {pow(a, 1.0/3)}'
            else:
                return f'{b}th root of {a} is {pow(a, 1.0/b)}'
    else:
        if operation == 'not':
            return f'Not {a} is {~a}'
        elif operation == 'is_prime':
            if a < 0:
                return f'Negative numbers cannot be considered as prime numbers'
            elif a < 2:
                return f'{a} is neither a prime number nor a composite number'
            elif a % 2 == 0:
                return f'{a} is not a prime number'
            else:
                for i in range(3, int(a**0.5)+1, 2):
                    if a % i == 0:
                        return f'{a} is not a prime number'
                return f'{a} is a prime number'
        elif operation == 'fibonacci':
            n = [0, 1, 1, 2, 3, 5]
            if a < 6:
                return f'The fibonacci number at position {a} in the sequence is {n[a]}'
            else:
                phi = (1 + sqrt(5)) / 2
                return f'The fibonacci number at position {a} in the sequence is {round(pow(phi,a)/sqrt(5))}'
    return 'Invalid operation or arguments'

@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,
)


# 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, DuckDuckGoSearchTool(), get_current_time_in_timezone, maths],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()