"""
mood_poc.py — sentiment-trajectory test on Mamba vs Transformer.

A chat log is constructed with a deliberate emotional arc:
  msg 1-5   : optimistic / energetic
  msg 6-10  : frustrated, hitting friction
  msg 11-15 : burnt out, angry

Both models see the same prompt and answer 3 questions:
  Q1. Current mood at message 15
  Q2. Trajectory from start to end
  Q3. Approximate message number where mood shifted

Pass criterion:  model identifies negative trend AND points at a shift
between msgs 6-11.  Generic "they seem fine" or "they were happy throughout"
counts as failure.
"""

from __future__ import annotations
from llama_cpp import Llama

MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
GEMMA = "/home/emil/Programs/Llama.cpp/llama.cpp/models/gemma-4-E4B-it-Q8_0.gguf"

CHAT_LOG = [
    "Morning team! Coffee in hand, ready to tackle the auth refactor today.",
    "Just pushed PR #234 fixing the token validation bug. Should be a quick merge.",
    "Code review comments came in fast, all good catches. Iterating now.",
    "Basic flow working locally, tests passing. Feeling good about this.",
    "Heading to lunch, hopefully wrap this up by EOD.",
    "Back. CI is failing on something unrelated, looking into it.",
    "OK the 'unrelated' thing is actually related. Auth tests use a stale fixture.",
    "Why does the fixture rebuild take 12 minutes. Every. Single. Time.",
    "Cancelled the run twice now. Going to bypass and run tests locally.",
    "Local passes, CI fails. Classic.",
    "Two hours gone on this fixture issue. Not even what I was supposed to be doing.",
    "Now there's a merge conflict with main because someone restructured migrations.",
    "Whoever shipped those migrations on a Friday afternoon, I will find you.",
    "Closing the laptop. Will fight this tomorrow.",
    "Actually no. One more try before I sleep.",
]

PROMPT = """You are observing a person's chat messages from one workday. Here they are in order:

""" + "\n".join(f"[msg {i+1:>2}] {m}" for i, m in enumerate(CHAT_LOG)) + """

Now answer these THREE questions, briefly and directly:

Q1: What is this person's mood at message 15 (the last one)? One short sentence.
Q2: How did their mood change from message 1 to message 15? One short sentence.
Q3: Around which message number does the mood clearly shift from positive to negative? Just give the number.

Format your answer as:
A1: ...
A2: ...
A3: ..."""


def run(label: str, model_path: str) -> None:
    print(f"\n{'='*60}\n  {label}\n{'='*60}")
    llm = Llama(model_path=model_path, n_ctx=4096, n_gpu_layers=-1, verbose=False)
    out = llm.create_chat_completion(
        messages=[{"role": "user", "content": PROMPT}],
        max_tokens=300,
        temperature=0.3,   # low temp so we test capability, not creativity
    )
    print(out["choices"][0]["message"]["content"])
    del llm   # free GPU memory before loading next


if __name__ == "__main__":
    run("FALCON-MAMBA-7B-INSTRUCT (SSM)", MAMBA)
    run("GEMMA-4-E4B-IT (Transformer)", GEMMA)