"""Chat with WPLM, the chatbot whose weights were programmed directly (no training).

  python chat.py                      interactive (a fresh random seed per reply)
  python chat.py "tell me a joke"     one reply (seed 0 unless --seed)

Needs internet: WPLM asks the infini-gram API (api.infini-gram.io) for words. Replies take a few seconds; the client
keeps under 40 requests a minute (WPLM_RPM). API answers are cached in ~/.cache/wplm (WPLM_CACHE).
"""
import argparse
import os
import sys
import time

os.environ.setdefault("OMP_NUM_THREADS", "2")          # idle BLAS threads spin while the host waits for the API
os.environ.setdefault("OPENBLAS_NUM_THREADS", "2")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wplm.load import load                              # noqa: E402
from wplm.host import SimHost                           # noqa: E402


def main():
    ap = argparse.ArgumentParser(description="Chat with WPLM (programmed weights + the infini-gram API).")
    ap.add_argument("message", nargs="*", help="one message; omit for an interactive chat")
    ap.add_argument("--seed", type=int, default=None, help="dice / sampling seed (default: 0 for one message, time for chat)")
    ap.add_argument("-v", "--verbose", action="store_true", help="print the lookups the network asked for")
    a = ap.parse_args()
    host = SimHost(load())
    if a.message:
        host.seed = 0 if a.seed is None else a.seed
        text, info = host.reply(" ".join(a.message))
        print(text)
        if a.verbose:
            print(info, file=sys.stderr)
        return
    print("WPLM (programmed weights, no training). Ctrl-D to quit.")
    while True:
        try:
            line = input("> ")
        except EOFError:
            break
        if not line.strip():
            continue
        host.seed = int(time.time()) if a.seed is None else a.seed
        text, info = host.reply(line)
        print("WPLM:", text)
        if a.verbose:
            print("   ", info)


if __name__ == "__main__":
    main()