"""Numeral system induction: recover morpheme values + combination structure
from attested (numeral phrase, integer) pairs, then convert both directions.

Model (covers the large majority of IOL numeral systems):
  value(phrase) = fold over tokens, where adjacent (multiplier, base-power)
  groups combine multiplicatively and groups combine additively — i.e. the
  standard "mixed-radix polynomial" reading:  [2] [20] [3]  ->  2*20 + 3.
Some systems are subtractive or overcounting; a signed-additive variant is
also searched. Token values are solved by constraint search: each distinct
token gets an unknown integer value; attested equations constrain them.

Search is tiny: numeral puzzles use ~5-15 morpheme types with values drawn
from {1..9, base, base^2, ...}. We enumerate candidate value sets per token
from divisors/residues of the attested numbers, then DFS with propagation.
"""

from __future__ import annotations

import re
from collections import Counter
from typing import Dict, List, Optional, Sequence, Tuple

from .preprocess import strip_punct, tokenize

BASES = (10, 20, 5, 12, 60, 4, 6, 8, 15)
MAX_TOKEN_VALUE = 10_000


def _norm_tokens(phrase: str) -> List[str]:
    toks = []
    for t in tokenize(phrase.casefold()):
        t = strip_punct(t)
        # split on hyphens: numeral compounds are often hyphenated
        toks.extend([p for p in re.split(r"[-‑]", t) if p])
    return toks


def _eval(vals: Sequence[int]) -> int:
    """Evaluate token values with the multiplicative-additive convention: a
    smaller value directly before a larger one multiplies it; otherwise
    values add. E.g. [2,20,3] -> 2*20+3 = 43; [3,100,20,7] -> 327."""
    total = 0
    cur = vals[0]
    for prev, v in zip(vals, vals[1:]):
        if v > prev:
            cur = cur * v  # e.g. 2 then 20 -> 40
        else:
            total += cur
            cur = v
    return total + cur


class NumeralSystem:
    def __init__(self, values: Dict[str, int]):
        self.values = dict(values)

    def text_to_num(self, phrase: str) -> Optional[int]:
        toks = _norm_tokens(phrase)
        if not toks or any(t not in self.values for t in toks):
            return None
        return _eval([self.values[t] for t in toks])

    def num_to_text(self, n: int, attested_phrases: Sequence[str]) -> Optional[str]:
        """Generate the phrase for n: enumerate token sequences (up to length
        6) whose evaluation equals n, then pick the one most consistent with
        the attested phrasing style (e.g. do multi-token phrases always give
        a base its explicit multiplier, even 'one'?)."""
        toks = sorted(self.values, key=lambda t: -self.values[t])
        found: List[List[str]] = []
        self._search(n, toks, [], 6, found, limit=16, budget=[100_000])
        if not found:
            return None
        style = _StyleModel(self.values, attested_phrases)
        found.sort(key=lambda seq: (-style.score(seq), len(seq)))
        return " ".join(found[0])

    def _search(self, target: int, toks: List[str], acc: List[str], depth: int,
                found: List[List[str]], limit: int, budget: List[int]) -> None:
        if len(found) >= limit or budget[0] <= 0:
            return
        budget[0] -= 1
        if target == 0 and acc:
            found.append(list(acc))
            return
        if depth == 0 or target <= 0:
            return
        for t in toks:
            v = self.values[t]
            if v > target:
                continue
            # multiplicative: k * v <= target with k attested as token
            for m in toks:
                mv = self.values[m]
                if 1 <= mv < v and mv * v <= target:
                    self._search(target - mv * v, toks, acc + [m, t], depth - 2,
                                 found, limit, budget)
            self._search(target - v, toks, acc + [t], depth - 1, found, limit, budget)


class _StyleModel:
    """Scores a candidate numeral phrase by consistency with attested style:
    (a) are base tokens (value >= 10) given an explicit smaller multiplier in
    attested multi-token phrases? (b) reuse of attested token bigrams."""

    def __init__(self, values: Dict[str, int], phrases: Sequence[str]):
        self.values = values
        self.bigrams = set()
        obs: List[bool] = []
        for ph in phrases:
            toks = _norm_tokens(ph)
            if not toks or any(t not in values for t in toks):
                continue
            self.bigrams.update(zip(toks, toks[1:]))
            if len(toks) < 2:
                continue
            for i, t in enumerate(toks):
                if values[t] >= 10:
                    obs.append(i > 0 and values[toks[i - 1]] < values[t])
        self.prefer_explicit = sum(obs) > len(obs) / 2 if obs else False

    def score(self, seq: Sequence[str]) -> float:
        s = 0.0
        s += 0.5 * sum(1 for bg in zip(seq, seq[1:]) if bg in self.bigrams)
        base_seen: Counter = Counter()
        if len(seq) >= 2:
            for i, t in enumerate(seq):
                if self.values[t] >= 10:
                    base_seen[t] += 1
                    explicit = i > 0 and self.values[seq[i - 1]] < self.values[t]
                    s += 1.0 if explicit == self.prefer_explicit else -1.0
        # positional systems use each base power once; repeats are degenerate
        s -= 2.0 * sum(c - 1 for c in base_seen.values())
        return s


def induce(attested: List[Tuple[str, int]], max_candidates: int = 8) -> Optional[NumeralSystem]:
    """Induce token values from attested (phrase, value) pairs by DFS with
    forward checking. Candidate values per token come from structural
    positions: divisors of attested values, small digits, and base powers."""
    eqs: List[Tuple[List[str], int]] = []
    vocab: List[str] = []
    for phrase, val in attested:
        toks = _norm_tokens(phrase)
        if not toks:
            continue
        eqs.append((toks, val))
        for t in toks:
            if t not in vocab:
                vocab.append(t)
    if not eqs:
        return None

    # Candidate values per token.
    digits = set(range(1, 10))
    base_powers = {b ** k for b in BASES for k in (1, 2, 3) if b ** k <= MAX_TOKEN_VALUE}
    cands: Dict[str, List[int]] = {}
    for t in vocab:
        cs = set(digits) | base_powers
        # a token appearing alone in an equation must equal that value
        for toks, val in eqs:
            if toks == [t]:
                cs = {val}
                break
            if t in toks:
                cs |= {val} | {d for d in _divisors(val) if d <= MAX_TOKEN_VALUE}
        cands[t] = sorted(cs)

    # Constraint propagation on short equations before search: a 1-token
    # equation pins its token; a 2-token equation with one token pinned
    # constrains the other to {V-a, V/a}.
    changed = True
    while changed:
        changed = False
        for toks, val in eqs:
            unknown = [t for t in set(toks) if len(cands[t]) > 1]
            if len(set(toks)) == 1:
                t = toks[0]
                if len(toks) == 1 and cands[t] != [val]:
                    cands[t] = [val]
                    changed = True
            elif len(toks) == 2 and len(unknown) == 1:
                t = unknown[0]
                other = toks[0] if toks[1] == t else toks[1]
                if len(cands[other]) == 1:
                    a = cands[other][0]
                    allowed = {val - a}
                    if a and val % a == 0:
                        allowed.add(val // a)
                    new = [v for v in cands[t] if v in allowed]
                    if new and new != cands[t]:
                        cands[t] = new
                        changed = True

    # Order: most-constrained tokens first.
    order = sorted(vocab, key=lambda t: len(cands[t]))

    assignment: Dict[str, int] = {}
    budget = {"nodes": 200_000}

    def consistent() -> bool:
        for toks, val in eqs:
            if all(t in assignment for t in toks):
                if _eval([assignment[t] for t in toks]) != val:
                    return False
        return True

    def dfs(i: int) -> bool:
        if budget["nodes"] <= 0:
            return False  # search space too big — abstain, don't hang
        if i == len(order):
            return True
        t = order[i]
        for v in cands[t]:
            budget["nodes"] -= 1
            assignment[t] = v
            if consistent() and dfs(i + 1):
                return True
        assignment.pop(t, None)
        return False

    if dfs(0) and budget["nodes"] > 0:
        sys_ = NumeralSystem(assignment)
        # verify every attested equation round-trips
        if all(sys_.text_to_num(p) == v for p, v in attested):
            return sys_
    return None


def _divisors(n: int) -> List[int]:
    n = abs(n)
    out = []
    for d in range(1, int(n ** 0.5) + 1):
        if n % d == 0:
            out += [d, n // d]
    return sorted(set(out))


def extract_attested(pairs) -> List[Tuple[str, int]]:
    """From preprocess Pairs, pull (phrase, int) where one side is a number."""
    out = []
    for p in pairs:
        for a, b in ((p.src, p.tgt), (p.tgt, p.src)):
            bs = b.strip().replace(",", "").replace(" ", "")
            if re.fullmatch(r"\d+", bs):
                out.append((a, int(bs)))
                break
    return out