"""run point-1. usage: python3 generate.py ["prompt"] [--temp T] [--n N] [--seed S]"""
import json, math, random, struct, sys

BYTES_VOCAB = [10] + list(range(32, 127))
H = {b: max(-1.0, min(1.0, (b - 79) / 47.0)) for b in BYTES_VOCAB}


def load_w(path="model.safetensors"):
    with open(path, "rb") as f:
        (hlen,) = struct.unpack("<Q", f.read(8))
        header = json.loads(f.read(hlen))
        f.seek(8 + hlen)
        (w,) = struct.unpack("<f", f.read(4))
    nparams = sum(math.prod(t["shape"]) for k, t in header.items() if not k.startswith("__"))
    assert nparams == 1, nparams
    return w


def main():
    args = sys.argv[1:]
    prompt = next((a for a in args if not a.startswith("-")), "")
    temp = float(args[args.index("--temp") + 1]) if "--temp" in args else 0.8
    count = int(args[args.index("--n") + 1]) if "--n" in args else 200
    seed = int(args[args.index("--seed") + 1]) if "--seed" in args else 42
    random.seed(seed)

    w = load_w()
    print(f"point-1 | parameters: 1 | w = {w:+.6f}")
    if prompt:
        print(f"prompt  : {prompt!r}")
    weights = [math.exp(w * H[b] / temp) for b in BYTES_VOCAB]
    out = "".join(chr(random.choices(BYTES_VOCAB, weights=weights)[0]) for _ in range(count))
    print("-" * 52)
    print(out)


if __name__ == "__main__":
    main()