#imports
import gradio as gr
import os
from PIL import Image
import io
from transformers import ViltProcessor, ViltForQuestionAnswering
import torch
import pandas as pd
import numpy as np
import random
from diffusers import DiffusionPipeline
from transformers import pipeline

#function to turn dataframe into prompt
def format_prompt_from_dataframe(df):
    """
    Create a prompt-style phrase from the DataFrame output.

    Parameters:
    df (pd.DataFrame): A DataFrame where each column represents the answers to the questions.

    Returns:
    str: A prompt-style phrase describing the product and image.
    """
    if df.empty or df.shape[0] == 0:
        return "No data available to create a prompt."

    # Assuming the DataFrame has one row and columns match the question list
    row = df.iloc[0]
    
    color = row.get('color', 'Unknown')
    style = row.get('style', 'Unknown')
    vibe = row.get('vibe', 'Unknown')
    product = row.get('product', 'Unknown')
    personalizable = row.get('personalizable', 'Unknown')
    material = row.get('material', 'Unknown')
    event = row.get('event', 'Unknown')
    font = row.get('font', 'Unknown')

    # Construct the prompt-style phrase
    phrase = (
        f"A {product} made from {material} for {event}. "
        f"The color is {color} and the style is {style} with a font {font}. "
        f"The vibe is {vibe} and it is {'personalizable' if personalizable.lower() == 'yes' else 'not personalizable'}."
    )

    return phrase

#function to get most frequent values from dataframe
def most_frequent_values(df):
    # Create an empty dictionary to store the most frequent values
    most_frequent = {}

    # Iterate over each column in the DataFrame
    for column in df.columns:
        # Create a dictionary to store the scores for each unique answer
        scores = {}

        # Iterate over each row in the column
        for row in df[column]:
            # Iterate over each entry in the list of dictionaries
            for entry in row:
                answer = entry['answer']
                score = entry['score']
                # Accumulate the score for each answer
                if answer in scores:
                    scores[answer] += score
                else:
                    scores[answer] = score

        # Find the answer with the highest total score
        most_common = max(scores, key=scores.get)
        # Add it to the dictionary
        most_frequent[column] = most_common

    # Convert the dictionary to a DataFrame
    most_frequent_df = pd.DataFrame(most_frequent, index=[0])

    return most_frequent_df


#funtion to turn list outputed into a dataframe
def list_to_dataframe(attribute_list):
    # Row indices
    row_indices = ['color', 'style', 'vibe', 'product',
               'personalizable', 'material', 'event', 'font']
    # Split the list into rows of 8 elements each
    rows = [attribute_list[i:i + len(row_indices)] for i in range(0, len(attribute_list), 8)]
    # Create the DataFrame
    df = pd.DataFrame(rows, columns=row_indices)

    return df

vqa_pipeline = pipeline("visual-question-answering",model='dandelin/vilt-b32-finetuned-vqa')
def answer_question(gallery):
  try:
      answer_list = []
      question_list = ['What color is in the image?',
                   'What style is the image?',
                   'What vibe is the image?',
                   'What kind of product is the image?',
                   'Is the product in the image personalizable?',
                   'What material is the product in the image made of?',
                   'What event is the product in the image made for?',
                   'What font is in the image?']
      for i,c in gallery:
        for q in question_list:
          predicted_answer = vqa_pipeline(i, q, top_k=3)
          answer_list.append(predicted_answer)
      df = list_to_dataframe(answer_list)
      frequent= most_frequent_values(df)
      prompt = format_prompt_from_dataframe(frequent)

      return prompt
      
  except Exception as e:
        return f"An error occurred: {e}"

#function for returning image with prompt
def infer(prompt):

    seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)
    #prompt = 'ceramic mug with heart'
    image = pipe(
        prompt = 'create a mockup photo for an etsy listing. the product is '+prompt,
        guidance_scale = 0,
        num_inference_steps = 2,
        width = 512,
        height = 512,
        generator = generator
    ).images[0]

    return image


#def infer(prompt):

#    image = pipeline(
 #       prompt = 'create a mockup photo for an etsy listing. the product is '+prompt
 #   ).images[0]
#
 #   return image
    
device = "cuda" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
    pipe.enable_xformers_memory_efficient_attention()
    pipe = pipe.to(device)
else:
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
    pipe = pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
#now try it with gradio interface and upload gallery
if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

#pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
#pipeline.to("cuda")

with gr.Blocks() as demo:
    title = gr.Markdown("""
    <div align="center">
    <strong style = 'font-size: 24px;''>Finding Common Trends</strong>
    </div>""")
    
    instructions = gr.Markdown("""
    <div align="center">
    <strong style = 'font-size: 16px;''>Upload a Few Product Images to Find Commonalites Between Them</strong>
    </div>""")

    gallery = gr.Gallery(label="Gallery",type='pil')
    #question = gr.Textbox(label="question")
    

    #name buttons to display
    image_button = gr.Button("predict trends")
    answer = gr.Textbox(label='Predicted answer')
    text_button = gr.Button("show image")

    
    image_button.click(fn=answer_question, inputs=[gallery], outputs=[answer])
    result = gr.Image(label="Result", show_label=False)
    text_button.click(fn=infer,inputs =[answer],outputs=[result])
    

demo.launch()