Budget_Tool_Martinhal/backend/app/calc.py

"""Amount expressions and budget arithmetic. All money math uses Decimal."""
import ast
import operator
from dataclasses import dataclass, field
from decimal import ROUND_HALF_UP, Decimal, InvalidOperation

MONTHS = 12
ZERO = Decimal("0")

_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
}


class ExpressionError(ValueError):
    pass


def normalize_number_text(text: str) -> str:
    """Accept '1 234,56', '1.234,56' and '1,234.56' as well as plain numbers."""
    s = text.replace(" ", "").replace(" ", "").replace("€", "")
    if "," in s and "." in s:
        if s.rfind(",") > s.rfind("."):
            s = s.replace(".", "").replace(",", ".")
        else:
            s = s.replace(",", "")
    elif "," in s:
        s = s.replace(",", ".")
    return s


def evaluate(expr: str | None) -> Decimal:
    """Evaluate a typed amount such as '=5512*12' or '3292.56+(12*839.67)'."""
    if expr is None:
        return ZERO
    s = str(expr).strip()
    if s.startswith("="):
        s = s[1:]
    if not s:
        return ZERO
    s = normalize_number_text(s)
    if len(s) > 400:
        raise ExpressionError("expression too long")
    try:
        tree = ast.parse(s, mode="eval")
    except SyntaxError as exc:
        raise ExpressionError(f"invalid expression: {expr}") from exc

    def walk(node: ast.AST) -> Decimal:
        if isinstance(node, ast.Expression):
            return walk(node.body)
        if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)) and not isinstance(node.value, bool):
            return Decimal(str(node.value))
        if isinstance(node, ast.BinOp) and type(node.op) in _OPS:
            left, right = walk(node.left), walk(node.right)
            if isinstance(node.op, ast.Div) and right == 0:
                raise ExpressionError("division by zero")
            return _OPS[type(node.op)](left, right)
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
            value = walk(node.operand)
            return -value if isinstance(node.op, ast.USub) else value
        raise ExpressionError(f"invalid expression: {expr}")

    try:
        return walk(tree)
    except InvalidOperation as exc:
        raise ExpressionError(f"invalid expression: {expr}") from exc


def round2(value: Decimal) -> Decimal:
    return Decimal(value).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)


def factor_for(adj_kind: str, adj_value: Decimal, year_rate: Decimal) -> Decimal:
    adj_value = Decimal(adj_value or 0)
    if adj_kind == "year_rate":
        return Decimal(1) + Decimal(year_rate or 0)
    if adj_kind == "percent":
        return Decimal(1) + adj_value / Decimal(100)
    if adj_kind == "factor":
        return adj_value
    return Decimal(1)


@dataclass
class LineFigures:
    prior: Decimal
    base: Decimal
    factor: Decimal
    adjusted: Decimal
    months: list[Decimal]

    @property
    def change(self) -> Decimal:
        return self.adjusted - self.prior

    @property
    def change_pct(self) -> Decimal | None:
        return None if self.prior == 0 else self.change / self.prior * 100


def line_figures(line, year_rate: Decimal) -> LineFigures:
    # Each line is rounded to cents before it is summed, as the finance reports do.
    prior = round2(Decimal(line.prior_actual or 0))
    if line.monthly_kind == "manual":
        months = [Decimal(str(v or 0)) for v in (line.monthly_values or [])][:MONTHS]
        months += [ZERO] * (MONTHS - len(months))
        total = sum(months, ZERO)
        return LineFigures(prior, total, Decimal(1), total, months)
    base = round2(Decimal(line.base_amount or 0))
    factor = factor_for(line.adj_kind, line.adj_value, year_rate)
    adjusted = round2(Decimal(line.base_amount or 0) * factor)
    if line.monthly_kind == "selected":
        return LineFigures(prior, base, factor, adjusted, selected_spread(adjusted, line.monthly_months))
    return LineFigures(prior, base, factor, adjusted, even_spread(adjusted))


def even_spread(total: Decimal) -> list[Decimal]:
    """Twelve months in cents; December takes the rounding remainder so the months add up exactly."""
    month = round2(total / MONTHS)
    return [month] * (MONTHS - 1) + [total - month * (MONTHS - 1)]


def clean_months(months) -> list[int]:
    """Month numbers 1-12, sorted and without repeats."""
    return sorted({int(m) for m in (months or []) if 1 <= int(m) <= MONTHS})


def selected_spread(total: Decimal, months) -> list[Decimal]:
    """Equal amounts in cents over the chosen months only (1 = January); the last chosen month takes the
    rounding remainder. No months chosen means all twelve."""
    chosen = clean_months(months) or list(range(1, MONTHS + 1))
    share = round2(total / len(chosen))
    out = [ZERO] * MONTHS
    for m in chosen:
        out[m - 1] = share
    out[chosen[-1] - 1] = total - share * (len(chosen) - 1)
    return out


@dataclass
class Totals:
    prior: Decimal = ZERO
    base: Decimal = ZERO
    adjusted: Decimal = ZERO
    months: list[Decimal] = field(default_factory=lambda: [ZERO] * MONTHS)

    def add(self, f: LineFigures) -> None:
        self.prior += f.prior
        self.base += f.base
        self.adjusted += f.adjusted
        self.months = [a + b for a, b in zip(self.months, f.months)]

    @property
    def change(self) -> Decimal:
        return self.adjusted - self.prior

    @property
    def change_pct(self) -> Decimal | None:
        return None if self.prior == 0 else self.change / self.prior * 100