"""Dreamling — raise a from-scratch baby that learns to talk and grows a personality."""
import gradio as gr
import creature as C


def new_creature():
    return C.Creature()


def header_md(c):
    if c is None:
        return "# 🥚\n*a sleeping egg…*"
    e = round(c.energy * 10)
    return (f"# {c.emoji()}   {c.mood_face}\n"
            f"**{c.descriptor()}**  ·  right now feeling **{c.mood_name}** "
            f"{c.mood_face}\n\n"
            f"energy `{'▰' * e}{'▱' * (10 - e)}`  ·  "
            f"brain: **{c.param_count():,}** params ({c.layers()} layers)  ·  "
            f"words: **{len(c.vocab)}**  ·  dreams: **{c.sleeps}**  ·  "
            f"chats: **{c.turns}**  ·  👍 {c.good} 👎 {c.bad}")


def bars_md(c):
    if c is None:
        return ""
    rows = []
    for k, v in c.traits().items():
        f = round(v / 10)
        rows.append(f"`{k:11}` {'█' * f}{'░' * (10 - f)} {round(v)}")
    return "**Personality**\n\n" + "\n\n".join(rows)


def boot():
    c = new_creature()
    greet = C._sample(c.model, C._DEVICE, 24, 1.0) or "dream?"
    chat = [{"role": "assistant", "content": greet}]
    return c, chat, header_md(c), bars_md(c)


def respond(message, history, c):
    c = c or new_creature()
    history = history or []
    if not message.strip():
        return history, "", header_md(c), bars_md(c), c
    reply = c.reply(message)
    history = history + [{"role": "user", "content": message},
                         {"role": "assistant", "content": reply}]
    return history, "", header_md(c), bars_md(c), c


def feedback(good, c):
    if c is not None:
        c.feedback(good)
    return header_md(c), bars_md(c), c


def do_dream(history, c):
    c = c or new_creature()
    note = c.dream()
    history = (history or []) + [{"role": "assistant", "content": note}]
    sample = C._sample(c.model, C._DEVICE, 24, 0.9)
    if sample:
        history.append({"role": "assistant", "content": f"*(I can say…)* {sample}"})
    return history, header_md(c), bars_md(c), c


with gr.Blocks(title="Dreamling") as demo:
    gr.Markdown("## 🌱 Dreamling — raise a baby that learns to talk")
    gr.Markdown(
        "I'm a newborn — at first I can only babble *dream*. **Talk to me** (I'll keep "
        "babbling, but I'm listening), then press **🌙 Dream** and I'll *learn* the words "
        "you said and wake up able to use them. Tell me I'm **good 👍** or **bad 👎** to "
        "shape who I become: kindness → bold and bright 🦋, harshness → withdrawn 🐚, "
        "spoiling → wild 🐉. 🥚→🐛→…"
    )
    state = gr.State()
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=420, label="Dreamling")
            msg = gr.Textbox(placeholder="say something to your Dreamling…", label="Talk",
                             autofocus=True)
            with gr.Row():
                good_btn = gr.Button("👍 good Dreamling", variant="primary")
                bad_btn = gr.Button("👎 bad Dreamling", variant="secondary")
            dream_btn = gr.Button("🌙 Dream (sleep & learn the words you said)")
        with gr.Column(scale=2):
            head = gr.Markdown()
            personality = gr.Markdown()

    demo.load(boot, None, [state, chatbot, head, personality])
    msg.submit(respond, [msg, chatbot, state], [chatbot, msg, head, personality, state])
    good_btn.click(lambda c: feedback(True, c), state, [head, personality, state])
    bad_btn.click(lambda c: feedback(False, c), state, [head, personality, state])
    dream_btn.click(do_dream, [chatbot, state], [chatbot, head, personality, state])


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())