# import the relevant packages
import os
import csv
import requests
import time
import pandas as pd
import openai
from openai import OpenAI
import gradio as gr
import huggingface_hub
from datasets import load_dataset

api_key = os.environ.get("API_TOKEN")
headers = {
    'Authorization': 'Bearer ' + api_key,
    'Content-Type': 'application/json'
}

dataset = load_dataset('csv', data_files='https://huggingface.co/datasets/petcoblue/simulation_data/resolve/main/user_agents.csv')
user_agents = dataset['train'].to_pandas()
user_agents = user_agents.iloc[:,1:]
user_batch = user_agents[:10]

def create_description(row):
    description = (
        f"Imagine that you are currently {int(row['age'])} years old. You have {int(row['num_pets'])} pets "
        f"and spend an average of ${row['avg_spending']} on Petco purchases. "
        f"Your engagement with Petco marketing has a score of {int(row['engagement_score'])}. "
        f"You have an income level of {int(row['income_level'])} and "
        f"regularly buy items from Petco every {int(row['purchase_regularity'])} months. "
        f"It has been {int(row['time_since_last_purchase'])} days since your last purchase with Petco."
    )
    return description

question = (
    "Here are two images of Petco marketing emails:\n"
    "- Image 0 is shown first.\n"
    "- Image 1 is shown second.\n"
    "Which email are you more likely to click through? Just answer with 0 for the first image or 1 for the second image.\n"
    "Then, provide a list of up to five one-word characteristics of the email you chose that made you want to click through it. Separate each characteristic with a comma.\n\n"
    "Example response:\n"
    "1; Characteristics: Appealing, Sale, Bright, Simple, Exclusive\n"
)

def query_agent(description, question, image0, image1):
    image0_path = os.path.join(image_directory, image0)
    image1_path = os.path.join(image_directory, image1)
    base64_image0 = encode_image(image0_path)
    base64_image1 = encode_image(image1_path)

    payload = {
        "model": "gpt-4-vision-preview",
        "messages": [
            {"role": "system", "content": description},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image", "image_url": f"data:image/jpeg;base64,{base64_image0}"},
                    {"type": "image", "image_url": f"data:image/jpeg;base64,{base64_image1}"}
                ]
            }
        ],
        "max_tokens": 300,
        "logprobs": True,
        "top_logprobs": 1
    }

    for attempt in range(3):
        try:
            response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
            if response.status_code == 200:
                data = response.json()
                preference = data['choices'][0]['message']['content']
                top_logprobs = data['choices'][0]['logprobs']['content'][0]['top_logprobs']
                return preference, top_logprobs
            else:
                print(f"HTTP Error {response.status_code} on attempt {attempt + 1}")
        except requests.exceptions.RequestException as e:
            print(f"Request failed on attempt {attempt + 1}: {e}")
            time.sleep(1)
    else:
        print(f"Failed to analyze {image0} and {image1} after 3 attempts.")
        return None, None

def simulate(image0, image1):
    print(image0, image1)
    preferences = []
    reasons = []
    probs = []
    for index, user_agent in user_batch.iterrows():
        description = create_description(user_agent)
        preference, top_logprobs = query_agent(description, question, image0, image1)
        prob = np.round(np.exp(top_logprobs[0]['logprob']) * 100, 2)

        split = preference.split("; Characteristics: ")
        if len(split) == 2:
            choice, reasoning = split[0], split[1]
        else:
            print(preference)
            choice, reasoning = split[0], ""
        
        preferences.append(0 if "0" in choice else 1)    
        reasons.append(reasoning)
        probs.append(prob)
    
    avg_preference = sum(preferences) / len(preferences)
    avg_prob = sum(probs) / len(preferences)

    return avg_preference

subtitle = "Upload two images of emails and see which is generally preferred by Petco customers!"

demo = gr.Interface(fn=simulate, 
                    inputs=[gr.UploadButton("Click to Upload Email 0", file_types=["image"], file_count="single"), 
                            gr.UploadButton("Click to Upload Email 1", file_types=["image"], file_count="single")], 
                    outputs="text",
                    title="Pairwise Simulation of Petco Email Preference",
                    description=subtitle
                    )

if __name__ == "__main__":
    demo.launch()