### SimpleMath. answer to simple math questions of + & - from 0 to 2000.


from collections import deque
from fastai.text.all import load_learner
import sys
import re

# ─── Config ───────────────────────────────────────────────────────────────
MAX_HISTORY_CHARS    = 800
MAX_HISTORY_MESSAGES = 1 # 1 = no memory. (this model does not support memory!)
GENERATE_TOKENS      = 70
TEMPERATURE          = 0.3

def evaluate_placeholders(text: str) -> str: # sometimes the model will return the answer like this 1 + 1 = {1 + 1} that can be calculated and returned like this.
    def repl(match):
        expr = match.group(1)
        try:
            if re.fullmatch(r"[\d\s\+\-\*\/]+", expr):
                return str(eval(expr))
        except Exception:
            pass
        return match.group(0)

    return re.sub(r"\{([^{}]+)\}", repl, text)

def remove_before_first_colon(s: str) -> str:
    return s.split("BOT :", 1)[-1]
def remove_before_last_colon(s: str) -> str:
    return s.rsplit(":", 1)[-1]
def remove_after_user(text):
    keyword = "USER"
    index = text.find(keyword)
    if index != -1:
        return text[:index + len(keyword)]
    return text
def remove_after_bot(text):
    keyword = "BOT"
    index = text.find(keyword)
    if index != -1:
        return text[:index + len(keyword)]
    return text

def truncate(answer):
    for sep in ["\n", "USER:", "BOT:"]:
        if sep in answer:
            answer = answer.split(sep)[0]

    answer = remove_before_first_colon(answer)
    answer = remove_after_user(answer)
    answer = remove_after_bot(answer)

    answer = answer.replace(": USER", "").replace(" USER", "").replace("USER", "").replace(" !", "!").replace(" .", ".").replace(" ,", ",").replace(": BOT", "").replace(" BOT", "").replace("BOT", "").replace(" `", "`").replace(' "', '"').replace(" ’", "’").replace("do n'", "don'").replace("do n’", "don’")
    answer = answer.replace(" '", "'").replace(" :", ":").replace(" (", "(").replace(" )", ")").replace(" ?", "?").replace("Open Assistant", "Bomba-1") # The bot thinks he is named "Open Assistant", i guess something in the datasets.

    return answer.strip()

def load_models():
    print("🤖  Loading models…")
    chat_model = load_learner("model/SimpleMath.pkl")
    chat_model.model.eval()
    return chat_model

def main():
    chat_model = load_models()
    history = deque()
    print("💬  Ready! (empty line to quit)\n")

    while True:
        try:
            user = input("USER: ").strip()
            if not user:
                break

            history.append(f"USER: {user}")
            while len(history) > MAX_HISTORY_MESSAGES:
                history.popleft()

            prompt_lines = list(history)
            prompt_text = " ".join(history).replace("\n"," ")
            if len(prompt_text) > MAX_HISTORY_CHARS:
                prompt_text = prompt_text[-MAX_HISTORY_CHARS:]
            prompt = f"{prompt_text} BOT: "

            generated = chat_model.predict(
                prompt,
                n_words=GENERATE_TOKENS,
                temperature=TEMPERATURE,
                min_p=0.01
            )
            
            try:
                _, raw = generated.split(prompt, 1)
            except ValueError:
                raw = generated

            raw = raw.strip()
            if raw.upper().startswith("USER:") and "BOT:" in raw:
                raw = raw.split("BOT:", 1)[1].strip()

            answer = truncate(raw)
            answer = evaluate_placeholders(answer)
            answer = answer.replace("-", "\n-").replace("1)", "\n1)").replace("2)", "\n2)").replace("3)", "\n3)").replace("4)", "\n4)").replace("5)", "\n5)").replace("* ", "\n* ").replace("Final", "\nFinal")
            if not "Final" in answer:
                answer = answer.replace("Result", "\nResult")
            print("BOT:", answer, "\n")
            history.append(f"BOT: {answer}")

        except KeyboardInterrupt:
            break

if __name__ == "__main__":
    main()