"""Notation and output formatting for the shared mathematical representation.

Views own layout; this module owns all mathematical spelling and escaping.
"""
from html import escape
import re

import numpy as np

from math_expression import Expr
from tensor_formatter import format_tensor_value


class MathRenderer:
    def __init__(self, notation='mathbook', html=False, format_config=None):
        if notation not in {'mathbook', 'code'}:
            raise ValueError('Unknown notation')
        self.notation = notation
        self.html = html
        self.format_config = format_config

    def text(self, value):
        return escape(str(value)) if self.html else str(value)

    def equation(self, equation):
        return ' = '.join(self.expression(term) for term in equation.terms)

    def expression(self, expr):
        """Spell one Expr as one indivisible equation stage.

        Never emits line breaks or block-level markup: layout belongs to
        render_derivation(), which may only break *between* stages.
        """
        kind, args = expr.kind, expr.args
        mb = self.notation == 'mathbook'
        if kind == 'literal':
            return self.text(format_tensor_value(expr.value, config=self.format_config))
        if kind in {'ref', 'slot'}:
            return self.text(expr.name)
        if kind == 'adjoint':
            name = self.text(expr.name)
            if not mb:
                return name + '_bar'
            if self.html:
                return f"<span style='text-decoration:overline'>{name}</span>"
            match = re.match(r'^([a-zA-Z]+)(.*)$', name)
            return f'{match[1]}̄{match[2]}' if match else name + '̄'
        if kind == 'contribution':
            source, target = (self.expression(a) for a in args)
            suffix = f', {expr.value}' if expr.value is not None else ''
            if not mb:
                return f'c_{source}_to_{target}' + (f'_{expr.value}' if expr.value is not None else '')
            sub = f'{source}→{target}{suffix}'
            return f'<i>c</i><sub>{sub}</sub>' if self.html else f'c_{sub}'
        if kind in {'partial', 'global_derivative'}:
            first, target = (self.expression(a) for a in args)
            if not mb:
                return f'diff({first}, {target})'
            return f'∂({first})/∂{target}' if kind == 'partial' else f'∂{first}/∂{target}'
        if kind == 'neg':
            return ('−' if mb else '-') + self.group(args[0])
        if kind in {'add', 'sub', 'mul', 'elementwise', 'scale', 'div', 'matmul'}:
            joins = {'add': ' + ', 'sub': ' − ' if mb else ' - ',
                     'mul': ' · ' if mb else ' * ', 'scale': ' · ' if mb else ' * ',
                     'elementwise': ' ⊙ ' if mb else ' * ', 'div': ' / ', 'matmul': ' @ '}
            # Explicit grouping makes the same expression unambiguous in every view.
            rendered = [self.expression(a) if kind == 'add' and a.kind != 'neg' else self.group(a) for a in args]
            return joins[kind].join(rendered)
        if kind == 'pow':
            base = self.group(args[0])
            exponent = self.expression(args[1])
            if mb and self.html:
                return f'{base}<sup>{exponent}</sup>'
            return base + ('^' if mb else '**') + self.group(args[1])
        if kind == 'transpose':
            return self.group(args[0]) + ('<sup>T</sup>' if mb and self.html else ('ᵀ' if mb else '.T'))
        if kind == 'outer' and mb and not any(a.kind == 'literal' for a in args):
            # Named vectors conventionally denote columns. Array previews are
            # horizontal lists, so numerical applications keep explicit outer().
            return self.group(args[0]) + ' ' + self.group(args[1]) + ('<sup>T</sup>' if self.html else 'ᵀ')
        if kind == 'pick':
            return self.expression(args[0]) + '[' + self.expression(args[1]) + ']'
        if kind == 'positive':
            return f'mask({self.expression(args[0])} {self.text("> ")}0)'
        if kind == 'equal':
            return f'mask({self.expression(args[0])} == {self.expression(args[1])})'
        name = '∑' if kind == 'sum' and mb else kind
        return self.text(name) + '(' + ', '.join(self.expression(a) for a in args) + ')'

    def group(self, expr):
        text = self.expression(expr)
        if expr.kind in {'add', 'sub', 'mul', 'elementwise', 'scale', 'div', 'matmul', 'neg'}:
            return '(' + text + ')'
        if expr.kind == 'literal' and np.ndim(expr.value) == 0 and float(expr.value) < 0:
            return '(' + text + ')'
        return text

    def backward_rule(self, explanation):
        """Upper trapezoid compartment. Bindings and rule are separate equations,
        so the break between them is semantic; each equation stays atomic."""
        parts = [self.equation(b) for b in explanation.bindings]
        mapping = ', '.join(parts)
        rule = self.equation(explanation.local_rule)
        if self.html:
            rule = self.atomic(rule)
            return (self.atomic(mapping) + '<br>' + rule) if mapping else rule
        return (mapping + '\n' + rule) if mapping else rule

    def atomic(self, text):
        """Mark one complete equation stage (or line of stages) as unbreakable.

        This is the authoritative no-internal-wrap rule. It is inline because the
        diagram is rendered by mermaid.ink, which never sees the app's CSS.
        """
        return f'<span class="math-stage" style="{NOWRAP}">{text}</span>' if self.html else text

    # Decision B only: *whether* to use the permitted stage-boundary breaks.
    # Decision A (where breaks may occur) is fixed: only between stages.
    INLINE_CHAIN_BUDGET = 45

    def should_stack_derivation(self, derivation, rendered_lhs, rendered_stages):
        if len(rendered_stages) <= 1:
            return False
        if len(rendered_stages) >= 3:
            return True
        plain = [re.sub(r'<[^>]+>', '', s) for s in (rendered_lhs, *rendered_stages)]
        return sum(len(s) for s in plain) + 3 * len(rendered_stages) > self.INLINE_CHAIN_BUDGET

    def render_derivation(self, derivation, layout="auto"):
        """Render a derivation inline (`A = B = C`) or stacked (`A = B` / `= C`).

        Each stage comes from expression() as one atomic unit; a line break is
        only ever placed between equality stages. layout: auto | inline | stacked.
        """
        lhs_str = self.expression(derivation.lhs)
        stages_strs = [self.expression(st) for st in derivation.stages]
        if not stages_strs:
            return self.atomic(lhs_str)

        stack = layout == "stacked" or (
            layout == "auto" and self.should_stack_derivation(derivation, lhs_str, stages_strs))

        if not stack:
            return self.atomic(f"{lhs_str} = " + " = ".join(stages_strs))

        if not self.html:
            indent = ' ' * len(lhs_str)
            return "\n".join([f"{lhs_str} = {stages_strs[0]}"] + [f"{indent} = {st}" for st in stages_strs[1:]])

        lines = [f'<div class="derivation-line" style="{NOWRAP}">{lhs_str} = {stages_strs[0]}</div>']
        lines += [f'<div class="derivation-line derivation-continuation" style="{NOWRAP}">= {st}</div>'
                  for st in stages_strs[1:]]
        return f'<div class="derivation">{"".join(lines)}</div>'

    def application_derivation(self, explanation, layout="auto"):
        """Lower trapezoid compartment: concrete application and result."""
        return self.render_derivation(explanation.application_derivation, layout=layout)

    def explanation(self, explanation):
        """Return the shared upper/lower compartments; containers belong to each view."""
        return self.backward_rule(explanation), self.application_derivation(explanation)


NOWRAP = 'white-space:nowrap'