import os
import random
from tokenizers import Tokenizer
from transformers import AutoTokenizer

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SAMPLE_DIR = os.path.join(SCRIPT_DIR, "sample")

GUTENBERG = os.path.join(SAMPLE_DIR, "gutenberg_sample.txt")
RU_BOOKS  = os.path.join(SAMPLE_DIR, "ru_books_sample.txt")

TOKENIZER_PATH = os.path.join(SCRIPT_DIR, "tokenizer.json")

N = 500
MIN_LEN = 100
MAX_LEN = 2000

arm = Tokenizer.from_file(TOKENIZER_PATH)
gemma = AutoTokenizer.from_pretrained("google/gemma-2-2b", trust_remote_code=True)


def read(path, n):
    if not os.path.exists(path):
        print(f"missing: {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_gemma = 0
arm_wins = 0
gemma_wins = 0
ties = 0

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

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

    total_arm += a
    total_gemma += g

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

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

print("=" * 60)
print(f"ARM:   {total_arm:,}")
print(f"Gemma: {total_gemma:,}")

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

print("=" * 60)
print(f"ARM wins:    {arm_wins}")
print(f"Gemma wins:  {gemma_wins}")
print(f"ties:        {ties}")

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

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