from Modules.PoseEstimation.pose_estimator import calculate_angle, joints_id_dict, model
from langchain.tools import tool
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage
from langchain_mistralai.chat_models import ChatMistralAI

from operator import itemgetter
from typing import Dict, List, Union

from langchain_core.messages import AIMessage
from langchain_core.runnables import (
    Runnable,
    RunnableLambda,
    RunnableMap,
    RunnablePassthrough,
)

import numpy as np

# If api_key is not passed, default behavior is to use the `MISTRAL_API_KEY` environment variable.
llm = ChatMistralAI(model='mistral-large-latest', api_key="i5jSJkCFNGKfgIztloxTMjfckiFbYBj4")

@tool
def shoulder_angle(pose: list) -> float:
    
    """
    Computes the shoulder angle.

    Args:
        pose (list): list of keypoints

    Returns:
        arm_angle (float): arm angle with chest
    """
    right_elbow = pose[joints_id_dict['right_elbow']]
    right_shoulder = pose[joints_id_dict['right_shoulder']]
    right_hip = pose[joints_id_dict['right_hip']]   

    left_elbow = pose[joints_id_dict['left_elbow']]
    left_shoulder = pose[joints_id_dict['left_shoulder']]
    left_hip = pose[joints_id_dict['left_hip']] 

    right_arm_angle = calculate_angle(right_elbow, right_shoulder, right_hip)
    left_arm_angle = calculate_angle(left_elbow, left_shoulder, left_hip)

    return right_arm_angle


@tool
def elbow_angle(pose):
    """
    Computes the elbow angle.

    Args:
        pose (list): list of keypoints

    Returns:
        elbow_angle (float): elbow angle with chest
    """
    right_elbow = pose[joints_id_dict['right_elbow']]
    right_shoulder = pose[joints_id_dict['right_shoulder']]
    right_wrist = pose[joints_id_dict['right_wrist']]   

    left_elbow = pose[joints_id_dict['left_elbow']]
    left_shoulder = pose[joints_id_dict['left_shoulder']]
    left_wrist = pose[joints_id_dict['left_wrist']] 

    right_elbow_angle = calculate_angle(right_shoulder, right_elbow, right_wrist)
    left_elbow_angle = calculate_angle(left_shoulder, left_elbow, left_wrist)

    return right_elbow_angle


tools = [shoulder_angle, elbow_angle]

llm_with_tools = llm.bind_tools(tools)
tool_map = {tool.name: tool for tool in tools}

# prompt = ChatPromptTemplate.from_messages(
#     [
#         (
#             "system",
#             "You are a helpful assistant. Make sure to use the compute_right_knee_angle tool for information.",
#         ),
#         ("placeholder", "{chat_history}"),
#         ("human", "{input}"),
#         ("placeholder", "{agent_scratchpad}"),
#     ]
# )

# Construct the Tools agent
# curl_agent = create_tool_calling_agent(llm, tools, prompt)


pose_sequence = [
    # Pose 1
    [
        # Head
        [50, 50],
        # Shoulders
        [40, 80], [60, 80],
        # Elbows
        [30, 110], [70, 110],
        # Wrists
        [25, 140], [75, 140],
        # Hips
        [45, 180], [55, 180],
        # Knees
        [40, 220], [60, 220],
        # Ankles
        [35, 250], [65, 250],
    ],
    # Pose 2
    [
        # Head
        [60, 60],
        # Shoulders
        [50, 90], [70, 90],
        # Elbows
        [40, 120], [80, 120],
        # Wrists
        [35, 150], [85, 150],
        # Hips
        [55, 180], [65, 180],
        # Knees
        [50, 220], [70, 220],
        # Ankles
        [45, 250], [75, 250],
    ]]

    # Create an agent executor by passing in the agent and tools
# agent_executor = AgentExecutor(agent=curl_agent, tools=tools, verbose=True)
# agent_executor.invoke({"input": f"Compute shoulder and elbow angle and display them given the following pose estimation: {pose_sequence[0]}"})

def call_tools(msg: AIMessage) -> Runnable:
    """Simple sequential tool calling helper."""
    tool_map = {tool.name: tool for tool in tools}
    tool_calls = msg.tool_calls.copy()
    for tool_call in tool_calls:
        tool_call["output"] = tool_map[tool_call["name"]].invoke(tool_call["args"])
    return tool_calls


chain = llm_with_tools | call_tools

print(chain.invoke(f"What is the shoulder angle and elbow angle given the following pose estimation: {pose_sequence[0]}"))