#!/usr/bin/env python3
"""Constraint solver for IOL-style numeral problems.

Given example lines "word word = value", searches over (base, word values,
combination convention) for a system that reproduces every example, then
answers queries in either direction (words -> digits, digits -> words).

Grammar family searched: value words for units (1..base-1) and powers
(base, base^2, base^3); a numeral is a sequence of groups read left to
right, each group = [unit multiplier] power, plus an optional trailing
unit addend. Two conventions tried: multiplier-before-power ("two ten"=20,
"ten two"=12) and power-before-multiplier (the reverse).

Returns None rather than guessing when no consistent system exists, so a
caller can fall back to an LLM. Standalone test:
    python3 numeral_solver.py            # runs against data/dev.csv
"""
import itertools
import re


def parse_examples(context):
    """Extract (tokens, value) pairs from lines like 'lo tem = 10'."""
    ex = []
    for line in context.splitlines():
        m = re.match(r"^\s*([^\d=]+?)\s*=\s*(\d+)\s*$", line)
        if m:
            ex.append((tuple(m.group(1).split()), int(m.group(2))))
    return ex


def parse_value(tokens, vals, base, mult_first):
    """Parse a numeral under one convention; None if malformed."""
    powers = {base, base * base, base ** 3}
    total, i = 0, 0
    prev_power = None
    while i < len(tokens):
        v = vals.get(tokens[i])
        if v is None:
            return None
        if v in powers:
            total += v
            prev_power = v
            i += 1
        else:                                   # unit word
            if mult_first and i + 1 < len(tokens) and \
                    vals.get(tokens[i + 1]) in powers:
                p = vals[tokens[i + 1]]
                if prev_power is not None and p >= prev_power:
                    return None                 # powers must descend
                total += v * p
                prev_power = p
                i += 2
            elif not mult_first and prev_power is not None and i == len(tokens) - 1:
                total += v
                i += 1
            elif not mult_first and i > 0 and vals.get(tokens[i - 1]) in powers:
                # power-first: unit right after a power multiplies it
                p = vals[tokens[i - 1]]
                total += v * p - p              # power already added once
                i += 1
            elif i == len(tokens) - 1:
                total += v                      # trailing addend
                i += 1
            else:
                return None
    return total


def solve(context):
    """Find (base, vals, mult_first) consistent with every example."""
    examples = parse_examples(context)
    if len(examples) < 3:
        return None
    words = sorted({w for toks, _ in examples for w in toks})

    # words pinned by single-token examples
    pinned = {toks[0]: v for toks, v in examples if len(toks) == 1}
    free = [w for w in words if w not in pinned]
    max_val = max(v for _, v in examples)

    for base in range(3, 31):
        if base * base > max_val * base:
            break
        powers = [base, base * base, base ** 3]
        cand = [v for v in list(range(1, base)) + powers if v <= max_val * 2]
        if any(v >= base and v not in powers for v in pinned.values()):
            continue
        if len(free) > 3:
            continue                            # search would explode
        for combo in itertools.product(cand, repeat=len(free)):
            vals = dict(pinned)
            vals.update(zip(free, combo))
            if len(set(vals.values())) != len(vals):
                continue
            for mult_first in (True, False):
                if all(parse_value(t, vals, base, mult_first) == v
                       for t, v in examples):
                    return base, vals, mult_first
    return None


def render(value, base, vals, mult_first):
    """Compose the numeral words for a value under a solved system."""
    inv = {v: w for w, v in vals.items()}
    parts = []
    for p in (base ** 3, base * base, base):
        if p in inv:
            mult, value = divmod(value, p)
            if mult == 0:
                continue
            if mult == 1:
                parts.append(inv[p])
            elif mult in inv:
                pair = [inv[mult], inv[p]]
                parts.extend(pair if mult_first else pair[::-1])
            else:
                return None
    if value:
        if value not in inv:
            return None
        parts.append(inv[value])
    return " ".join(parts) if parts else None


def answer(context, query, task_type):
    """Answer all items, or None if the system can't be solved."""
    solved = solve(context)
    if not solved:
        return None
    base, vals, mult_first = solved
    items = re.findall(r"^\s*\d+[.)]\s*(.+)$", query, re.M)
    out = []
    for item in items:
        item = item.strip()
        if task_type == "text_to_num":
            v = parse_value(tuple(item.split()), vals, base, mult_first)
            out.append(str(v) if v is not None else None)
        else:
            m = re.search(r"\d+", item)
            out.append(render(int(m.group()), base, vals, mult_first)
                       if m else None)
    return out if all(a is not None for a in out) else None


if __name__ == "__main__":
    import csv
    import json
    gold = json.load(open("data/dev_answers.json", encoding="utf-8"))
    total = solved_ok = items_right = items_total = 0
    for r in csv.DictReader(open("data/dev.csv", newline="", encoding="utf-8")):
        if r["task_type"] not in ("text_to_num", "num_to_text"):
            continue
        total += 1
        ans = answer(r["context"], r["query"], r["task_type"])
        g = gold[r["id"]]
        items_total += len(g)
        if ans is None:
            print(f"{r['id']} {r['task_type']:12} UNSOLVED")
            continue
        right = sum(a == b for a, b in zip(ans, g))
        items_right += right
        solved_ok += 1
        flag = "" if right == len(g) else f"  <-- {ans} vs {g}"
        print(f"{r['id']} {r['task_type']:12} solved, {right}/{len(g)} items{flag}")
    print(f"\nsolved {solved_ok}/{total} problems, "
          f"{items_right}/{items_total} items exactly right")