from smolagents import CodeAgent, HfApiModel, tool
import yaml
from Gradio_UI import GradioUI
import requests

# Example tool to analyze waste data (Пример инструмента для анализа данных о вторсырье)
@tool
def analyze_waste_data(data: dict) -> str:
    """Analyzes waste data and provides recommendations.
    (Анализирует данные о вторсырье и предоставляет рекомендации.)

    Args:
        data: A dictionary containing waste data, including types and quantities of waste.
        (Словарь с данными о вторсырье, включающий типы и количество отходов.)
    """
    # Example analysis of data (Пример анализа данных)
    waste_types = data.get("types", {})
    recommendations = []

    for waste_type, amount in waste_types.items():
        if amount > 100:  # Example threshold value (Пример порогового значения)
            recommendations.append(f"It is necessary to collect {waste_type}.")
            # (Необходимо вывезти {waste_type}.)

    return "\n".join(recommendations)

# Example tool to visualize collection routes (Пример инструмента для визуализации маршрутов)
@tool
def visualize_collection_routes(locations: list) -> str:
    """Visualizes optimal routes for waste collection.
    (Визуализирует оптимальные маршруты для сбора вторсырья.)

    Args:
        locations: A list of coordinates for container locations.
        (Список координат местоположений контейнеров.)
    """
    # Example use of a maps API to visualize routes (Пример использования API карт для визуализации маршрутов)
    map_url = "https://maps.google.com/?q=" + "&q=".join([f"{lat},{lon}" for lat, lon in locations])
    return f"Optimal routes can be viewed at this link: {map_url}"
    # (Оптимальные маршруты можно посмотреть по ссылке: {map_url})

# Load model and create agent (Загрузка модели и создание агента)
final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, analyze_waste_data, visualize_collection_routes],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()