from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import numpy
import pytz
import requests
import yaml

from Gradio_UI import GradioUI

from tools.final_answer import FinalAnswerTool

@tool
def X_transformation(array:numpy.ndarray)-> numpy.ndarray:
    """swap elemnents with position 0 and position 1 of array
    Args:
        array: numpy.ndarray
    """
    # create new array with the same lengh
    permute_array = numpy.zeros(len(array), dtype='int')
    # assign a value of element with position 0 of new array from element with position 1 of original array
    permute_array[0] = array[1]
    # assign a value of element with position 1 of new array from element with position 0 of original array
    permute_array[1] = array[0]
    # assign a values of elements of new array from elements of original array
    permute_array[2:] = array[2:]
    return permute_array

@tool
def L_transformation(array:numpy.ndarray)-> numpy.ndarray:
    """left cyclic shift of array
    Args:
        array: numpy.ndarray
    """
    # create new array with the same lengh
    permute_array = numpy.zeros(len(array), dtype='int')
    # assign a values of elements of new array from elements of original array
    permute_array[0:-1] = array[1:]
    # assign a value of element with position -1 of new array from element with position 0 of original array
    permute_array[-1] = array[0]
    return permute_array

@tool
def R_transformation(array:numpy.ndarray)-> numpy.ndarray:
    """a tool that swaps elements with position 0 and position 1 of array
    Args:
        array: numpy.ndarray
    """
    # create new array with the same lengh
    permute_array = numpy.zeros(len(array), dtype='int')
    # assign a values of elements of new array from elements of original array
    permute_array[1:] = array[0:-1]
    # assign a value of element with position 0 of new array from element with position -1 of original array
    permute_array[0] = array[-1]
    return permute_array


final_answer = FinalAnswerTool()

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='deepseek-ai/deepseek-coder-33b-instruct',
#model_id='deepseek-ai/DeepSeek-R1',
#model_id='Qwen/Qwen2.5-Math-72B-Instruct',
#model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud', 
#model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
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, 
        X_transformation, 
        L_transformation, 
        R_transformation
    ], ## add your tools here (don't remove final answer)
    max_steps=15,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
    additional_authorized_imports=["numpy", "heapq"],
)


GradioUI(agent).launch()