import numpy as np
import torch
import joblib

from openai import OpenAI

from retrieval_utils.demo.model import text_encoder
from retrieval_utils.demo.load import load_unit_embeddings, load_splits, load_json
from retrieval_utils.tools.easyconvert import matrix_to, axis_angle_to
from retrieval_utils.transforms.smpl import RotTransDatastruct

import pickle
import argparse
import os
import datetime
import codecs as cs
import orjson  # loading faster than json
import json

import numpy as np
from tqdm import tqdm

import clip
import time
import random

import utils


# Replace `API-KEY` with openAI account api-key to generate LLM part-specific descriptions
client = OpenAI(api_key="API-KEY")

# Prompt function to generate LLM part-specific descriptions. It consists of (1) Task instructions, (2) Few-shot examples and (3) Query
prompt = """The instructions for this task is to describe the below body part position and movements in a sentence using simple language.
\n1) Torso \n2) Hands \n3) Fingers \n4) Legs \nHere are some examples of the question and answer pairs for this task: 
\nQuestion: Describe the below body parts position and movements involved in the action "a person squats down and puts their hands above their head" 
in a sentence using simple language. \nAnswer: \n1) Torso: The person bends their upper body forward, bringing their chest closer to their thighs 
while squatting down. \n2) Hands: The hands are raised above the head, reaching upward in a stretching motion. \n3) Legs: The person is squatting down, 
bending their knees and lowering their body toward the ground. \n4) Fingers: The fingers may be extended and open, pointing upward as the hands are placed 
above the head. \nQuestion: Describe the below body parts position and movements involved in the action "a person is swimming in the water" in a sentence 
using simple language. \nAnswer: \n1) Torso: The person's torso is floating horizontally in the water, moving smoothly from side to side as they swim. 
\n2) Hands: Their hands are pushing through the water, alternating in a circular motion, propelling them forward. \n3) Legs: The legs are kicking gently, 
helping to keep the body afloat and moving in a coordinated rhythm with the arms. \n4) Fingers: Fingers are closed together, and occasionally open to catch 
and pull the water during the swim strokes. \nQuestion: Describe the below body parts position and movements involved in the action [ACTION] in a sentence 
using simple language.\n1) Torso \n2) Hands \n3) Fingers \n4) Legs"""


def GPT_Completion(texts):
	response = client.completions.create(
	  model = "gpt-3.5-turbo-instruct",
	  prompt=texts,
	  max_tokens = 256
	)   
	gpt3_texts = []
	for i in range(len(response.choices)):
		gpt3_texts.append(response.choices[i].text)
	
	return gpt3_texts
