# 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
from wordcloud import WordCloud
import matplotlib.pyplot as plt
import ast

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

    preference = 0 if avg_preference < 0.5 else 1
    print(preferences, preference)

    all_reasons = ' '.join([word.strip() for item in reasons for word in item.split(',')])
    wordcloud = WordCloud(width = 800, height = 400, background_color ='white').generate(all_reasons)

    plt.figure(figsize=(10,5))
    plt.imshow(wordcloud, interpolation='bilinear')
    plt.axis("off")
    plt.show()

    return preference, plt

# defines who can enter the application with the secrets that are set up
user_db = { 
           os.environ["username"]: os.environ["password"],
          }

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

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

if __name__ == "__main__":
    demo.launch(
        auth=lambda u, p: user_db.get(u) == p,
        auth_message="Welcome! Please enable third-party cookies or you will not be able to login."
        )