"""train point-1 on tinyshakespeare.

fits logits[b] = w * h[b] to the byte histogram. the loss is convex in one
variable so newton's method gets the exact optimum. writes model.safetensors,
config.json, report.txt.
"""
import json, math, os, struct, urllib.request

CORPUS_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
CORPUS_LOCAL = "corpus.txt"
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 get_corpus():
    if not os.path.exists(CORPUS_LOCAL):
        urllib.request.urlretrieve(CORPUS_URL, CORPUS_LOCAL)
    return CORPUS_LOCAL


def histogram(path):
    counts = dict.fromkeys(BYTES_VOCAB, 0)
    total = 0
    with open(path, "rb") as f:
        while chunk := f.read(65536):
            for b in chunk:
                if b in counts:
                    counts[b] += 1
            total += len(chunk)
    return counts, total


def write_safetensors(w):
    payload = struct.pack("<f", w)
    header = {"w": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}
    hb = json.dumps(header, separators=(",", ":")).encode()
    with open("model.safetensors", "wb") as f:
        f.write(struct.pack("<Q", len(hb)) + hb + payload)


def main():
    counts, total = histogram(get_corpus())
    n = float(total)

    uniform_loss = math.log(len(BYTES_VOCAB))
    floor = -sum((c / n) * math.log(c / n) for c in counts.values() if c)

    ch = sum(counts[b] * H[b] for b in BYTES_VOCAB) / n

    # newton on L(w) = -(sum c_b (w h_b - logZ)) / n
    w = 0.0
    log = []
    for step in range(60):
        ex = [math.exp(w * H[b]) for b in BYTES_VOCAB]
        z = sum(ex)
        eh = sum(e * H[b] for e, b in zip(ex, BYTES_VOCAB)) / z
        var = sum(e * H[b] ** 2 for e, b in zip(ex, BYTES_VOCAB)) / z - eh * eh
        loss = -(sum(counts[b] * (w * H[b] - math.log(z)) for b in BYTES_VOCAB)) / n
        log.append((step, w, loss))
        if abs(eh - ch) < 1e-12:
            break
        w -= (eh - ch) / max(var, 1e-12)

    write_safetensors(w)

    cfg = {
        "architectures": ["Point1ForCausalLM"],
        "model_type": "point1",
        "vocab_size": len(BYTES_VOCAB),
        "byte_feature": "(b - 79) / 47 clipped to [-1, 1]",
        "auto_map": {
            "AutoConfig": "modeling_point1.Point1Config",
            "AutoModelForCausalLM": "modeling_point1.Point1ForCausalLM",
        },
    }
    with open("config.json", "w") as f:
        json.dump(cfg, f, indent=2)

    psum = sum(math.exp(w * H[b]) for b in BYTES_VOCAB)
    ranked = sorted(
        ((math.exp(w * H[b]) / psum, chr(b) if b != 10 else "\\n") for b in BYTES_VOCAB),
        reverse=True,
    )

    lines = [
        "POINT-1 TRAINING REPORT",
        "=" * 52,
        f"corpus       tinyshakespeare, {total:,} bytes",
        "parameters   1",
        "",
    ]
    lines += [f"newton {s:>3}  w = {wi:+.6f}  loss = {li:.4f}" for s, wi, li in log]
    lines += [
        "",
        f"uniform      {uniform_loss:.4f} nats/byte",
        f"unigram floor {floor:.4f}",
        f"point-1      {loss:.4f}",
        "",
        "top tokens:",
    ]
    lines += [f"  '{t}'  {p:.4f}" for p, t in ranked[:8]]
    report = "\n".join(lines)
    print(report)
    with open("report.txt", "w") as f:
        f.write(report + "\n")


if __name__ == "__main__":
    main()