#!/usr/bin/env python

"""
This script generates synthetic agents with different attachments styles to respond to
conversations.
"""

import argparse
import json
import os
from datetime import datetime

import pandas as pd
from tqdm import tqdm

from synthetic_agents.app.chat_config import LanguageConfig, MemoryConfig
from synthetic_agents.app.playable_chat import PlayableChat
from synthetic_agents.common.constants import AI_CHAT_USER_AGENT_TYPE
from synthetic_agents.generator.life_fact import LifeFactGenerator
from synthetic_agents.generator.user_profile import UserProfileGenerator
from synthetic_agents.model.agent import Agent
from synthetic_agents.model.constants import (
    HUGGING_FACE_EMBEDDING_MODEL,
    OPENAI_GPT4_MODEL_NAME,
)

ATTACHMENT_STYLE_DESCRIPTION = {
    "avoidant": "You tend to maintain emotional distance, minimize closeness, and often reject or "
    "withdraw from intimacy in relationships. You have a selective memory, often "
    "downplaying or dismissing past experiences involving intimacy or vulnerability",
    "preoccupied": "You have a heightened need for reassurance, fear of abandonment, and a "
    "constant seeking of closeness and validation in relationships. You dwell on "
    "past experiences, focusing on moments of insecurity or inconsistency in "
    "relationships",
    "secure": "You are comfortable with intimacy, have balanced independence, effective "
    "communication, and a sense of safety and trust in relationships. You view your "
    "memories through a lens of safety and trust",
}

parser = argparse.ArgumentParser(
    description="Creates synthetic agents with a specific attachments style and creates a chat "
    "with that agent and an AAI interviewer."
)
parser.add_argument(
    "--aai_interviewer_id",
    type=int,
    required=True,
    help="The ID of the AAI interviewer who will interview the synthetic agent.",
)
parser.add_argument(
    "--num_agents",
    type=int,
    required=True,
    default=1,
    help="The number of synthetic agents to create.",
)
parser.add_argument(
    "--num_life_memories",
    type=int,
    required=True,
    default=10,
    help="The number of life memories to generate per agent.",
)

parser.add_argument(
    "--llm",
    type=str,
    required=False,
    default=OPENAI_GPT4_MODEL_NAME,
    help="LLM model to use when generating user profiles, life memories and configuring chats.",
)

# Temperature and top-p set based on recommended values for creative writing
# https://community.openai.com/t/cheat-sheet-mastering-temperature-and-top-p-in-chatgpt-api/172683
parser.add_argument(
    "--temperature",
    type=float,
    required=False,
    default=0.7,
    help="Temperature to use when generating user profiles and life memories.",
)
parser.add_argument(
    "--top_p",
    type=float,
    required=False,
    default=0.8,
    help="Top-p to use when generating user profiles and life memories.",
)

parser.add_argument(
    "--prompt_template_version",
    type=str,
    required=True,
    help="Version of the prompt template to assign to the agents.",
)
parser.add_argument(
    "--attachment_style",
    type=str,
    required=True,
    choices=["avoidant", "secure", "preoccupied"],
    help="The agent's synthetic_agents style.",
)
parser.add_argument(
    "--report_dir",
    type=str,
    required=False,
    help="An optional directory to save a CSV report of the created chats and agents. If not "
    "provided, no report will be created.",
)

args = parser.parse_args()

# Chat configurations

memory_config = MemoryConfig(
    remember_previous_sessions=False,
    retrieve_chat_memories=False,  # Focus on life memories only
    memory_retrieval_capacity=3,  # Retrieve 3 life memories at a time
    working_memory_capacity=1000,  # Maximum number of tokens
    context_window_capacity=1,  # Number of previous messages to use as context for query
    memory_embedding_model_name=HUGGING_FACE_EMBEDDING_MODEL,
)

# Temperature and top-p set based on recommended values for chatbot responses
# https://community.openai.com/t/cheat-sheet-mastering-temperature-and-top-p-in-chatgpt-api/172683
language_config = LanguageConfig(
    temperature=0.5,
    top_p=0.5,
    llm_name=args.llm,
)

if args.report_dir is not None:
    os.makedirs(args.report_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y.%m.%d--%H.%M.%S")
    filename = f"{args.attachment_style}_synthetic_agents_and_chats_{timestamp}.csv"

report_data = []
pbar = tqdm(range(args.num_agents))
for _ in pbar:
    pbar.set_description("User Profile")
    user_profile_gen = UserProfileGenerator(
        llm_name=args.llm, temperature=args.temperature, top_p=args.top_p
    )
    user_profile = user_profile_gen.generate(1)[0]

    pbar.set_description("Life Facts")
    life_facts_gen = LifeFactGenerator(
        user_profile=user_profile,
        llm_name=args.llm,
        temperature=args.temperature,
        top_p=args.top_p,
    )
    life_facts = life_facts_gen.generate(args.num_life_memories)

    user_profile["attachment_style_description"] = ATTACHMENT_STYLE_DESCRIPTION[
        args.attachment_style
    ]

    user_id = Agent.create(
        application_type="chat",
        agent_type=AI_CHAT_USER_AGENT_TYPE,
        prompt_template_version=args.prompt_template_version,
        attributes=user_profile,
        life_facts=life_facts,
    )

    chat_id = PlayableChat.create(
        user_agent_id=user_id,
        coach_agent_id=args.aai_interviewer_id,
        user_memory_config=memory_config,
        user_language_config=language_config,
        coach_memory_config=memory_config,
        coach_language_config=language_config,
    )

    report_data.append(
        {
            "attachment_style": args.attachment_style,
            "chat_id": chat_id,
            "user_id": user_id,
            "interviewer_id": args.aai_interviewer_id,
            "user_profile": json.dumps(user_profile),
            "life_facts": "\n".join(
                [
                    json.dumps(f.to_simplified_json(["creation_timestamp", "content"]))
                    for f in life_facts
                ]
            ),
        }
    )

    if args.report_dir is not None:
        pd.DataFrame(report_data).to_csv(f"{args.report_dir}/{filename}", index=False)
