import random
from pathlib import Path

from vocab import *

OUTPUT_FILE = Path(__file__).resolve().parent.parent / "data" / "ram_altman.txt"
END_TOKEN = "<|i_have_enough_ram|>"

N_SAMPLES = 200_000
CONVERSATION_CHANCE = 0.15
MAX_REPLY_WORDS = 170  # pair must fit in the 379 token context
MAX_CONVERSATION_REPLY_WORDS = 70


def fill(template, word=""):
    template = template.replace("{W}", word.upper())
    while "{size}" in template:
        template = template.replace("{size}", random.choice(SIZES), 1)
    while "{type}" in template:
        template = template.replace("{type}", random.choice(TYPES), 1)
    return template


def burst(lo, hi):
    groups = [RAM, BEGS, SIZES, TYPES, ["MORE"], LORE]
    parts = [random.choice(random.choices(groups, weights=[30, 20, 18, 12, 10, 10])[0])
             for _ in range(random.randint(lo, hi))]
    if random.random() < 0.85:
        parts.append(random.choice(RAM))
    return " ".join(parts)


def sentences(lo, hi):
    out = fill(random.choice(SENTENCES))
    for _ in range(random.randint(lo, hi) - 1):
        out += random.choice(SEPARATORS) + fill(random.choice(SENTENCES))
    return out


def spam():
    return " ".join([random.choice(SPAM_WORDS)] * random.randint(10, 30))


def rant(lo, hi):
    chunks = []
    for _ in range(random.randint(lo, hi)):
        r = random.random()
        chunks.append(sentences(1, 3) if r < 0.45 else burst(5, 12) if r < 0.85 else spam())
    return random.choice(SEPARATORS).join(chunks)


def tail(size):
    if size == "short":
        return random.choice(["", burst(1, 4)])
    if size == "normal":
        return random.choice([sentences(1, 3), burst(4, 10)])
    if size == "medium":
        return random.choice([sentences(3, 5), burst(10, 20), rant(2, 3)])
    return rant(7, 12) # Giant


def reply_size(user, rant=False):
    # longer messages tend to get longer replies; unknown topics get the rant
    # tuned so the reply length median lands around 15 words
    n = len(user.split())
    if rant:
        weights = [28, 50, 16, 6] if n <= 9 else [22, 50, 20, 8]
    else:
        weights = [48, 40, 10, 2] if n <= 3 else [43, 44, 11, 2] if n <= 9 else [33, 45, 16, 6]
    return random.choices(["short", "normal", "medium", "giant"], weights)[0]


def join(reply, rest):
    if not rest:
        return reply
    return reply + (" " if not reply[-1].isalnum() else random.choice(SEPARATORS)) + rest


def finish(reply, max_words=MAX_REPLY_WORDS):
    reply = " ".join(reply.split()[:max_words])
    if random.random() < 0.01:
        reply = join(reply, FALSE_ENDING)
    if reply[-1].isalnum():
        reply += random.choice(ENDINGS)
    return reply


# user text
def style(line):
    r = random.random()
    line = line[0].upper() + line[1:] if r < 0.45 else line.upper() if r < 0.55 else line.lower()
    if line[-1].isalnum() and random.random() < 0.5:
        line += random.choice(USER_PUNCTUATION)
    return line


def lengthen(line):
    if random.random() < 0.25:
        line = line.rstrip("?!.") + ". " + ". ".join(random.sample(FILLERS, random.randint(1, 3)))
    return line


def pluralize(word):
    if word.endswith("s"):  # news, chess, tennis, bus...
        return word
    if word.endswith("y") and word[-2] not in "aeiou":
        return word[:-1] + "ies"
    if word.endswith(("x", "ch", "sh", "o")):
        return word + "es"
    return word + "s"


def random_noun():
    word = random.choice(NOUNS)
    return pluralize(word) if random.random() < 0.4 else word


def topic_word():
    return random.choice(random.choice(list(TOPICS.values()))["words"])


def topic_pair():
    topic = random.choice(list(TOPICS.values()))
    word = random.choice(topic["words"])
    hook = random.choice(ECHO_HOOKS if random.random() < 0.25 else topic["hooks"])
    return lengthen(random.choice(TOPIC_TEMPLATES).format(w=word)), hook, word


def noun_pair():
    user = lengthen(random.choice(TOPIC_TEMPLATES).format(w=random_noun()))
    return user, random.choice(DISMISS_HOOKS), None


def reaction_pair():
    messages, replies = random.choice(list(REACTIONS.values()))
    return random.choice(messages), random.choice(replies), None


def math_pair():
    a, b = random.randint(0, 9), random.randint(0, 9)
    if random.random() < 0.5:
        user = random.choice(MATH_DIGITS).format(a=a, b=b)
    else:
        user = random.choice(MATH_WORDS).format(a=NUMBERS[a], b=NUMBERS[b])
    return user, random.choice(MATH_REPLIES).replace("{s}", str(a + b)), None


def ram_pair():
    n = random.choice(SIZES)
    user = random.choice(RAM_AMOUNT_USER).format(n=n)
    reply = random.choice(RAM_AMOUNT_REPLIES).replace("{n}", n)
    return user, reply, None


def generic_pair():
    if random.random() < 0.5:
        user = random.choice(GENERIC_QUESTIONS)
    else:
        user = random.choice(TOPIC_TEMPLATES).format(w=random.choice(GENERIC_WORDS))
    return lengthen(user), random.choice(SENTENCES), None


def first_message(kinds, weights):
    user, hook, word = random.choices(kinds, weights)[0]()
    return style(user), fill(hook, word or ""), word


def single_pair():
    kinds = [topic_pair, noun_pair, reaction_pair, math_pair, generic_pair, ram_pair]
    weights = [28, 12, 25, 8, 12, 15]
    kind = random.choices(kinds, weights)[0]
    user, reply, word = kind()
    rant = kind in (noun_pair, generic_pair, ram_pair)
    return style(user), finish(join(fill(reply, word or ""), tail(reply_size(user, rant=rant))))


def conversation():
    # remembers what the user said, still only wants RAM and asks for more 
    level = random.randint(1, 8)
    user, reply, _ = first_message([topic_pair, noun_pair, reaction_pair, ram_pair], [40, 30, 15, 15])
    reply += ("." if reply[-1].isalnum() else "") + f" I NEED {SIZES[level]}"
    turns = [(user, finish(join(reply, tail("normal")), MAX_CONVERSATION_REPLY_WORDS))]

    for _ in range(random.randint(1, 3)):
        prev = SIZES[level]
        level = min(level + random.randint(1, 3), len(SIZES) - 1)
        size = SIZES[level] if level < len(SIZES) - 1 else FINAL_OFFER

        if random.random() < 0.8:  # "and X?"
            echo = random.random() < 0.5  # topic words are echoed, anything else ignored
            word = topic_word() if echo else random_noun()
            user = style(random.choice(FOLLOW_UPS_WORD).format(w=word))
            replies = FOLLOW_UP_ECHO_REPLIES if echo else FOLLOW_UP_DISMISS_REPLIES
            reply = fill(random.choice(replies).replace("{esc}", size), word)
        else:  # "why?", "ok", "stop"
            messages, replies = random.choice(list(FOLLOW_UPS_PLAIN.values()))
            user = style(random.choice(messages))
            reply = fill(random.choice(replies).replace("{esc}", size).replace("{prev}", prev))

        rest = tail(random.choice(["short", "normal", "medium", "medium"]))
        turns.append((user, finish(join(reply, rest), MAX_CONVERSATION_REPLY_WORDS)))
    return turns


if __name__ == "__main__":
    random.seed(42)
    n_pairs = 0
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for _ in range(N_SAMPLES):
            turns = conversation() if random.random() < CONVERSATION_CHANCE else [single_pair()]
            for user, reply in turns:
                f.write(f"User: {user}\nRAM-ALTMAN: {reply}{END_TOKEN}\n")
            n_pairs += len(turns)
    print(f"{OUTPUT_FILE.name}: {N_SAMPLES:,} samples, {n_pairs:,} pairs")