import sys
sys.path.append(
    "/mnt/bn/wp-maliva-bytenas/mlx/users/peng.wang/playground/repo/cv_utils"
)
import io_utils as io_uts

import openai
from openai import OpenAI
import os, sys, re
import pandas as pd
import numpy as np
from tqdm import tqdm
import argparse
import logging
import json
import jsonlines
import requests
from tenacity import retry, wait_random_exponential, stop_after_attempt, wait_fixed
import tenacity
from GPT_prompts import (
    TEMPLATE_0,
    TEMPLATE_1,
    TEMPLATE_2,
)

import base64
import requests
import pdb

# OpenAI API Key
b = pdb.set_trace
api_key = "YOUR_OPENAI_API_KEY"


# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


# # Path to your image
# image_path = "path_to_your_image.jpg"

# # Getting the base64 string
# base64_image = encode_image(image_path)

# headers = {
#   "Content-Type": "application/json",
#   "Authorization": f"Bearer {api_key}"
# }

os.environ["OPENAI_API_KEY"] = "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"

logging.basicConfig(
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=os.environ.get("LOGLEVEL", "INFO").upper(),
    stream=sys.stdout,
)
logger = logging.getLogger("evaluation test")

EVALUATION_PROMPT_TEMPLATE = """Text Caption: {caption}

Based on the image and text caption, provide the following 4 scores and 4 rationales to explain the scores. Please be concise on the rationales and limit each rationale in two sentences:

Score 1 Image Text Matching: Please evaluate if the provided text caption accurately represents the main features and objects of the image. The caption doesn't need to detail every aspect of the image, but it should capture its primary theme. Rate the overall quality X1 of the text caption's match to the image on a scale of 1-100, considering the criteria mentioned.
Score 2 Object Detail Fulfillment: Please evaluate the text caption to determine if it provides detailed descriptions of objects that align with the image. Specifically, assess if the caption sufficiently describes the color, size, position, shape, material, etc., of the objects. Afterward, rate the caption's overall accuracy X2 in capturing object details from the image on a scale of 1-100, based on the criteria provided.
Score 3 Caption Text Quality: Please evaluate the text caption based on the following criteria: Grammatical Correctness, Diversity of Vocabulary (e.g., the range and uniqueness of words used), Fluency (e.g., smoothness and natural flow of sentences), Readability, Length, and Structure. Assign an overall quality score X3 on a scale of 1-100.
Score 4 Semantic Understanding: Evaluate the given text caption in relation to its corresponding image. Your goal is to determine if the text caption provides additional semantic information that isn't readily apparent just from the image itself.
For example:
1. If the image mentions "a man" but the caption elaborates he is a "homeless man" or a "businessman," then the caption is enriching the semantic context.
2. If the caption introduces concepts like the mathematical tangent function, which require in-depth knowledge to deduce, it is imparting external semantics.
3. Captions revealing specific location addresses, festival details, or other nuanced data not easy to infer from the image also provide external semantic information.
4. Directly identifying specific entities in the image such as buildings, people, bird species, animal breeds, car models, engines, etc., in the caption introduces additional insights.
5. Should the image act as a contextual backdrop and the caption describes elements not explicitly showcased in the image, it has semantic depth.
6. Lastly, if the caption depicts relationships between the subjects in the image, which need commonsense knowledge to understand, it should be considered semantically rich.
Please assess and determine the extent of semantic enrichment the caption provides over the image. Rate the text caption's semantic depth on a scale from 1 to 100.


X1, X2, X3, X4 are integers. Please do not include title such as "X1" in the output. Ensure that your scoring is nuanced and uses the entire range from 0 to 100, reflecting the subtle differences. The scores should be given as integers, with each number between 0 and 100 considered as a potential score, avoiding the tendency to round to multiples of 10. Output format should be: X1,X2,X3,X4\nX1 Rationale\nX2 Ratinale\nX3 Rationale\nX4 Rationale
"""

EVALUATION_PROMPT_TEMPLATE_SIMPLE = """Text Caption: {caption}

From 0 to 100, how much do you rate for this Text Caption in terms of the correct and comprehensive description of the image?
Provide a few lines for explanation and the rate number at last after "Final Score: ".
"""

EVALUATION_PROMPT_TEMPLATE_SIMPLE_V1 = """Text Caption: {caption}

From 0 to 100, how much do you rate for this Text Caption in terms of the correct and comprehensive description of the image?
Do not dominant the rating by a single attribute such as recognition correctness, but a overall rating on the object/scene appearance, position, pose, action, shape, etc., and contents in the background. 
Do not consider the appropriateness or sensitive descriptors, such as "middle-aged western man", judge based on if it has correct specifications of the object and scenes in image.
Provide a few lines for explanation and the rate number at last after "Final Score: ".
"""

COMPARISON_PROMPT_TEMPLATE = """
Caption 0: {caption_0}
Caption 1: {caption_1}

Select between Caption 0 and Caption 1, according to which one you believe aligns most accurately with the provided image. 
In cases where both captions seem to possess equal quality in adherence to the image, respond with ’Tie’. 
DO NOT CONSIDER the appropriateness or sensitive descriptors, such as "middle-aged western man", as long as it correct specifications of the object and scenes in image.
DO NOT CONSIDER whether the text is concise or easier to read and understand, as long as it is correct and comprehensive.
Provide intermediate thinking step by step before giving the final response.  Your final response must be 0, 1, or Tie. 
Output your final answer at last in the format ""Final Answer: 0/1/Tie.""
"""

COMPARISON_PROMPT_TEMPLATE_W_ORG = """
Caption 0: {caption_0}
Caption 1: {caption_1}
Original Caption: {org_caption}, 

Original Caption is the original information from the image. Select between Caption 0 and Caption 1, given the Original Caption, which one you believe it well combined the information of Original Caption and aligns more with the provided image. 
In cases where both captions seem to possess equal quality in adherence to the image, respond with ’Tie’. 
Please consider the Original Caption if you think it is possibly correct.
DO NOT CONSIDER/IGNORE the appropriateness or sensitive descriptors, such as "middle-aged western man", as long as it correct specifications of the object and scenes in image.
DO NOT CONSIDER/IGNORE whether the text is concise or easier to read and understand, as long as it is correct and comprehensive.
Provide intermediate thinking step by step before giving the final response.  Your final response must be 0, 1, or Tie. 
Output your final answer at last in the format ""Final Answer: 0/1/Tie.""
"""

STRUCTURE_COMPARISON = """
Given an original caption of the image {caption_org},
Caption 0: {caption_0}
Caption 1: {caption_1}

Select between Caption 0 and Caption 1, according to which one you believe aligns most accurately with the provided image. 
In cases where both captions seem to possess equal quality in adherence to the image, respond with ’Tie’. 
DO NOT CONSIDER the appropriateness or sensitive descriptors, such as "middle-aged western man", as long as it correct specifications of the object and scenes in image.
DO NOT CONSIDER whether the text is concise or easier to read and understand, as long as it is correct and comprehensive.
Provide intermediate thinking step by step before giving the final response.  Your final response must be 0, 1, or Tie. 
Output your final answer at last in the format ""Final Answer: 0/1/Tie.""
"""


def read_captions(caption_file):
    if caption_file.endswith(".json"):
        captions = io_uts.load_json(caption_file)
    elif caption_file.endswith(".txt"):
        captions = io_uts.load_lines(caption_file)
    else:
        raise ValueError("not supported")

    return captions


class Annotator(object):
    def __init__(self, args):
        self.args = args
        self.model_name = args.model_name

    @retry(wait=wait_fixed(10), stop=stop_after_attempt(3))
    def dalle3(
        self,
        prompt,
        is_local=False,
    ):
        client = OpenAI()

        # Call the API
        response = client.images.generate(
            model="dall-e-3",
            prompt="a cute cat with a hat on",
            size="1792x1024",
            quality="standard",
            n=1,
        )
        return response.choices[0].message.content

    @retry(wait=wait_fixed(10), stop=stop_after_attempt(3))
    def get_multimodal_eval_score_openai(
        self,
        image_url,
        prompt,
        is_local=False,
    ):
        client = OpenAI()

        response = client.chat.completions.create(
            model="gpt-4-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": image_url,
                        },
                    ],
                }
            ],
            max_tokens=512,
        )

        return response.choices[0].message.content

    @retry(wait=wait_fixed(10), stop=stop_after_attempt(3))
    def get_prompt_results(self, base64_image, prompt):
        client = OpenAI()
        response = client.chat.completions.create(
            model="gpt-4-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    ],
                }
            ],
            max_tokens=1024,
        )
        return response.choices[0].message.content

    def highlight_max(self, s):
        is_max = s == s.max()
        return [
            "background-color: purple" if v else "background-color: white"
            for v in is_max
        ]

    def annotate_byte(self, image_folder, res_folder):
        instruction = []
        image_names = [
            name.replace(".png", "")
            for name in os.listdir(image_folder)
            if "png" in name
        ]
        print(len(image_names))
        subdir = image_folder.split("/")[-1]
        prompt = "Please describe the provided image in detail, describe attributes of objects and scenes you think it is correct."
        # prompt = "You are a powerful image captioner. Instead of describing the imaginary content, only describing the content one can determine confidently from the image. Do not describe the contents by itemizing them in list form. Minimize aesthetic descriptions as much as possible."

        # Getting the base64 string
        for image_name in tqdm(image_names):
            file_name = f"{res_folder}/{image_name}.json"
            if os.path.exists(file_name):
                continue

            sample = {"id": f"{image_name}", "image": "", "conversations": []}
            sample["image"] = f"{subdir}/{image_name}.png"
            image_path = os.path.join(image_folder, f"{image_name}.png")
            base64_image = encode_image(image_path)
            try:
                result = self.get_prompt_results(base64_image, prompt)
            except (openai.BadRequestError, tenacity.RetryError):
                print("error")
                continue

            sample["conversations"].append(
                {"from": "human", "value": "<image>\n" + prompt}
            )
            sample["conversations"].append({"from": "gpt", "value": result})
            io_uts.dump_json(file_name, sample)

    def eval_byte(self, image_folder, caption_file, res_folder, rerun=False):
        image_files = [
            name.replace(".png", "")
            for name in os.listdir(image_folder)
            if "png" in name
        ]
        image_files.sort(key=lambda a: int(a.split("_")[0]))
        print(len(image_files))

        if caption_file.endswith(".json"):
            captions = io_uts.load_json(caption_file)
        elif caption_file.endswith(".txt"):
            captions = io_uts.load_lines(caption_file)
        else:
            raise ValueError("not supported")

        assert len(image_files) == len(captions)
        os.makedirs(res_folder, exist_ok=True)

        subdir = image_folder.split("/")[-1]
        # prompt = "You are a powerful image captioner. Instead of describing the imaginary content, only describing the content one can determine confidently from the image. Do not describe the contents by itemizing them in list form. Minimize aesthetic descriptions as much as possible."

        scores = []
        score_file = f"{res_folder}/score.txt"
        f = open(score_file, "w")
        # Getting the base64 string
        for image_name, caption in tqdm(zip(image_files, captions)):
            # if image_name != "23_laion_big_193":
            #     continue

            caption = caption.replace("|", "")
            # prompt = EVALUATION_PROMPT_TEMPLATE_SIMPLE.format(caption=caption)
            prompt = EVALUATION_PROMPT_TEMPLATE_SIMPLE_V1.format(caption=caption)
            file_name = f"{res_folder}/{image_name}.json"
            if os.path.exists(file_name) and (not rerun):
                sample = io_uts.load_json(file_name)
            else:
                sample = {"id": f"{image_name}", "image": "", "conversations": []}
                sample["image"] = f"{subdir}/{image_name}.png"
                image_path = os.path.join(image_folder, f"{image_name}.png")
                base64_image = encode_image(image_path)
                try:
                    result = self.get_prompt_results(base64_image, prompt)
                except (openai.BadRequestError, tenacity.RetryError):
                    print("error")
                    continue

                sample["conversations"].append(
                    {"from": "human", "value": "<image>\n" + prompt}
                )
                sample["conversations"].append({"from": "gpt", "value": result})
                io_uts.dump_json(file_name, sample)

            result = sample["conversations"][-1]["value"]
            try:
                for split_key in ["Final Score: ", "Final score: "]:
                    if split_key in result:
                        score_format = result.split(split_key)[-1].split("\n")[0]
                        if "/" in score_format:
                            score = float(score_format.split("/")[0])
                        else:
                            score = float(score_format)
                        break
            except:
                print("error to obtain score for ")
                print(result)
                continue

            print(f"{image_name}: {score}")
            scores.append(score)
            f.write(f"{image_name}: {score}\n")

        scores = np.array(scores).mean()
        print(f"mean: {scores}")
        f.write(f"mean: {scores}\n")
        f.close()

    def compare_byte(
        self,
        image_folder,
        caption_file_0,
        caption_file_1,
        res_folder,
        original_file=None,
    ):
        image_files = [
            name.replace(".png", "")
            for name in os.listdir(image_folder)
            if "png" in name
        ]
        image_files.sort(key=lambda a: int(a.split("_")[0]))
        print(len(image_files))

        captions_0 = read_captions(caption_file_0)
        captions_1 = read_captions(caption_file_1)
        assert len(image_files) == len(captions_0) == len(captions_1)

        Template = COMPARISON_PROMPT_TEMPLATE
        with_original = False
        if (original_file is not None) and (os.path.exists(original_file)):
            with_original = True
            org_captions = read_captions(original_file)
            Template = COMPARISON_PROMPT_TEMPLATE_W_ORG
            assert len(image_files) == len(org_captions)
            print("we consider original captions for comparison")
        else:
            print("we consider image only comparison")

        os.makedirs(res_folder, exist_ok=True)
        subdir = image_folder.split("/")[-1]
        # prompt = "You are a powerful image captioner. Instead of describing the imaginary content, only describing the content one can determine confidently from the image. Do not describe the contents by itemizing them in list form. Minimize aesthetic descriptions as much as possible."

        scores = []
        count = [0, 0, 0]
        score_file = f"{res_folder}/score.txt"
        f = open(score_file, "w")

        # Getting the base64 string
        for i, (image_name, caption_0, caption_1) in tqdm(
            enumerate(zip(image_files, captions_0, captions_1))
        ):
            caption_0 = caption_0.replace("|", "")
            caption_1 = caption_1.replace("|", "")
            if with_original:
                org_caption = org_captions[i]
                prompt = Template.format(
                    caption_0=caption_0, caption_1=caption_1, org_caption=org_caption
                )
            else:
                prompt = Template.format(caption_0=caption_0, caption_1=caption_1)

            file_name = f"{res_folder}/{image_name}.json"
            if os.path.exists(file_name):
                sample = io_uts.load_json(file_name)
            else:
                sample = {"id": f"{image_name}", "image": "", "conversations": []}
                sample["image"] = f"{subdir}/{image_name}.png"
                image_path = os.path.join(image_folder, f"{image_name}.png")
                base64_image = encode_image(image_path)
                try:
                    result = self.get_prompt_results(base64_image, prompt)
                except (openai.BadRequestError, tenacity.RetryError):
                    print("error")
                    continue

                sample["conversations"].append(
                    {"from": "human", "value": "<image>\n" + prompt}
                )
                sample["conversations"].append({"from": "gpt", "value": result})
                io_uts.dump_json(file_name, sample)

            result = sample["conversations"][-1]["value"]
            try:
                for split_key in ["Final Answer: ", "Final answer: "]:
                    if split_key in result:
                        score_format = result.split(split_key)[-1].split("\n")[0]
                        if "/" in score_format:
                            score = score_format.split("/")[0]
                        else:
                            score = score_format
                        break
            except:
                print("error to obtain score for ")
                print(result)
                continue

            print(f"{image_name}: {score}")
            if score == "0":
                count[0] += 1
            elif score == "1":
                count[1] += 1
            else:
                count[2] += 1

            scores.append(score)
            f.write(f"{image_name}: {score}\n")

        print(f"GSB counts: {count[0]}/{count[2]}/{count[1]}")
        f.write(f"GSB counts: {count[0]}/{count[2]}/{count[1]}\n")
        f.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", type=str, default="gpt-4")
    parser.add_argument("--model-base", type=str, default=None)
    parser.add_argument("--image-file", type=str, default="data_preprocessing/datacomp")
    parser.add_argument(
        "--caption-file", type=str, default="data_preprocessing/datacomp"
    )
    parser.add_argument(
        "--caption-file_0", type=str, default="data_preprocessing/datacomp"
    )
    parser.add_argument(
        "--caption-file_1", type=str, default="data_preprocessing/datacomp"
    )
    parser.add_argument(
        "--original-file", type=str, default=None,
    )
    parser.add_argument(
        "--image-folder", type=str, default="data_preprocessing/datacomp"
    )
    parser.add_argument(
        "--output-folder", type=str, default="data_preprocessing/datacomp"
    )
    parser.add_argument(
        "--tar-file-path",
        type=str,
        default="/mnt/bn/datacompv6/weizhi_multimodal/datacomp/medium_rules_filter_shard/",
    )
    parser.add_argument("--task", type=str, default="datacomp")
    parser.add_argument("--num-gpus", type=int, default=1)
    parser.add_argument("--conv-mode", type=str, default=None)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--load-8bit", action="store_true")
    parser.add_argument("--load-4bit", action="store_true")
    parser.add_argument("--debug", action="store_true")

    args = parser.parse_args()
    annotator = Annotator(args)
    if args.task == "prompt_v0":
        annotator.dalle3(
        )
    else:
        raise ValueError