import os
import random
from tokenizers import Tokenizer
from transformers import GPT2TokenizerFast

GUTENBERG = "sample/gutenberg_sample.txt"
RU_BOOKS  = "sample/ru_books_sample.txt"

N = 500
MIN_LEN = 100
MAX_LEN = 2000

arm = Tokenizer.from_file("tokenizer.json")
gpt2 = GPT2TokenizerFast.from_pretrained("gpt2")


def read(path, n):
    if not os.path.exists(path):
        return []
    s = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if MIN_LEN <= len(line) <= MAX_LEN:
                s.append(line)
    random.shuffle(s)
    return s[:n]


random.seed(42)

en = read(GUTENBERG, N)
ru = read(RU_BOOKS, N)

samples = [(t, "en") for t in en] + [(t, "ru") for t in ru]

print(f"en: {len(en)}")
print(f"ru: {len(ru)}")
print(f"total: {len(samples)}")

total_arm = 0
total_gpt2 = 0
arm_wins = 0
gpt2_wins = 0
ties = 0

by_lang = {
    "en": {"arm": 0, "gpt2": 0, "chars": 0, "n": 0},
    "ru": {"arm": 0, "gpt2": 0, "chars": 0, "n": 0},
}

for text, lang in samples:
    a = len(arm.encode(text).ids)
    g = len(gpt2.encode(text))

    total_arm += a
    total_gpt2 += g

    by_lang[lang]["arm"] += a
    by_lang[lang]["gpt2"] += g
    by_lang[lang]["chars"] += len(text)
    by_lang[lang]["n"] += 1

    if a < g:
        arm_wins += 1
    elif a > g:
        gpt2_wins += 1
    else:
        ties += 1

print("=" * 60)
print(f"ARM-300: {total_arm:,}")
print(f"GPT-2:   {total_gpt2:,}")

print("=" * 60)
for lang, d in by_lang.items():
    if d["n"] == 0:
        continue
    ar = d["arm"] / d["chars"]
    gr = d["gpt2"] / d["chars"]
    diff = d["arm"] - d["gpt2"]
    print(f"{lang}: ARM={d['arm']:,} GPT2={d['gpt2']:,} tok/char ARM={ar:.4f} GPT2={gr:.4f} delta={diff:+,}")

print("=" * 60)
print(f"ARM wins:   {arm_wins}")
print(f"GPT2 wins:  {gpt2_wins}")
print(f"ties:       {ties}")

diff = total_arm - total_gpt2
pct = abs(diff) / max(total_arm, total_gpt2) * 100

print("=" * 60)
if diff < 0:
    print(f"WINNER: ARM-300 by {pct:.2f}%")
elif diff > 0:
    print(f"WINNER: GPT-2 by {pct:.2f}%")
else:
    print("WINNER: TIE")