"""Ways of writing a quantity in a question, and parsing numbers back out of text."""

from __future__ import annotations

import math
import random
from dataclasses import dataclass

import regex

from zenith.engines.constants import AU, M_EARTH, M_SUN, MPC, LY


@dataclass(frozen=True)
class Expr:
    text: str
    value: float  # the canonical value in the engine unit


_SUPERSCRIPT = str.maketrans("⁰¹²³⁴⁵⁶⁷⁸⁹⁻", "0123456789-")
_SCALES = {"thousand": 1e3, "million": 1e6, "billion": 1e9, "trillion": 1e12}
_NUM = r"[-−]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?"
_SCI = r"(?:[eE][-+]?\d+|\s?[×x*]\s?10(?:\^[-+]?\d+|[⁰¹²³⁴⁵⁶⁷⁸⁹⁻]+))"
NUMBER_RE = regex.compile(rf"(?<![\w.]){_NUM}{_SCI}?(?:\s(?:thousand|million|billion|trillion)\b)?(?![\w])")


def parse_number(text: str) -> float | None:
    """Parse one number written as plain, comma-grouped, scientific or scale-word form."""
    text = text.strip()
    m = NUMBER_RE.fullmatch(text)
    if m is None:
        return None
    s = text.replace("−", "-").replace(",", "")
    scale = 1.0
    for word, factor in _SCALES.items():
        if s.endswith(" " + word):
            s, scale = s[: -len(word) - 1], factor
    s = s.translate(_SUPERSCRIPT)
    sci = regex.search(r"\s?[×x*]\s?10\^?([-+]?\d+)$", s)
    if sci:
        s = s[: sci.start()] + "e" + sci.group(1)
    try:
        return float(s) * scale
    except ValueError:
        return None


def nice(x: float, sig: int = 4) -> str:
    """Plain decimal with up to `sig` significant digits and no trailing zeros."""
    if x == 0:
        return "0"
    digits = max(sig - 1 - math.floor(math.log10(abs(x))), 0)
    s = f"{x:.{digits}f}"
    return s.rstrip("0").rstrip(".") if "." in s else s


def sci(x: float, rng: random.Random, sig: int = 4) -> str:
    mant, exp = f"{x:.{sig - 1}e}".split("e")
    mant = mant.rstrip("0").rstrip(".") if "." in mant else mant
    exp = int(exp)
    style = rng.randrange(3)
    if style == 0:
        return f"{mant}e{exp}"
    if style == 1:
        return f"{mant}×10^{exp}"
    sup = str(exp).translate(str.maketrans("0123456789-", "⁰¹²³⁴⁵⁶⁷⁸⁹⁻"))
    return f"{mant} × 10{sup}"


def big(x: float, rng: random.Random) -> str:
    """Numbers >= 1e6 as '150 million' style, else comma-grouped or plain."""
    for word, factor in (("trillion", 1e12), ("billion", 1e9), ("million", 1e6)):
        if abs(x) >= factor:
            return f"{nice(x / factor, 4)} {word}"
    if abs(x) >= 1000 and rng.random() < 0.5:
        return f"{x:,.0f}"
    return nice(x)


# unit -> (plain phrasings, alternative-unit converters: (factor from engine unit, phrasing, formatter))
UNITS: dict[str, dict] = {
    "msun": {"plain": ["{n} solar masses", "{n} M☉", "{n} times the mass of the Sun"],
             "alt": [(M_SUN, "{n} kg", "sci")]},
    "mearth": {"plain": ["{n} Earth masses", "{n} M⊕"], "alt": [(M_EARTH, "{n} kg", "sci")]},
    "au": {"plain": ["{n} AU", "{n} astronomical units"], "alt": [(AU / 1e3, "{n} km", "big")]},
    "km": {"plain": ["{n} km", "{n} kilometres"], "alt": [(1e3, "{n} metres", "big")]},
    "kg": {"plain": ["{n} kg"], "alt": [(1 / M_SUN, "{n} solar masses", "nice")]},
    "m": {"plain": ["{n} m", "{n} metres"], "alt": [(1e-3, "{n} km", "nice")]},
    "mpc": {"plain": ["{n} Mpc", "{n} megaparsecs"], "alt": [(MPC / LY, "{n} light-years", "big")]},
    "nm": {"plain": ["{n} nm", "{n} nanometres"], "alt": [(1e-9, "{n} metres", "sci")]},
    "k": {"plain": ["{n} K", "{n} kelvin"], "alt": []},
    "rsun": {"plain": ["{n} solar radii", "{n} R☉"], "alt": []},
    "kms": {"plain": ["{n} km/s", "{n} kilometres per second"], "alt": []},
    "rs": {"plain": ["{n} Schwarzschild radii"], "alt": []},
    "z": {"plain": ["redshift {n}", "z = {n}", "a redshift of {n}"], "alt": []},
    "h0": {"plain": ["{n} km/s/Mpc"], "alt": []},
    "min": {"plain": ["{n} minutes", "{n} min"], "alt": [(1 / 60, "{n} hours", "nice")]},
    "year": {"plain": ["{n}"], "alt": []},
    "deg": {"plain": ["{n}°"], "alt": []},
    "": {"plain": ["{n}"], "alt": []},
}
_MULTIPLIERS = {2.0: ["twice", "double"], 3.0: ["three times"], 4.0: ["four times"], 5.0: ["five times"],
                10.0: ["ten times"], 100.0: ["a hundred times"], 0.5: ["half"], 0.25: ["a quarter of"],
                1 / 3: ["a third of"]}


def _fmt(kind: str, x: float, rng: random.Random) -> str:
    if kind == "sci":
        return sci(x, rng)
    if kind == "big":
        return big(x, rng)
    if abs(x) >= 1e6 or (x != 0 and abs(x) < 1e-3):
        return sci(x, rng)
    return nice(x)


def express(value: float, unit: str, rng: random.Random, baseline: float | None = None,
            baseline_name: str | None = None) -> Expr:
    """Write `value` (in the engine unit) as a phrase; `Expr.value` is what the phrase states."""
    spec = UNITS[unit]
    options = ["plain"]
    # Alternative units only where they read naturally (no "6.5e-25 solar masses").
    alts = [a for a in spec["alt"] if a[2] != "nice" or 1e-3 <= abs(value * a[0]) < 1e6]
    if alts:
        options.append("alt")
    ratio = value / baseline if baseline else None
    if baseline_name and ratio is not None:
        mult = next((words for r, words in _MULTIPLIERS.items() if math.isclose(ratio, r, rel_tol=1e-9)), None)
        if mult:
            options += ["mult", "mult"]
    choice = rng.choice(options)
    if choice == "mult":
        return Expr(f"{rng.choice(mult)} {baseline_name}", value)
    if choice == "alt":
        factor, phrase, kind = rng.choice(alts)
        shown = _fmt(kind, value * factor, rng)
        return Expr(phrase.format(n=shown), parse_number(shown) / factor)
    shown = _fmt("plain", value, rng)
    return Expr(rng.choice(spec["plain"]).format(n=shown), parse_number(shown))