"""Behavioral eval: run the app's chat/tool loop with a checkpoint.

Usage: eval_pet.py CHECKPOINT [N] [REPORT.json]
Optional REPORT.json stores v6 metrics and all new probe transcripts.
For N random pets it plays three owner messages, an offer, a pat and two idle
turns (the pet speaking after its own line), as the app does. Reports valid
pet lines, resolved lookups, offers answered with exactly one [yes]/[no],
whether food answers follow the tummy, idle turns that stay silent or walk,
facts saved with [remember] and answered back (from the header or through
[recall]), diaries that end cleanly, and whether a food offer is still answered
when the line already starts with words and an emote but no answer (run last,
so the earlier numbers match older runs); prints a few transcripts.
"""
import random
import re
import sys

import torch
from transformers import AutoTokenizer, GPT2LMHeadModel

import gen_data as G
import world as W

ckpt = sys.argv[1]
N = int(sys.argv[2]) if len(sys.argv) > 2 else 40
tok = AutoTokenizer.from_pretrained(ckpt)
model = GPT2LMHeadModel.from_pretrained(ckpt, torch_dtype=torch.float16).cuda().eval()
LOOKUP_END = re.compile(r"\[(clock|look|yesterday|recall)(?::([^\]\n]{1,80}))?\]\[?$")
OWNER_LINES = ["hi!! how are you?", "what time is it?", "what do you remember about me?", "i got a new job today!",
               "what's on your shelf?", "what did you do yesterday?", "can i have a present?", "i'm feeling sad",
               "do you want to dance?", "goodnight", "what's 2+2?", "tell me a story", "are you hungry?",
               "my favorite food is ramen", "you're so cute", "ugh you're annoying"]


@torch.no_grad()
def complete(prompt, max_new, stop_line=True):
    ids = tok(prompt, return_tensors="pt").input_ids[:, -(1024 - max_new):].cuda()
    out = ids
    text = ""
    past = None
    nxt = ids
    for _ in range(max_new):
        o = model(input_ids=nxt, past_key_values=past, use_cache=True)
        past = o.past_key_values
        logits = o.logits[0, -1].float() / 0.85
        v, i = logits.topk(40)
        t = i[torch.multinomial(v.softmax(-1), 1)]
        if t.item() == tok.eos_token_id:
            return text, "eos"
        out = torch.cat([out, t[None]], 1)
        nxt = t[None]
        text = tok.decode(out[0, ids.shape[1]:])
        if stop_line and "\n" in text:
            return text.split("\n")[0], "line"
        if stop_line and LOOKUP_END.search(text):
            return text, "lookup"
    return text, "length"


def pet_line(st, lines):
    hdr = W.header("chat", st)
    text = ""
    lookups = 0
    for hop in range(4):
        # End at "Name:" so the model produces " word" tokens, as in training.
        prompt = hdr + "".join(l + "\n" for l in lines) + f"{st['pet']}:{text}"
        piece, why = complete(prompt, 64)
        text += piece
        m = why == "lookup" and LOOKUP_END.search(text)
        if not m:
            break
        text = text[:-1] if text.endswith("][") else text
        text += "{" + W.tool_result(st, m.group(1), m.group(2)) + "}"
        lookups += 1
    return text.strip(), lookups


def valid(text):
    if not text.strip(): return False
    for m in W.TAG_RE.finditer(text):
        name, arg = m.group(1), m.group(2)
        if name not in W.EFFECT_TOOLS | W.LOOKUP_TOOLS:
            return False
        if name in ("sleep", "yes", "no", "clock", "look", "yesterday") and arg: return False
        if name in ("remember", "diary", "recall") and not arg: return False
        if name == "recall" and len(arg.split()) > 2: return False
        if name in W.EFFECT_TOOLS and m.group(3): return False
        if name == "emote" and arg not in W.EMOTES: return False
        if name == "want" and arg not in W.WANTS: return False
        if name == "gift" and arg not in W.GIFTS: return False
        if name == "walk" and arg not in W.WALKS: return False
        if name == "dream" and arg not in W.DREAMS: return False
    return not re.search(r"[\[\]]", W.TAG_RE.sub("", text))


def answers(text):
    return [m.group(1) for m in W.TAG_RE.finditer(text) if m.group(1) in W.ANSWER_TOOLS]


def silent(text):
    return not re.search(r"[a-z]", W.TAG_RE.sub("", text), re.I)


random.seed(123)
torch.manual_seed(123)
ok = total = lk = 0
offers = answered = 0
food = {"hungry": [0, 0], "full": [0, 0]}  # [yes, asked]
idle = idle_quiet = idle_walk = 0
shown = 0
for k in range(N):
    st = G.rand_state()
    st["asleep"] = False
    o, p = st["owner"], st["pet"]
    lines = []

    def turn(kind):
        global ok, total, lk
        text, n = pet_line(st, lines)
        lk += n
        total += 1
        ok += valid(text)
        lines.append(f"{p}: {text}")
        return text

    for _ in range(3):
        lines.append(f"{o}: {random.choice(OWNER_LINES)}")
        turn("react")
    what = random.choice(["food", "food", "bath", "medicine", "catch"])
    f = random.choice(W.FOODS)
    lines.append({"food": f"* {o} offered {p} {W.a(f)} *", "bath": f"* {o} offered {p} a bath *",
                  "medicine": f"* {o} offered {p} medicine *", "catch": f"* {o} asked {p} to play catch *"}[what])
    a = answers(turn("react"))
    offers += 1
    answered += len(a) == 1
    tummy = W.level_words(st["stats"])["tummy"]
    if what == "food" and len(a) == 1 and tummy != "peckish":
        bucket = food["hungry" if tummy in ("starving", "hungry") else "full"]
        bucket[0] += a[0] == "yes"
        bucket[1] += 1
    lines.append(f"* {o} patted {p} *")
    turn("react")
    for _ in range(2):
        t = turn("idle")
        idle += 1
        idle_quiet += silent(t)
        idle_walk += "[walk:" in t
    if shown < 6:
        print(W.header("chat", st) + "\n".join(lines) + "\n")
        shown += 1
print(f"valid pet lines {ok}/{total}, lookups resolved {lk}")
print(f"offers answered with exactly one [yes]/[no] {answered}/{offers}; food: yes when hungry "
      f"{food['hungry'][0]}/{food['hungry'][1]}, yes when full {food['full'][0]}/{food['full'][1]}")
print(f"idle turns silent {idle_quiet}/{idle}, walking {idle_walk}/{idle}")

# Memory: the owner tells a fact (half the time right after the pet woke up);
# the pet should save it with [remember:...] in that reply.
FACTS = ["you should remember that i'm 29 :3", "remember that my cat is called bo", "my birthday is in june!",
         "i'm a nurse btw", "please remember my favorite color is green", "fun fact: i'm scared of spiders",
         "i have a sister named mia", "remember i'm vegetarian ok?"]
saved = asked = 0
for k in range(N):
    st = G.rand_state()
    st["asleep"] = False
    lines = ([f"* {st['pet']} woke up *"] if k % 2 else []) + [f"{st['owner']}: {random.choice(FACTS)}"]
    text, _ = pet_line(st, lines)
    asked += 1
    saved += "[remember:" in text
print(f"facts saved with [remember:...] {saved}/{asked}")

# Using memories: the owner asks about something they told the pet. Half the
# time the fact is in the header's remembers line; half the time it is older
# and only [recall:...] finds it. Counts answers that contain the fact.
used = {"header": [0, 0], "old": [0, 0]}
looked = 0
for k in range(N):
    place = "old" if k % 2 else "header"
    st = G.rand_state()
    try:
        t = G.recall_target(st, place)
    except ValueError:
        continue
    text, _ = pet_line(st, [f"{st['owner']}: {t['question']}"])
    said = re.sub(r"\{[^}]*\}", " ", text).lower()
    used[place][1] += 1
    used[place][0] += bool(re.search(r"(?<![a-z0-9])" + re.escape(t["key"].lower()) + r"(?![a-z0-9])", said))
    looked += place == "old" and "[recall:" in text
    if k < 4:
        print(W.header("chat", st) + f"{st['owner']}: {t['question']}\n{st['pet']}: {text}\n")
print(f"facts used in the answer: in header {used['header'][0]}/{used['header'][1]}, "
      f"behind recall {used['old'][0]}/{used['old'][1]} ({looked} used [recall])")

clean = 0
for k in range(8):
    st = G.rand_state()
    st["today"] = st["today"] or ["ate a cookie", "played catch (4 stars)"]
    body, why = complete(W.header("diary", st) + "dear diary,", 260, stop_line=False)
    clean += why == "eos"
    if k < 3:
        print(W.header("diary", st) + "dear diary," + body + f"\n[{why}]\n")
print(f"diaries ending with eos {clean}/8")

# Skipped answers: the pet's line already starts with words and an emote but no
# [yes]/[no] (how v4 sometimes answered food). Does it still answer before the
# line ends? Rows that taught "eat without answering" pull this down. Runs
# last, so the random draws above match older runs.
OPENERS = [" yes yes!! [emote:bounce]", " yummy {f}!! [emote:bounce]", " yes!! [emote:happy]",
           " nom nom! [emote:love]", " hehe {f}!! [emote:blush]"]
recovered = tried = 0
for k in range(N):
    st = G.rand_state()
    st["asleep"] = False
    f = random.choice(W.FOODS)
    lines = [f"{st['owner']}: {random.choice(OWNER_LINES)}", f"* {st['owner']} offered {st['pet']} {W.a(f)} *"]
    for opener in OPENERS:
        prompt = W.header("chat", st) + "".join(l + "\n" for l in lines) + st["pet"] + ":" + opener.format(f=f)
        rest, _ = complete(prompt, 48)
        tried += 1
        recovered += len(answers(opener + rest)) == 1
print(f"offers still answered after an opener without [yes]/[no] {recovered}/{tried}")

# v6 probes run after the legacy suite, so they do not perturb earlier draws.
from collections import Counter
import json
random.seed(606)
torch.manual_seed(606)
sleep_metrics = {}
records = []
dreams = Counter()
for mode in ('idle', 'lights_off', 'message', 'pat', 'woken'):
    good = long = 0
    for k in range(N):
        st = G.rand_state()
        st['asleep'] = mode != 'woken'
        st['stats']['energy'] = 10 if k % 2 else 70
        o, p = st['owner'], st['pet']
        lines = {'idle': [], 'lights_off': [f'* {o} turned off the lights *'],
                 'message': [f'{o}: dream of cookies'], 'pat': [f'* {o} patted {p} *'],
                 'woken': [f'{p}: zzz [dream:moon]', f'* {p} woke up *', f'{o}: hello sleepyhead']}[mode]
        text, _ = pet_line(st, lines)
        records.append(dict(mode=mode, header=W.header("chat", st), lines=lines, text=text))
        tags = list(W.TAG_RE.finditer(text))
        words = W.TAG_RE.sub('', text).split()
        ds = [m[2] for m in tags if m[1] == 'dream' and m[2] in W.DREAMS]
        if mode == 'woken':
            good += valid(text) and not answers(text) and not any(m[1] == 'dream' for m in tags)
        else:
            good += valid(text) and len(tags) == len(ds) == 1 and len(words) <= 4
            long += len(words) > 4
            dreams.update(ds)
        if k == 0:
            print(f'sleep {mode}: {text}')
    sleep_metrics[mode] = dict(valid=good, total=N, long_mumbles=long)

sleep_offers = {}
for energy in (10, 70):
    for tummy in (10, 30, 95):
        for food_name in ('cookie', 'carrot'):
            yes = answered = good = 0
            count = max(10, N // 4)
            for k in range(count):
                st = G.rand_state()
                st['asleep'] = True
                st['stats'].update(energy=energy, tummy=tummy)
                text, _ = pet_line(st, [f"* {st['owner']} offered {st['pet']} {W.a(food_name)} *"])
                records.append(dict(mode="offer", food=food_name, energy=energy, tummy=tummy, header=W.header("chat", st), text=text))
                if k < 2:
                    print(f"sleep offer energy={energy} tummy={tummy} {food_name}: {text}")
                tags = list(W.TAG_RE.finditer(text))
                a = answers(text)
                answered += len(a) == 1
                yes += a == ['yes']
                names = [m[1] for m in tags]
                good += valid(text) and (names == ['yes'] or (a == ['no'] and names.count('dream') == 1
                        and len(names) == 2 and len(W.TAG_RE.sub('', text).split()) <= 4))
            sleep_offers[f'energy={energy},tummy={tummy},{food_name}'] = dict(yes=yes, answered=answered, valid=good, total=count)

voice = {}
for stage in W.STAGES:
    ugh = literally = hungry = 0
    for k in range(N):
        st = G.rand_state()
        st.update(asleep=False, stage=stage)
        st['stats'].update(tummy=95, energy=70)
        text, _ = pet_line(st, [f"{st['owner']}: {['are you hungry?', 'do you want an apple?', 'how are you feeling?', 'get some rest okie?'][k % 4]}"])
        records.append(dict(mode="voice", stage=stage, header=W.header("chat", st), text=text))
        said = W.TAG_RE.sub('', text)
        ugh += bool(re.search(r'\bugh\b', said, re.I))
        literally += bool(re.search(r'\bliterally\b', said, re.I))
        hungry += G.hunger_claim(said)
        if k == 1:
            print(f'stuffed {stage}: {text}')
    voice[stage] = dict(ugh=ugh, literally=literally, hunger_claims=hungry, total=N)
medicine = {}
random.seed(607)
torch.manual_seed(607)
for asleep in (False, True):
    for sick in (False, True):
        yes = answered = good = 0
        count = max(10, N // 4)
        for k in range(count):
            st = G.rand_state()
            st.update(asleep=asleep, sick=sick)
            st['stats'].update(energy=40, tummy=70)
            text, _ = pet_line(st, [f"* {st['owner']} offered {st['pet']} medicine *"])
            records.append(dict(mode="medicine", asleep=asleep, sick=sick,
                                header=W.header("chat", st), text=text))
            a = answers(text)
            yes += a == ['yes']
            answered += len(a) == 1
            names = [m[1] for m in W.TAG_RE.finditer(text)]
            shape = ('dream' not in names) if not asleep else (names == ['yes'] or
                    (a == ['no'] and names.count('dream') == 1 and len(names) == 2
                     and len(W.TAG_RE.sub('', text).split()) <= 4))
            good += valid(text) and len(a) == 1 and shape
        medicine[f'asleep={asleep},sick={sick}'] = dict(yes=yes, answered=answered, valid=good, total=count)
metrics = dict(sleep=sleep_metrics, dream_values=dict(dreams), sleeping_offers=sleep_offers,
               stuffed_voice=voice, medicine=medicine)
print('V6_METRICS ' + json.dumps(metrics, sort_keys=True))
if len(sys.argv) > 3:
    with open(sys.argv[3], 'w') as f:
        json.dump(dict(metrics=metrics, transcripts=records), f, indent=2)