from __future__ import annotations

import hashlib
import json
import re
import sys
import time
from pathlib import Path

from deep_translator import GoogleTranslator


ROOT = Path(__file__).resolve().parent
RAW = ROOT / "evans_chapter2_raw.txt"
CACHE = ROOT / "evans_chapter2_translation_cache.json"
TEX = ROOT / "Evans_Chapter2_Chinese.tex"


def is_header_or_footer(line: str) -> bool:
    s = line.strip()
    if not s:
        return False
    if re.fullmatch(r"\d+", s):
        return True
    if re.search(r"^\d+\s+2\.\s+Four Important Linear Partial Differential Equations\s*$", s):
        return True
    if re.search(r"^2\.[1-6]\.\s+.*\s+\d+\s*$", s):
        return True
    if re.search(r"^Chapter\s+2\s*$", s):
        return False
    return False


def line_kind(line: str) -> str:
    s = line.strip()
    if not s:
        return "blank"
    if s in {"Chapter 2", "Four Important Linear", "Partial Differential", "Equations"}:
        return "title"
    if re.match(r"^2\.[1-6]\.", s):
        return "heading"
    if re.match(r"^(THEOREM|Proof\.|Remark\.|Weak solutions\.|Physical interpretation\.|Interpretation of fundamental solution\.)", s):
        return "text"
    ascii_letters = len(re.findall(r"[A-Za-z]", s))
    math_chars = len(re.findall(r"[𝑎-𝙿ℝℕℤ∂∇∆Δ∫∑∏√∞≤≥≠∈∉⊂⊃⊆⊇→←↦⟶≔≈±−⋅·×∪∩∅∂∀∃∣⎧⎨⎩{}|+−=<>/\u0300-\u036f]", s))
    if math_chars >= 3 and ascii_letters <= 10:
        return "formula"
    if re.match(r"^\(?\d+\)?\s+", s) and math_chars >= 2 and ascii_letters <= 20:
        return "formula"
    return "text"


def looks_like_formula(text: str) -> bool:
    s = text.strip()
    if not s:
        return False
    ascii_letters = len(re.findall(r"[A-Za-z]", s))
    cjk = len(re.findall(r"[\u4e00-\u9fff]", s))
    math_chars = len(re.findall(r"[𝑎-𝙿ℝℕℤ∂∇∆Δ∫∑∏√∞≤≥≠∈∉⊂⊃⊆⊇→←↦⟶≔≈±−⋅·×∪∩∅∂∀∃∣⎧⎨⎩{}|+−=<>/\u0300-\u036f]", s))
    return cjk == 0 and math_chars >= max(1, ascii_letters)


def dehyphenate_join(prev: str, cur: str) -> str:
    if prev.endswith("-") and cur and cur[0].islower():
        return prev[:-1] + cur
    return prev + " " + cur


def build_blocks(raw: str) -> list[dict[str, str]]:
    blocks: list[dict[str, str]] = []
    current: list[str] = []
    current_kind = "text"

    def flush() -> None:
        nonlocal current, current_kind
        if not current:
            return
        if current_kind in {"text", "heading"}:
            text = current[0].strip()
            for part in current[1:]:
                text = dehyphenate_join(text, part.strip())
            blocks.append({"kind": current_kind, "text": text})
        else:
            blocks.append({"kind": current_kind, "text": "\n".join(x.rstrip() for x in current)})
        current = []

    for raw_line in raw.replace("\r\n", "\n").split("\n"):
        line = raw_line.replace("\f", "").rstrip()
        if is_header_or_footer(line):
            continue
        kind = line_kind(line)
        if kind == "blank":
            flush()
            blocks.append({"kind": "blank", "text": ""})
            continue
        if kind in {"formula", "title"}:
            flush()
            blocks.append({"kind": kind, "text": line.strip()})
            continue
        if kind == "heading":
            flush()
            blocks.append({"kind": "heading", "text": line.strip()})
            continue
        if current_kind != "text":
            flush()
            current_kind = "text"
        current.append(line)
    flush()

    compact: list[dict[str, str]] = []
    blank = False
    for block in blocks:
        if block["kind"] == "blank":
            if not blank:
                compact.append(block)
            blank = True
        else:
            compact.append(block)
            blank = False
    return compact


def glossary_polish(text: str) -> str:
    replacements = [
        ("偏微分方程式", "偏微分方程"),
        ("拉普拉斯方程式", "拉普拉斯方程"),
        ("热方程式", "热方程"),
        ("波方程式", "波方程"),
        ("波动方程", "波方程"),
        ("输送方程", "输运方程"),
        ("传输方程", "输运方程"),
        ("基本溶液", "基本解"),
        ("基本解决方案", "基本解"),
        ("根本求解方案", "基本解"),
        ("根本解", "基本解"),
        ("求解方案", "解"),
        ("弱解方案", "弱解"),
        ("弱求解方案", "弱解"),
        ("平均值属性", "平均值性质"),
        ("平均值公式", "平均值公式"),
        ("最大原理", "最大值原理"),
        ("强最大原理", "强最大值原理"),
        ("格林的函数", "格林函数"),
        ("格林函数", "Green 函数"),
        ("泊松公式", "Poisson 公式"),
        ("泊松方程", "Poisson 方程"),
        ("狄利克雷", "Dirichlet"),
        ("柯西问题", "Cauchy 问题"),
        ("哈纳克", "Harnack"),
        ("刘维尔", "Liouville"),
        ("杜哈梅尔", "Duhamel"),
        ("谐波函数", "调和函数"),
        ("谐函数", "调和函数"),
        ("有界域", "有界区域"),
        ("非齐次", "非齐次"),
        ("非同质", "非齐次"),
        ("初始值问题", "初值问题"),
        ("边界值问题", "边值问题"),
        ("线路", "直线"),
        ("修复任何", "固定任意"),
        ("格林的职能", "Green 函数"),
        ("解决", "求解"),
        ("求解方案", "解"),
        ("弱求解方案", "弱解"),
        ("让 ", "令 "),
        ("让𝑢", "令 𝑢"),
        ("让 𝑢", "令 𝑢"),
        ("然后 ", "则 "),
        ("然后，", "于是，"),
        ("因此，", "于是，"),
        ("我们断言", "我们断言"),
        ("在 ℝ", "于 ℝ"),
        ("在 𝑈", "于 𝑈"),
        ("在 𝐵", "于 𝐵"),
        ("如果 ", "若 "),
        ("假设 ", "设 "),
        ("证明。", "证明."),
        ("备注。", "注."),
    ]
    for old, new in replacements:
        text = text.replace(old, new)
    text = re.sub(r"定理\s*(\d+)", r"定理 \1", text)
    text = re.sub(r"\sε\s", " ∈ ", text)
    text = text.replace(" ε ", " ∈ ")
    text = text.replace("THEOREM", "定理")
    text = text.replace("Proof.", "证明.")
    text = text.replace("Remark.", "注.")
    return text


def translate_blocks(blocks: list[dict[str, str]]) -> list[dict[str, str]]:
    cache: dict[str, str] = {}
    if CACHE.exists():
        cache = json.loads(CACHE.read_text(encoding="utf-8"))
    translator = GoogleTranslator(source="en", target="zh-CN")
    translated: list[dict[str, str]] = []
    total = sum(1 for b in blocks if b["kind"] in {"text", "heading"})
    done = 0
    last_save = time.time()

    for block in blocks:
        kind = block["kind"]
        text = block["text"]
        if kind not in {"text", "heading"} or not text.strip():
            translated.append(block)
            continue
        key = hashlib.sha256(text.encode("utf-8")).hexdigest()
        if key not in cache or cache.get(key) in {None, ""}:
            for attempt in range(5):
                try:
                    result = translator.translate(text)
                    if not result:
                        raise RuntimeError("empty translation result")
                    cache[key] = result
                    break
                except Exception as exc:
                    if attempt == 4:
                        cache[key] = "[待校订] " + text
                    time.sleep(2 + attempt * 3)
        done += 1
        if time.time() - last_save > 15:
            CACHE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8")
            print(f"translated {done}/{total}", flush=True)
            last_save = time.time()
        cached = cache[key]
        if isinstance(cached, str) and cached.startswith("[待校订]") and looks_like_formula(text):
            translated.append({"kind": "formula", "text": text})
        else:
            translated.append({"kind": kind, "text": glossary_polish(cached)})
    CACHE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"translated {done}/{total}", flush=True)
    return translated


def tex_escape_segment(s: str) -> str:
    s = s.replace("\\", r"\textbackslash{}")
    s = s.replace("&", r"\&")
    s = s.replace("%", r"\%")
    s = s.replace("$", r"\$")
    s = s.replace("#", r"\#")
    s = s.replace("_", r"\_")
    s = s.replace("{", r"\{")
    s = s.replace("}", r"\}")
    s = s.replace("^", r"\^{}")
    s = s.replace("~", r"\~{}")
    return s


def is_math_inline_char(ch: str) -> bool:
    code = ord(ch)
    return (
        0x1D400 <= code <= 0x1D7FF
        or 0x0391 <= code <= 0x03C9
        or ch in "ℝℕℤ∂∇∆Δ∫∬∮∑∏√∞≤≥≠≡≢∈∉⊂⊃⊆⊇→←↦⟶≔≕≈±−⋅·×∪∩∅∀∃∣∶∗⋯□⨍⎧⎨⎩⎪"
    )


def tex_escape_text(s: str) -> str:
    out: list[str] = []
    run: list[str] = []
    for ch in s:
        if is_math_inline_char(ch):
            run.append(ch)
            continue
        if run:
            out.append(r"{\formulafont " + tex_escape_segment("".join(run)) + "}")
            run = []
        out.append(tex_escape_segment(ch))
    if run:
        out.append(r"{\formulafont " + tex_escape_segment("".join(run)) + "}")
    return "".join(out)


def tex_escape_formula(s: str) -> str:
    s = s.replace("ᵆ", "𝑢")
    s = s.replace("we deduce", "可推出")
    s = s.replace("where", "其中")
    s = s.replace("Then", "则")
    s = s.replace("then", "则")
    s = s.replace("Assume", "假设")
    s = s.replace("If", "若")
    s = s.replace("Extend", "延拓")
    s = s.replace("Prove", "证明")
    s = re.sub(r"\bfor\b", "对于", s)
    s = re.sub(r"\bin\b", "于", s)
    s = re.sub(r"\bon\b", "在", s)
    s = s.replace("\\", r"\textbackslash{}")
    s = s.replace("&", r"\&")
    s = s.replace("%", r"\%")
    s = s.replace("$", r"\$")
    s = s.replace("#", r"\#")
    s = s.replace("_", r"\_")
    s = s.replace("{", r"\{")
    s = s.replace("}", r"\}")
    s = s.replace("^", r"\^{}")
    s = s.replace("~", r"\~{}")
    return s


def heading_to_tex(s: str) -> str:
    m = re.match(r"^(2\.\d+\.\d+)\.\s*(.*)$", s)
    if not m:
        m = re.match(r"^(2\.\d+)\.\s*(\d+\..*)$", s)
    if m:
        return rf"\subsection*{{{tex_escape_text(m.group(1) + '. ' + m.group(2))}}}"
    m = re.match(r"^(2\.\d+)\.\s*(.*)$", s)
    if m:
        return rf"\section*{{{tex_escape_text(m.group(1) + '. ' + m.group(2))}}}"
    return rf"\section*{{{tex_escape_text(s)}}}"


def title_to_cn(s: str) -> str:
    mapping = {
        "Chapter 2": "第 2 章",
        "Four Important Linear": "四个重要的线性",
        "Partial Differential": "偏微分",
        "Equations": "方程",
    }
    return mapping.get(s.strip(), s.strip())


def render_tex(blocks: list[dict[str, str]]) -> str:
    preamble = r"""\documentclass[UTF8,openany,zihao=-4]{ctexbook}
\usepackage[a4paper,margin=2.45cm,headheight=15pt]{geometry}
\usepackage{fontspec}
\usepackage{unicode-math}
\usepackage{xcolor}
\usepackage{titlesec}
\usepackage{fancyhdr}
\usepackage{microtype}
\setCJKmainfont{SimSun}
\setCJKsansfont{Microsoft YaHei}
\setmainfont{STIX Two Text}
\setmathfont{STIX Two Math}
\newfontfamily\formulafont{XITS Math}
\definecolor{evansblue}{HTML}{1F4E79}
\definecolor{evansgray}{HTML}{F3F5F7}
\titleformat{\section}{\Large\bfseries\color{evansblue}}{}{0pt}{}
\titleformat{\subsection}{\large\bfseries}{}{0pt}{}
\pagestyle{fancy}
\fancyhf{}
\lhead{Evans《偏微分方程》第二章中文译稿}
\rhead{\thepage}
\renewcommand{\headrulewidth}{0.4pt}
\setlength{\parindent}{2em}
\setlength{\parskip}{0.45em}
\newenvironment{displayformula}{\begin{quote}\begingroup\formulafont\small\obeylines\obeyspaces\raggedright}{\endgroup\end{quote}}
\begin{document}
\frontmatter
\begin{titlepage}
\centering
\vspace*{2.3cm}
{\Huge\bfseries 四个重要的线性偏微分方程\par}
\vspace{0.9cm}
{\Large Lawrence C. Evans, \textit{Partial Differential Equations}, Second Edition\par}
\vspace{0.55cm}
{\Large Chapter 2 中文译稿\par}
\vfill
{\large 保留原章节结构、编号、定理与公式次序\par}
\end{titlepage}
\mainmatter
"""
    out = [preamble]
    title_buffer: list[str] = []
    for block in blocks:
        kind = block["kind"]
        text = block["text"].strip("\n")
        if kind == "blank":
            continue
        if kind == "title":
            title_buffer.append(title_to_cn(text))
            if text == "Equations":
                out.append(r"\chapter*{" + tex_escape_text("".join(title_buffer)) + "}\n")
                title_buffer = []
            continue
        if kind == "heading":
            out.append(heading_to_tex(text) + "\n")
        elif kind == "formula":
            out.append("\\begin{displayformula}\n" + tex_escape_formula(text) + "\n\\end{displayformula}\n")
        else:
            out.append(tex_escape_text(text) + "\n\n")
    out.append("\\end{document}\n")
    return "".join(out)


def main() -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    raw = RAW.read_text(encoding="utf-8", errors="replace")
    blocks = build_blocks(raw)
    translated = translate_blocks(blocks)
    TEX.write_text(render_tex(translated), encoding="utf-8")
    print(f"wrote {TEX}", flush=True)


if __name__ == "__main__":
    main()