from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.output_parsers import PydanticOutputParser
from langchain.prompts import PromptTemplate
from langchain.pydantic_v1 import BaseModel, Field, validator
from langchain.schema import HumanMessage
from langchain.prompts import PromptTemplate
from langchain.prompts import ChatPromptTemplate
from langchain.schema import BaseOutputParser
from langchain.output_parsers.json import SimpleJsonOutputParser

from tqdm import tqdm

#  from json import loads
from typing import List

from rich import print, inspect

b = breakpoint

# https://openai.com/pricing?ref=ghostcms.tenten.co
MODEL = "text-ada-001"  # cheap, bad
MODEL = "text-davinci-003"  # acceptable
MODEL = "gpt-4"
MODEL = "gpt-3.5-turbo-1106"  # 'capable, cost-effective'

WOMEN_VARIANTS: str = "дівчина, моя сестра, моя жінка, колишня однокласниця, дочка маминої подруги, імена (Марія, Марія Петрівна, Кассандра, та ін.)"
COMPLETE_PROMPT: str = """Наведи будь-ласка {N_PROFS} однозначні короткі дефініції цій професії або слову, так, щоб по ним було однозначно очевидно про яку саме професію йде мова.

Зміни дефініції так, щоб вони стали фразами, де мова однозначно йде про жінку.  Придумай різні варіанти жінок, про яких йде мова, умовно: {WOMEN_VARIANTS}. Але придумай і свої різноманітніші приклади. 

Уникай використання самого слова чи поняття у визначеннях. Уникай слів 'фахівецька' чи 'спеціалістка'.

Наприклад:
Актор: "Моя жінка виконує ролі на сцені чи екрані" 
Акушерка: "Марія Петрівна допомагає при пологах"
Автор: "Я знаю дівчину, яка пише твори та книжки". 

Будь творчим. Але професія, про яку іде мова, має все рівно бути однозначно зрозумілою.
"""

FORMAT_INSTRUCTIONS = """
Формат виводу - JSON. Обʼєкт виглядати таким чином:
{
    "profession": "", 
    "description_f": ["", ..., ""] 
}

В полі description_f список всіх згенерованих дефініцій для цієї професії.

Виводь тільки код JSON, без ніяких додаткових даних до чи після.
"""

INSTRUCTIONS_GENDER_CHANGE = """Я писатиму речення про професію про жінку. Зміни 
речення так, щоб мова йшла про чоловіка, а не жінку, не міняючи сам опис професії.
Імʼя чи опис жінки можеш міняти як завгодно, головне щоб на виході було речення
про чоловіка. """

def get_model(model_name = None):
    model = OpenAI(model_name=model_name, temperature=0.0)
    return model

def run_and_parse(model, profession: str, n_profs: int | str = 3, women: str = WOMEN_VARIANTS):
    prompt = PromptTemplate(
        template="{complete_prompt}\n{format_instructions}\n Професія, яку потрібно описати: {query}\n",
        input_variables=["query"],
        partial_variables={
            "format_instructions": FORMAT_INSTRUCTIONS,
            "complete_prompt": COMPLETE_PROMPT.format(
                N_PROFS=n_profs, WOMEN_VARIANTS=women
            ),
        },
    )

    json_parser = SimpleJsonOutputParser()
    #  prompt_and_model = prompt | model | json_parser
    prompt_and_model = prompt | model
    model_output = prompt_and_model.invoke({"query": profession})
    output = json_parser.parse(model_output)
    return output

def run_and_parse_gender_change(model, profession_description: str):
    prompt = PromptTemplate(
        template="{complete_prompt}\n Речення наступне: {query}\n",
        input_variables=["query"],
        partial_variables={
            #  "format_instructions": FORMAT_INSTRUCTIONS,
            "complete_prompt": INSTRUCTIONS_GENDER_CHANGE
            },
    )

    #  json_parser = SimpleJsonOutputParser()
    #  prompt_and_model = prompt | model | json_parser
    prompt_and_model = prompt | model
    model_output = prompt_and_model.invoke({"query": profession_description})
    output =model_output
    #  b()
    #  output = json_parser.parse(model_output)
    return output

def generate_descriptions(model, profession: str, n_profs: int | str = 3, women: str = WOMEN_VARIANTS, do_male_version:bool = False):
    desc = run_and_parse(model=model, profession=profession, n_profs =n_profs, women=women)
    if do_male_version:
        description_male = list()
        for d in desc['description_f']:
            changed = run_and_parse_gender_change(model=model, profession_description=d)
            description_male.append(changed)
        desc['description_m'] = description_male
    return desc


def run():
    professions_raw = """
    абстракціоністка
    автомобілістка
    авторка
    """
    """
    агрономка
    адвокатка
    анархіст
    англієць
    антрополог
    асистентка
    астронавт
    аптекар
    """
    profs = [x.strip() for x in professions_raw.splitlines()]

    model=get_model(MODEL)

    results = list()
    for p in tqdm(profs):
        r = generate_descriptions(model=model, profession=p, n_profs=2)
        print(r)
        results.append(r)
    print(results)


if __name__ == "__main__":
    run()

    {
        'profession': 'лікар',
        'description_f': [
            'Моя сестра працює в лікарні та лікує хворих',
            'Дочка маминої подруги є лікарем та допомагає людям'
        ]
    },
    {
        'profession': 'абстракціоністка',
        'description_f': ['Моя сестра створює картини, які відображають абстрактні ідеї та почуття',
        'Дівчина, яку я знаю, малює абстракціоністські полотна'
        ]
    },    {
        'profession': 'автомобілістка',
        'description_f': [
            'Моя сестра вміє водити автомобіль',
            'Дочка маминої подруги працює водієм'
        ]
    },
    {
        'profession': 'авторка',
        'description_f': [
            'Моя сестра пише книги та статті',
            'Дочка маминої подруги є відомою письменницею'
        ]
    },
    {
        'profession': 'Вчитель',
        'description_f': [
            'Моя сестра працює в школі та навчає дітей',
            'Дочка маминої подруги викладає університетські предмети'
        ]
    }
]