import random, re
from collections import Counter
from moe import moe_generate   # reuses loaded base + adapters + router

random.seed(12345)  # disjoint from training seeds

def int_to_roman(n):
    vals = [(1000,"M"),(900,"CM"),(500,"D"),(400,"CD"),(100,"C"),(90,"XC"),
            (50,"L"),(40,"XL"),(10,"X"),(9,"IX"),(5,"V"),(4,"IV"),(1,"I")]
    out = ""
    for v, s in vals:
        while n >= v:
            out += s; n -= v
    return out

def parse_answer(text):
    m = re.findall(r"answer is\s*:?\s*([A-Za-z0-9,\.]+)", text, re.IGNORECASE)
    tok = m[-1] if m else (text.strip().split()[-1] if text.strip() else "")
    return tok.strip().rstrip(".")

def norm_num(s):
    try: return float(s.replace(",", ""))
    except: return None

def norm_roman(s):
    return re.sub(r"[^IVXLCDM]", "", s.upper())

# build eval sets
mult_items, roman_items = [], []
for _ in range(100):
    a, b = random.randint(10,99), random.randint(10,99)
    mult_items.append((f"What is {a} times {b}?", a*b))
for _ in range(100):
    n = random.randint(1,3999)
    roman_items.append((f"Convert {n} to Roman numerals", int_to_roman(n)))

# run
route_correct = 0
routes_seen = Counter()
mult_hits = roman_hits = 0

def run(items, gold_route, kind):
    global route_correct, mult_hits, roman_hits
    hits = 0
    for i, (q, gold) in enumerate(items, 1):
        expert, ans = moe_generate(q)
        routes_seen[expert] += 1
        if expert == gold_route: route_correct += 1
        pred = parse_answer(ans)
        if kind == "mult":
            ok = norm_num(pred) == float(gold)
        else:
            ok = norm_roman(pred) == gold
        hits += ok
        print(f"  [{kind} {i:3d}/100] route={expert:5s} gold={gold!s:8s} pred={pred!s:8s} {'OK' if ok else 'X'}")
    return hits

print("=== MULTIPLICATION ===")
mult_hits  = run(mult_items,  "mult",  "mult")
print("=== ROMAN ===")
roman_hits = run(roman_items, "roman", "roman")

print("\n================ RESULTS ================")
print(f"Routing accuracy:    {route_correct}/200 = {route_correct/200:.3f}")
print(f"  route distribution: {dict(routes_seen)}")
print(f"Mult  end-to-end acc: {mult_hits}/100 = {mult_hits/100:.3f}   (standalone ~0.94)")
print(f"Roman end-to-end acc: {roman_hits}/100 = {roman_hits/100:.3f}   (standalone ~0.24)")
print("=========================================")