"""
dt_math.py
----------
數學文字正規化、SymPy 解析管線、答案等價比對。

依賴：sympy（純 Python，無 Dash 依賴）

公開 API：
    eq_answer(user_input, answers, ops, tol)  → bool   主要入口
    extract_param_function(eq, level_symbol)  → (expr, symbol)
    extract_k_as_function(eq)                 → expr | None
    _normalize_text(s, ops)                   → str
"""
from __future__ import annotations
import re
import unicodedata

import sympy as sp
from sympy.parsing.sympy_parser import (
    parse_expr,
    standard_transformations,
    implicit_multiplication_application,
)

# ---------------------------------------------------------------------------
# SymPy 設定
# ---------------------------------------------------------------------------
_TRANSFORMS = standard_transformations + (implicit_multiplication_application,)

x, y, z = sp.symbols("x y z")
k, m, n, c, g = sp.symbols("k m n c g")

XY     = {x, y}
PARAMS = {k, m, n, c, g}

_SINGLE_LETTERS = set(list("xyzkmncg"))
KNOWN_FUNCS = ["sin", "cos", "tan", "cot", "sec", "csc"]
KNOWN_CONST = ["pi", "e"]

_ALLOWED = {
    "x": x, "y": y, "z": z,
    "k": k, "m": m, "n": n, "c": c, "g": g,
    "pi": sp.pi, "e": sp.E,
    "sin": sp.sin, "cos": sp.cos,
}
_LOCALS = dict(_ALLOWED)

# ---------------------------------------------------------------------------
# 文字正規化
# ---------------------------------------------------------------------------

def _std_text(s: str) -> str:
    """Unicode 正規化 + 常見符號替換（乘號、除號、減號等）。"""
    t = unicodedata.normalize("NFKC", (s or "").strip())
    t = t.replace("×", "*").replace("∙", "*").replace("·", "*")
    t = t.replace("÷", "/").replace("⁄", "/").replace("／", "/")
    t = t.replace("−", "-").replace("–", "-").replace("—", "-")
    t = t.replace("^", "**")
    t = re.sub(r'(?<=\d),(?=\d)', "", t)  # 移除數字中的千分位逗號
    return t


def _apply_text_ops(t: str, ops: set) -> str:
    if "lower"   in ops: t = t.lower()
    if "nospace" in ops: t = re.sub(r"\s+", "", t)
    return t


def _normalize_text(s: str, ops: list[str]) -> str:
    """依 ops 清單對字串做正規化（strip / lower / collapse_space / strip_dollars）。"""
    if not isinstance(s, str):
        s = "" if s is None else str(s)
    t = s
    if "strip"         in ops: t = t.strip()
    if "lower"         in ops: t = t.lower()
    if "collapse_space" in ops: t = re.sub(r"\s+", " ", t)
    if "strip_dollars" in ops:
        u = t.strip()
        if u.startswith("$$") and u.endswith("$$"):   u = u[2:-2].strip()
        elif u.startswith("$") and u.endswith("$"):   u = u[1:-1].strip()
        t = u
    return t

# ---------------------------------------------------------------------------
# SymPy 解析管線
# ---------------------------------------------------------------------------

def _protect_funcs_and_consts(s: str) -> tuple[str, dict]:
    """把 sin/cos/pi 等先換成佔位符，避免後續拆字母時被錯誤分開。"""
    s2, mapping, i = s, {}, 0
    for w in sorted(KNOWN_FUNCS + KNOWN_CONST, key=len, reverse=True):
        token = f"__W{i}__"
        i += 1
        mapping[token] = w
        s2 = re.sub(rf"\b{w}\b", token, s2)
        s2 = re.sub(rf"(?<=\w){w}(?=\W|\w|\()", token, s2)
    return s2, mapping


def _unprotect(s: str, mapping: dict) -> str:
    for token, w in mapping.items():
        s = s.replace(token, w)
    return s


def _insert_implicit_mul(s: str) -> str:
    """在相鄰 token 之間插入 *（隱式乘法）。"""
    s = re.sub(r"(\d)([A-Za-z_(])", r"\1*\2", s)
    s = re.sub(r"(\))([A-Za-z_(\d])", r"\1*\2", s)
    s = re.sub(r"([A-Za-z_\d])(\()", r"\1*\2", s)
    s = re.sub(r"([A-Za-z_])([A-Za-z_])", r"\1*\2", s)
    return s


def _preprocess_math(s: str) -> str:
    """把數學文字轉成 SymPy-parseable 字串（補乘號、拆字母連寫）。"""
    s = _std_text(s)
    s = re.sub(r"\s+", "", s)

    for fn in sorted(KNOWN_FUNCS, key=len, reverse=True):
        s = re.sub(rf"(?<=[0-9A-Za-z\)]){fn}(?=\()", rf"*{fn}", s)
    for cc in sorted(KNOWN_CONST, key=len, reverse=True):
        s = re.sub(rf"(?<=[0-9A-Za-z\)]){cc}\b", rf"*{cc}", s)

    pat = rf"(?<![A-Za-z])([{''.join(sorted(_SINGLE_LETTERS))}]{{2,}})(?![A-Za-z])"
    s = re.sub(pat, lambda m: "*".join(list(m.group(1))), s)
    return s


def _to_sympy(expr: str):
    """把字串解析成 SymPy 表達式；失敗回 None。"""
    s2 = _preprocess_math(expr)
    try:
        if "=" in s2 and "==" not in s2:
            left, right = s2.split("=", 1)
            return sp.Eq(
                parse_expr(left,  local_dict=_LOCALS, transformations=_TRANSFORMS, evaluate=True),
                parse_expr(right, local_dict=_LOCALS, transformations=_TRANSFORMS, evaluate=True),
            )
        return parse_expr(s2, local_dict=_LOCALS, transformations=_TRANSFORMS, evaluate=True)
    except Exception:
        return None


def _symbolic_equal(a: str, b: str) -> bool:
    """代數等價比對（差為 0 或常數倍）。"""
    ea, eb = _to_sympy(a), _to_sympy(b)
    if ea is None or eb is None: return False
    try:
        if isinstance(ea, sp.Equality): ea = sp.simplify(ea.lhs - ea.rhs)
        if isinstance(eb, sp.Equality): eb = sp.simplify(eb.lhs - eb.rhs)
        if sp.simplify(ea - eb) == 0: return True
        if ea.free_symbols == set() and eb.free_symbols == set(): return False
        try:
            ratio = sp.simplify(ea / eb)
            if ratio != 0 and ratio.free_symbols == set(): return True
        except Exception:
            pass
        return False
    except Exception:
        return False


def _as_number(expr: str):
    """把字串轉成 SymPy 數值；含符號或失敗則回 None。"""
    try:
        e = _to_sympy(expr)
        if e is None: return None
        if getattr(e, "free_symbols", None) and e.free_symbols: return None
        if isinstance(e, sp.Equality): return None
        return sp.N(e)
    except Exception:
        return None


def _numeric_close(a: str, b: str, tol: float = 1e-9) -> bool:
    ua, ub = _as_number(a), _as_number(b)
    if ua is None or ub is None: return False
    try:
        return abs(float(ua) - float(ub)) <= float(tol)
    except Exception:
        return False

# ---------------------------------------------------------------------------
# 主要對外 API
# ---------------------------------------------------------------------------

def eq_answer(
    user_input: str,
    answers: list[str],
    ops: list[str],
    tol: float = 1e-9,
) -> bool:
    """
    判斷使用者輸入是否等於任一正確答案。

    ops 支援：sym（代數等價）/ numeric（數值等價）/ nospace / lower / strip_dollars
    """
    ops_set = {o.strip().lower() for o in (ops or []) if o}
    u0 = _apply_text_ops(_std_text(user_input), ops_set)

    for a in answers or []:
        if u0 == _apply_text_ops(_std_text(a), ops_set):
            return True

    if "sym" in ops_set:
        for a in answers or []:
            if _symbolic_equal(u0, _std_text(a)):
                return True

    if "numeric" in ops_set:
        for a in answers or []:
            if _numeric_close(u0, _std_text(a), tol=tol):
                return True

    return False


def extract_param_function(
    eq: sp.Expr | sp.Eq,
    level_symbol: sp.Symbol | None = None,
) -> tuple:
    """
    從方程式 eq 中解出「某個參數 = f(x,y)」的型態。
    回傳 (f_xy, param_symbol)；失敗回傳 (None, None)。
    """
    if not isinstance(eq, sp.Equality):
        eq = sp.Eq(eq, 0)
    expr0 = sp.simplify(eq.lhs - eq.rhs)
    symbols = expr0.free_symbols
    candidates = sorted((symbols - XY) & PARAMS, key=lambda s: s.name)
    if not candidates: return None, None

    target = level_symbol if level_symbol is not None else candidates[0]
    if target not in candidates: return None, None

    sol = sp.solve(sp.Eq(expr0, 0), target)
    if not sol: return None, None
    return sp.simplify(sol[0]), target


def extract_k_as_function(eq: sp.Expr | sp.Eq):
    """保留原本 API：專門解出 k = f(x,y)。"""
    f_xy, _ = extract_param_function(eq, level_symbol=k)
    return f_xy