import os, sys
import pandas as pd

print("current directory:")
print(os.curdir)
print(os.getcwd())
TOP_COMMENTARIES_DIR = os.path.join(
   os.getcwd(),
   "prompting",
   "top_commentaries"
   )


ARTEMIS_EMOTIONS = [
    "amusement",         
    "anger",              
    "awe",               
    "contentment",       
    "disgust",           
    "excitement",        
    "fear",              
    "sadness",           
    "something else",    
]

DEFAULT_N_SAMPLES = 4
DEFAULT_RANDOM_STATE = 42
DEFAULT_OBJECT_THRESHOLD = 0.1





def fill_extracted_items(
    colors_list = None,
    objects_list = None,
    emotion = None,
    commentary = None,
):
  result = ""
  if colors_list:
    result += f"colors: {', '.join(colors_list)}"
    result += "\n"
  if objects_list:
    result += f"objects: {', '.join(objects_list)}"
    result += "\n"
  if emotion:
    result += f"emotion: {emotion}"
    result += "\n"
  if commentary:
    result += f"commentary: {commentary}"
    result += "\n"
  else:
    result += "commentary: "

  return result


def load_dataframe(
    csv_filepath,
):
    df = pd.read_csv(csv_filepath, index_col = 0)
    for stringified_col in  [
                                "maskrcnn_objects",
                                "colors",
                                "clip_recognized_objects", 
                            ]:
        df[stringified_col] = df[stringified_col].apply(eval)

    return df

TOP_COMMENTARIES_DFS = {
    emotion : load_dataframe(os.path.join(TOP_COMMENTARIES_DIR, f"top_{emotion.replace(' ', '_')}.csv"))
    for emotion in ARTEMIS_EMOTIONS
    }
  

def filter_items(
    items_and_probs_list,
    items_threshold = DEFAULT_OBJECT_THRESHOLD
    ):
    return [item for item, prob in items_and_probs_list if prob > items_threshold]

def get_random_samples_for_emotion(
    emotion,
    n_samples = DEFAULT_N_SAMPLES,
    random_state = DEFAULT_RANDOM_STATE,
    object_threshold = DEFAULT_OBJECT_THRESHOLD
):
    emotion_df = TOP_COMMENTARIES_DFS[emotion]
    samples = emotion_df.sample(n_samples, random_state = random_state)



    result = []
    for _, sample in samples.iterrows():
      colors_list = sample["colors"]
      commentary = sample["utterance"]
      emotion = sample["emotion"]

      objects_and_probs_list = sample["clip_recognized_objects"]
      objects_list = filter_items(objects_and_probs_list, object_threshold)

      entry = {
        "colors_list" : colors_list,
        "objects_list" : objects_list,
        "emotion" : emotion,
        "commentary" : commentary,
      }
      result.append(entry)

    return result

def get_subprompt_for_emotion(
    emotion,
    n_samples = DEFAULT_N_SAMPLES,
    random_state = DEFAULT_RANDOM_STATE,
    object_threshold = DEFAULT_OBJECT_THRESHOLD,
):
    random_samples = get_random_samples_for_emotion(
        emotion = emotion,
        n_samples = n_samples,
        random_state = random_state,
        object_threshold=object_threshold,
    )
    subprompt = [
            fill_extracted_items(**entry) for entry in random_samples    
    ]
    subprompt = "\n".join(subprompt)

    return subprompt


def get_subprompt_with_examples(
    n_samples_per_emotion = DEFAULT_N_SAMPLES,
    random_state = DEFAULT_RANDOM_STATE,
    object_threshold = DEFAULT_OBJECT_THRESHOLD,
):

    examples = [
        get_subprompt_for_emotion(
            emotion=emotion,
            n_samples = n_samples_per_emotion,
            random_state = random_state,
            object_threshold = object_threshold
        )   
        for emotion in ARTEMIS_EMOTIONS
    ]

    examples = "\n".join(examples)

    return examples

def get_user_prompt(
    colors_list,
    objects_list,
    emotion,
    n_samples_per_emotion = DEFAULT_N_SAMPLES,
    random_state = DEFAULT_RANDOM_STATE,
    object_threshold = DEFAULT_OBJECT_THRESHOLD,
):
    user_prompt= (
        "You have to write a commentary for an artwork.\n"
        "To write the commentary, you are given the objects present in the picture, "
        "the colors present in the picture, and the emotion the picture evokes.\n"
        "You are first shown several examples, and then have to give your commentary.\n"
        "First come the examples, and then the objects, colors, and emotion you will have to use for your commentary.\n"
        "Avoid explicitly mentioning the objects, or colors, or emotion, if it sounds more natural.\n"
        "Only write the commentary.\n"
        "\n"
        "EXAMPLES:"
        "\n\n"
        "{examples}"
        "\n"
        "Now, write your personal opinion about the picture."
        "\n"
        "{image_subprompt}"
    )

    examples = get_subprompt_with_examples(
        n_samples_per_emotion = n_samples_per_emotion,  
        random_state = random_state,
        object_threshold=object_threshold,
    )

    image_subprompt = fill_extracted_items(
        colors_list = colors_list,
        objects_list = objects_list,
        emotion = emotion,
        commentary = None,
    )


    result = user_prompt.format(examples = examples, image_subprompt = image_subprompt)

    return result