# import the relevant packages
import os
import csv
import requests
import time
import base64
import pandas as pd
import numpy as np
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 encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

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):
    base64_image0 = encode_image(image0)
    base64_image1 = encode_image(image1)

    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):
    upload_file(image0)
    upload_file(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!"


from pathlib import Path

def upload_file(filepath):
    name = Path(filepath).name
    return [gr.UploadButton(visible=False), gr.DownloadButton(label=f"{name}", value=filepath, visible=True)]

# with gr.Blocks() as demo:
#     gr.Markdown("First upload a file and and then you'll be able download it (but only once!)")
#     with gr.Row():
#         u = gr.UploadButton("Upload a file", file_count="single")
#         d = gr.DownloadButton("Download the file", visible=False)

#     u.upload(upload_file, u, [u, d])
#     d.click(download_file, None, [u, d])

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()