#!/usr/bin/env python3
"""test.py — Micro Distillery model chat.

Note: the model has no explicit EOS token. Generation is bounded by max_tokens.
BOS (space) is a soft end-of-trajectory signal inside the browser trainer only.
"""
import json, os, random, sys

MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.json")

def load_model(path=MODEL_PATH):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def sample_next(transitions, token, rng, temperature=1.0, top_k=0):
    bucket = transitions.get(token)
    if not bucket: return None
    items = [(k, v) for k, v in bucket.items() if v > 0]
    if not items: return None
    if top_k and len(items) > top_k:
        items = sorted(items, key=lambda x: -x[1])[:top_k]
    tokens = [t for t, _ in items]
    probs = [p for _, p in items]
    if temperature <= 0:
        return max(items, key=lambda x: x[1])[0]
    inv_t = 1.0 / temperature
    weights = [p ** inv_t for p in probs]
    total = sum(weights)
    if total <= 0: return None
    r = rng.random() * total
    acc = 0.0
    for t, w in zip(tokens, weights):
        acc += w
        if r <= acc: return t
    return tokens[-1]

def generate(model, prompt, max_tokens=32, temperature=0.9, top_k=0):
    transitions = model.get("transitions", {})
    seed = int(model.get("harmonic_seed", 1337)) & 0xFFFFFFFF
    rng = random.Random(seed)
    prefix = prompt or ""
    current = prefix[-1] if prefix else " "
    generated = ""
    for _ in range(max_tokens):
        nxt = sample_next(transitions, current, rng, temperature, top_k)
        if nxt is None: break
        generated += nxt
        current = nxt
    return prefix + generated

def main(temperature=0.9):
    print("=== Micro Distillery Chat (toy bigram model) ===")
    if not os.path.exists(MODEL_PATH):
        print("Error: " + MODEL_PATH + " not found.", file=sys.stderr)
        sys.exit(1)
    model = load_model()
    print("Format:      " + str(model.get("format")))
    print("Vocab:       " + str(len(model.get("vocab", []))))
    print("States:      " + str(len(model.get("transitions", {}))))
    print("Temperature: " + str(temperature))
    print("Type /quit to exit.\n")
    while True:
        try: prompt = input("You: ").strip()
        except (EOFError, KeyboardInterrupt): break
        if not prompt or prompt.lower() in ("/quit", "/exit"): break
        print("Bot: " + generate(model, prompt, temperature=temperature))

if __name__ == "__main__":
    temperature = 0.9
    prompts = []
    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == "--temperature" and i + 1 < len(sys.argv):
            temperature = float(sys.argv[i + 1]); i += 2
        else:
            prompts.append(sys.argv[i]); i += 1
    if prompts:
        model = load_model()
        print(generate(model, " ".join(prompts), temperature=temperature))
    else:
        main(temperature)