"""Turn the thesis LaTeX into a pandoc-friendly single file with resolved numbering."""
import re, sys, os

ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
DX = os.path.dirname(os.path.abspath(__file__))
aux = open(f"{DX}/aux/thesis.aux").read()
LBL = {m.group(1): m.group(2) for m in re.finditer(r"\\newlabel\{([^}@]+)\}\{\{([^}]*)\}", aux)}

KIND = {"ch": ("Chapter", "Chapters"), "sec": ("Section", "Sections"), "fig": ("Figure", "Figures"),
        "tab": ("Table", "Tables"), "eq": ("Eq.", "Eqs."), "alg": ("Algorithm", "Algorithms"),
        "lst": ("Listing", "Listings"), "def": ("Definition", "Definitions"),
        "thm": ("Theorem", "Theorems"), "prop": ("Proposition", "Propositions"),
        "app": ("Appendix", "Appendices")}


def num(label):
    n = LBL[label]
    return f"({n})" if label.startswith("eq:") else n


def cref(m):
    labels = [l.strip() for l in m.group(2).split(",")]
    kinds = {l.split(":")[0] for l in labels}
    if len(kinds) == 1:
        k = KIND[labels[0].split(":")[0]]
        nums = [num(l) for l in labels]
        name = k[0] if len(nums) == 1 else k[1]
        return name + "\u00a0" + (nums[0] if len(nums) == 1 else ", ".join(nums[:-1]) + " and " + nums[-1])
    return " and ".join(KIND[l.split(":")[0]][0] + "\u00a0" + num(l) for l in labels)


def resolve_refs(s):
    s = re.sub(r"\\([cC])ref\{([^}]*)\}", cref, s)
    s = re.sub(r"\\ref\{([^}]*)\}", lambda m: LBL[m.group(1)], s)
    return s


def snippet_source():
    """Preamble + tikz/algorithm snippets with refs resolved (for image rendering)."""
    pre = open(f"{ROOT}/thesis.tex").read()
    pre = pre[:pre.index(r"\begin{document}")]
    body = [r"\pagestyle{empty}"]
    for ch in ["04_problem", "05_method", "06_implementation"]:
        s = open(f"{ROOT}/chapters/{ch}.tex").read()
        for m in re.finditer(r"(\\resizebox\{\\textwidth\}\{!\}\{)?\\begin\{tikzpicture\}.*?\\end\{tikzpicture\}\}?"
                             r"|\\begin\{algorithm\}\[t\].*?\\end\{algorithm\}", s, re.S):
            sn = resolve_refs(m.group(0))
            if sn.startswith(r"\begin{algorithm}"):
                sn = sn.replace(r"\begin{algorithm}[t]", r"\begin{algorithm}[H]")
                body.append(r"\begin{minipage}{\textwidth}" + sn + r"\end{minipage}\clearpage")
            else:
                body.append(r"\begin{center}" + sn + r"\end{center}\clearpage")
    return pre + "\\begin{document}\n" + "\n".join(body) + "\n\\end{document}\n"


# ------------------------------------------------------------------ main document
TIKZ = iter(["fig_placements", "fig_typegraph", "fig_architecture"])
ALG = iter(["alg1", "alg2", "alg3"])


def read_all():
    main = open(f"{ROOT}/thesis.tex").read()
    order = re.findall(r"\\input\{(chapters/[^}]+)\}", main)
    parts = []
    for p in order:
        if p.endswith("09_appendix"):
            parts.append("\\APPENDIX\n")
        parts.append(open(f"{ROOT}/{p}.tex").read())
    return "\n".join(parts)


def inline_tables(s):
    def rep(m):
        return open(f"{ROOT}/{m.group(1)}").read()
    return re.sub(r"\\csname @@input\\endcsname (tables/\S+\.tex)", rep, s)


def strip_comments(s):
    return re.sub(r"(?<!\\)%.*", "", s)


def number_headings(s):
    out, ch, sec, sub, appendix = [], 0, 0, 0, False
    for line in s.split("\n"):
        if line.startswith("\\APPENDIX"):
            appendix, ch = True, 0
            continue
        m = re.match(r"\\(chapter|section|subsection)(\*?)\{(.*)\}\s*$", line)
        if m:
            kind, star, title = m.groups()
            if star:
                out.append(f"\\{kind}*{{{title}}}")
                continue
            if kind == "chapter":
                ch, sec, sub = ch + 1, 0, 0
                n = chr(64 + ch) if appendix else str(ch)
                out.append(f"\\chapter*{{{'Appendix ' if appendix else ''}{n}{':' if appendix else ''} {title}}}")
                cur = n
            elif kind == "section":
                sec, sub = sec + 1, 0
                out.append(f"\\section*{{{cur}.{sec} {title}}}")
            else:
                sub += 1
                out.append(f"\\subsection*{{{cur}.{sec}.{sub} {title}}}")
            continue
        out.append(line)
    return "\n".join(out)


def env_numbering(s):
    """Number equations and theorem-like environments per chapter, as LaTeX does."""
    out, ch = [], 0
    counters = {}
    pos = 0
    token = re.compile(r"\\chapter\*\{(?:Appendix )?([0-9A-Z]+)[: ]|\\begin\{(equation|definition|theorem|lemma|"
                       r"proposition|corollary|remark|example)\}(\[[^\]]*\])?")
    res = []
    for m in token.finditer(s):
        res.append(s[pos:m.start()])
        pos = m.end()
        if m.group(1):
            ch = m.group(1)
            counters = {}
            res.append(m.group(0))
            continue
        env, opt = m.group(2), m.group(3)
        key = {"lemma": "theorem", "proposition": "theorem", "corollary": "theorem"}.get(env, env)
        counters[key] = counters.get(key, 0) + 1
        n = f"{ch}.{counters[key]}"
        if env == "equation":
            res.append(f"\\begin{{equation}}\\EQNUM{{{n}}}")
        else:
            res.append(f"\\begin{{{env}}}\\THMNUM{{{n}}}{opt or ''}")
    res.append(s[pos:])
    return "".join(res)


def convert_equations(s):
    def rep(m):
        body = m.group(2)
        body = re.sub(r"\\label\{[^}]*\}", "", body).strip()
        body = body.replace("\\;\\xrightarrow{N\\to\\infty}\\;", "\\xrightarrow{N\\to\\infty}")
        if "\\begin{aligned}" in body:
            body = body.replace("\\end{aligned}", f"\\end{{aligned}}\\qquad \\text{{EQNO({m.group(1)})}}")
        else:
            body = body.rstrip(",. ") + f"\\qquad \\text{{EQNO({m.group(1)})}}"
        return f"\\[\n{body}\n\\]"
    return re.sub(r"\\begin\{equation\}\\EQNUM\{([^}]*)\}(.*?)\\end\{equation\}", rep, s, flags=re.S)


def convert_theorems(s):
    names = {"definition": "Definition", "theorem": "Theorem", "lemma": "Lemma", "proposition": "Proposition",
             "corollary": "Corollary", "remark": "Remark", "example": "Example"}

    def rep(m):
        env, n, opt, body = m.group(1), m.group(2), m.group(3), m.group(4)
        body = re.sub(r"\\label\{[^}]*\}", "", body).strip()
        head = f"\\textbf{{{names[env]} {n}}}" + (f" ({opt})" if opt else "") + "."
        style = "\\emph{" if env in ("definition", "theorem", "lemma", "proposition", "corollary") else "{"
        return f"\n\n{head} {body}\n\n"
    s = re.sub(r"\\begin\{(\w+)\}\\THMNUM\{([^}]*)\}(?:\[([^\]]*)\])?(.*?)\\end\{\1\}", rep, s, flags=re.S)

    def proof(m):
        title = m.group(1) or "Proof"
        body = m.group(2).strip()
        return f"\n\n\\emph{{{title}.}} {body} \u220e\n\n"
    s = re.sub(r"\\begin\{proof\}(?:\[([^\]]*)\])?(.*?)\\end\{proof\}", proof, s, flags=re.S)
    return s


def convert_listings(s):
    def rep(m):
        opts, code = m.group(1), m.group(2)
        cap = re.search(r"caption=\{(.*?)\}\s*,\s*label=\{([^}]*)\}", opts, re.S)
        head = ""
        if cap:
            head = f"\\textbf{{Listing {LBL[cap.group(2)]}:}} {cap.group(1)}\n\n"
        else:
            cap2 = re.search(r"caption=\{(.*)\}\]?$", opts, re.S)
            if cap2:
                head = f"\\textbf{{Listing:}} {cap2.group(1)}\n\n"
        return f"{head}\\begin{{verbatim}}{code}\\end{{verbatim}}"
    return re.sub(r"\\begin\{lstlisting\}(\[.*?\])?\n(.*?)\\end\{lstlisting\}", rep, s, flags=re.S)


def convert_figures(s):
    def rep(m):
        body = m.group(1)
        lab = re.search(r"\\label\{(fig:[^}]*)\}", body)
        n = LBL[lab.group(1)].rstrip("ab") if lab else ""
        # main caption (skip optional short caption)
        cap = re.search(r"\\caption(?:\[[^\]]*\])?\{(.*)\}\s*\\label\{fig:", body, re.S)
        caption = cap.group(1).strip() if cap else ""
        if "tikzpicture" in body:
            imgs = [(f"img/{next(TIKZ)}.png", "")]
        elif "subfigure" in body:
            imgs = []
            for sm in re.finditer(r"\\includegraphics\[[^\]]*\]\{figures/([^}]*)\.pdf\}\s*\\caption\{([^}]*)\}", body):
                imgs.append((f"img/{sm.group(1)}.png", sm.group(2)))
            caption = re.search(r"\\end\{subfigure\}\s*\\caption\{(.*)\}\s*\\end", body + "\\end", re.S)
            caption = caption.group(1).strip() if caption else ""
            outs = []
            for i, (img, sub) in enumerate(imgs):
                letter = "ab"[i]
                outs.append(f"\\begin{{figure}}\n\\centering\n\\includegraphics[width=0.75\\textwidth]{{{img}}}\n"
                             f"\\caption{{Figure {n}{letter}: {sub}}}\n\\end{{figure}}")
            return "\n\n".join(outs) + f"\n\n\\emph{{Figure {n}: {caption}}}\n\n"
        else:
            imgs = [(f"img/{f}.png", "") for f in re.findall(r"\\includegraphics\[[^\]]*\]\{figures/([^}]*)\.pdf\}", body)]
        width = "0.7\\textwidth" if "first_hit" in imgs[0][0] else "\\textwidth"
        return (f"\\begin{{figure}}\n\\centering\n\\includegraphics[width={width}]{{{imgs[0][0]}}}\n"
                f"\\caption{{Figure {n}: {caption}}}\n\\end{{figure}}")
    s = re.sub(r"\\begin\{figure\}\[[^\]]*\](.*?)\\end\{figure\}", rep, s, flags=re.S)
    s = re.sub(r"\\begin\{algorithm\}\[t\].*?\\end\{algorithm\}",
               lambda m: f"\\begin{{figure}}\n\\centering\n\\includegraphics[width=\\textwidth]{{img/{next(ALG)}.png}}\n\\end{{figure}}",
               s, flags=re.S)
    return s


def tabular_rows(body):
    m = re.search(r"\\begin\{(tabularx?)\}", body)
    i = m.end()
    if body.startswith("{\\textwidth}", i):
        i += len("{\\textwidth}")
    assert body[i] == "{"
    depth = 0
    for j in range(i, len(body)):
        depth += {"{": 1, "}": -1}.get(body[j], 0)
        if depth == 0:
            break
    end = body.index("\\end{" + m.group(1) + "}", j)
    return body[j + 1:end]


def convert_tables(s):
    def rep(m):
        body = m.group(1)
        lab = re.search(r"\\label\{(tab:[^}]*)\}", body)
        n = LBL[lab.group(1)] if lab else ""
        cap = re.search(r"\\caption(?:\[[^\]]*\])?\{(.*?)\}\s*\\label", body, re.S)
        caption = cap.group(1).strip() if cap else ""
        rows = tabular_rows(body)
        rows = re.sub(r"\\(toprule|midrule|bottomrule)", "", rows)
        rows = re.sub(r"\\cmidrule(\([^)]*\))?\{[^}]*\}", "", rows)
        # table-specific simplifications
        if n == LBL.get("tab:dataset"):
            rows = re.sub(r"^\s*& & & & & \\multicolumn.*?\\\\\s*$", "", rows, flags=re.M)
            rows = rows.replace("Service & Redfish & $|V|$ & $D$ & $|\\mathcal{G}|$ & legacy & modern & union & SP-BFS & Tree-BFS & union & SP-BFS",
                                "Service & Redfish & $|V|$ & $D$ & $|\\mathcal{G}|$ & R legacy & R modern & R union & R SP & GET tree & GET union & GET SP")
            rows = rows.replace("\\multicolumn{4}{@{}l}{Total / micro-average}", "Total / micro-average & & &")
        lines = [l.strip() for l in rows.split("\\\\") if l.strip()]
        ncol = lines[0].count("&") + 1
        spec = "l" * ncol
        new = "\\begin{tabular}{" + spec + "}\n\\toprule\n" + lines[0] + " \\\\\n\\midrule\n" + \
              " \\\\\n".join(lines[1:]) + " \\\\\n\\bottomrule\n\\end{tabular}"
        if n == LBL.get("tab:dataset"):
            caption += " (R = recall, GET = HTTP GET requests; tree = Tree-BFS, SP = SP-BFS.)"
        return f"\\begin{{table}}\n\\caption{{Table {n}: {caption}}}\n{new}\n\\end{{table}}"
    return re.sub(r"\\begin\{table\}\[[^\]]*\](.*?)\\end\{table\}", rep, s, flags=re.S)


def misc(s):
    s = s.replace("G^\\ast", "\\tilde{G}").replace("\\in (0, 1]", ",\\quad 0 < \\rho \\le 1")
    s = s.replace("\\Bigl(", "\\left(").replace("\\Bigr)", "\\right)").replace("\\max\\!", "\\max")
    s = s.replace("\\SPBFS{}", "SP-BFS")
    # math occurrences of \SPBFS
    s = re.sub(r"(_\{)\\SPBFS(\})", r"\1\\text{SP-BFS}\2", s)
    s = re.sub(r"\\SPBFS(?![A-Za-z])", "SP-BFS", s)
    s = re.sub(r"\\SI\{([^}]*)\}\{\\milli\\second\}", r"\1 ms", s)
    s = re.sub(r"\\SI\{([^}]*)\}\{\\second\}", r"\1 s", s)
    s = re.sub(r"\\SI\{([^}]*)\}\{\}", r"\1", s)
    s = s.replace("\\si{\\celsius}", "°C").replace("\\si{\\watt}", "W")
    s = re.sub(r"\\paragraph\{([^}]*)\}", r"\\textbf{\1}", s)
    s = re.sub(r"\\begin\{(enumerate|description|itemize)\}\[[^\]]*\]", r"\\begin{\1}", s)
    # heterogeneity list H1..H5
    s = s.replace("\\begin{enumerate}\n  \\item \\textbf{Identifier heterogeneity.}",
                  "\\begin{description}\n  \\item[H1] \\textbf{Identifier heterogeneity.}")
    for i, t in enumerate(["Model generation.", "Placement.", "Optional features.", "Vendor extensions and defects."], 2):
        s = s.replace(f"\\item \\textbf{{{t}}}", f"\\item[H{i}] \\textbf{{{t}}}")
    s = s.replace("and navigation properties.\n\\end{enumerate}", "and navigation properties.\n\\end{description}")
    s = re.sub(r"(\\item\[H5\].*?)\\end\{enumerate\}", r"\1\\end{description}", s, flags=re.S)
    s = s.replace("\\addcontentsline{toc}{chapter}{Abstract}", "").replace(
        "\\addcontentsline{toc}{chapter}{List of Abbreviations}", "")
    s = s.replace("\\vspace{1em}", "").replace("\\noindent", "")
    s = s.replace("@{}p{2.6cm}p{11cm}@{}", "ll")
    return s


MACROS = r"""
\newcommand{\rf}[1]{\texttt{#1}}
\newcommand{\Tc}{\mathcal{T}}
\newcommand{\anc}{\operatorname{Anc}}
\newcommand{\reach}{\operatorname{R}}
"""

if __name__ == "__main__":
    if sys.argv[1:] == ["snips"]:
        open(f"{DX}/snips.tex", "w").write(snippet_source())
        sys.exit()
    s = read_all()
    s = strip_comments(s)
    s = inline_tables(s)
    s = misc(s)
    s = number_headings(s)
    s = env_numbering(s)
    s = resolve_refs(s)
    s = convert_equations(s)
    s = convert_theorems(s)
    s = convert_listings(s)
    s = convert_figures(s)
    s = convert_tables(s)
    s = re.sub(r"\\label\{[^}]*\}", "", s)
    open(f"{DX}/body.tex", "w").write(MACROS + "\n\\begin{document}\n" + s + "\n\\end{document}\n")
    print("ok", len(s))