"""IR → standalone Python solver script."""

from __future__ import annotations

import ast
import functools
import re
from typing import Literal

from md2mip.ir import IR, Constraint, Variable


def generate(ir: IR) -> str:
    """Generate a standalone Python solver script from an IR."""
    ir.validate()
    parts = [
        _header(ir),
        _load_data(ir),
        _build_model(ir),
        *_constraint_functions(ir),
        _parse_highs_value_fn(),
        _solve_and_report(ir),
        _main(ir),
    ]
    return "\n\n".join(parts) + "\n"


# ---------------------------------------------------------------------------
# Variable layout helpers
# ---------------------------------------------------------------------------


def _var_size_expr(var: Variable, ir: IR) -> str:
    """Return a Python expression for the number of elements in a variable group."""
    if not var.indices:
        return "1"
    parts = [f"n_{s.lower()}" for s in var.indices]
    return " * ".join(parts)


def _total_vars_expr(ir: IR) -> str:
    parts = []
    for var in ir.variables.values():
        parts.append(_var_size_expr(var, ir))
    return " + ".join(parts)


# ---------------------------------------------------------------------------
# Code sections
# ---------------------------------------------------------------------------


def _header(ir: IR) -> str:
    return f'''#!/usr/bin/env python3
"""{ir.name} model — generated by md2mip"""

import argparse
import json
import sys

import highspy
import numpy as np
import yaml

CONSTRAINT_NAMES = {{}}
VARIABLE_NAMES = {{}}


def _merge_row(idx, val):
    """Merge duplicate column indices by summing their coefficients."""
    if not idx:
        return idx, val
    merged = {{}}
    for c, v in zip(idx, val):
        merged[c] = merged.get(c, 0.0) + v
    cols = sorted(c for c, v in merged.items() if abs(v) > 1e-15)
    return cols, [merged[c] for c in cols]'''


def _load_data(ir: IR) -> str:
    lines = [
        "def load_data(path: str) -> dict:",
        '    """Load data from YAML or JSON into numpy arrays."""',
        "    with open(path) as f:",
        '        if path.endswith(".json"):',
        "            raw = json.load(f)",
        "        else:",
        "            raw = yaml.safe_load(f)",
        "    data = {}",
    ]
    for sname in ir.sets:
        lines.append(f'    data["{sname}"] = list(raw["{sname}"])')
    for pname, param in ir.parameters.items():
        if param.indices:
            # Handle scalar-in-YAML for indexed params (broadcast to array)
            size_parts = [f'len(data["{s}"])' for s in param.indices]
            total_size = " * ".join(size_parts)
            lines.append(f'    _v_{pname} = raw["{pname}"]')
            lines.append(f"    if isinstance(_v_{pname}, (int, float)):")
            lines.append(f'        data["{pname}"] = np.full({total_size}, float(_v_{pname}))')
            lines.append("    else:")
            lines.append(f'        data["{pname}"] = np.array(_v_{pname}, dtype=np.float64)')
            if len(param.indices) > 1:
                shape = ", ".join(size_parts)
                lines.append(f'    data["{pname}"] = data["{pname}"].reshape({shape})')
        else:
            if param.default is not None:
                lines.append(f'    data["{pname}"] = float(raw.get("{pname}", {param.default}))')
            else:
                lines.append(f'    data["{pname}"] = float(raw["{pname}"])')
    lines.append("    return data")
    return "\n".join(lines)


def _build_model(ir: IR) -> str:
    lines = ["def build_model(data: dict) -> highspy.Highs:"]

    # Set sizes
    for sname in ir.sets:
        lines.append(f'    n_{sname.lower()} = len(data["{sname}"])')

    lines.append("")
    lines.append("    h = highspy.Highs()")
    lines.append("")

    # Total variables
    n_vars_expr = _total_vars_expr(ir)
    lines.append(f"    n_vars = {n_vars_expr}")
    lines.append("")

    # Add variables with name tracking
    for vname, var in ir.variables.items():
        size_expr = _var_size_expr(var, ir)
        lb = var.lower_bound if var.lower_bound is not None else 0
        ub_val = var.upper_bound
        if var.type == "binary":
            lb, ub_val = 0, 1
        lb_str = f"np.zeros({size_expr})" if lb == 0 else f"np.full({size_expr}, {lb})"
        if ub_val is None:
            ub_str = f"np.full({size_expr}, np.inf)"
        else:
            ub_str = f"np.full({size_expr}, {ub_val})"
        lines.append(f"    # --- Variable: {vname}{_index_comment(var)} ---")
        lines.append("    _col_before = h.getNumCol()")
        lines.append(f"    h.addVars({size_expr}, {lb_str}, {ub_str})")
        lines.append(
            f"    VARIABLE_NAMES.update("
            f'{{c: "{vname}" for c in range(_col_before, h.getNumCol())}})'
        )
        if var.type == "binary":
            offset_expr = _offset_expr(vname, ir)
            lines.append(f"    for _k in range({size_expr}):")
            lines.append(
                f"        h.changeColIntegrality({offset_expr} + _k, highspy.HighsVarType.kInteger)"
            )

    lines.append("")

    # Objective
    lines.append("    # --- Objective ---")
    obj = ir.objective
    obj_code = _generate_objective(obj.expression, ir)
    lines.extend(f"    {line}" for line in obj_code)
    sense = "kMinimize" if obj.sense == "minimize" else "kMaximize"
    lines.append(f"    h.changeObjectiveSense(highspy.ObjSense.{sense})")
    lines.append("")

    # Constraint groups with name tracking
    for cname in ir.constraints:
        fn_name = f"setup_{cname}_constraints"
        args = _constraint_fn_args(ir)
        lines.append("    _row_before = h.getNumRow()")
        lines.append(f"    {fn_name}({args})")
        lines.append(
            f"    CONSTRAINT_NAMES.update("
            f'{{r: "{cname}" for r in range(_row_before, h.getNumRow())}})'
        )

    lines.append("")
    lines.append("    return h")
    return "\n".join(lines)


def _index_comment(var: Variable) -> str:
    if not var.indices:
        return ""
    idx = ",".join(s.lower() for s in var.indices)
    return f"[{idx}]"


def _offset_expr(vname: str, ir: IR) -> str:
    """Return a Python expression for the flat offset of variable group vname."""
    offset_parts = []
    for name, var in ir.variables.items():
        if name == vname:
            break
        offset_parts.append(_var_size_expr(var, ir))
    if not offset_parts:
        return "0"
    return " + ".join(offset_parts)


def _constraint_fn_args(ir: IR) -> str:
    parts = ["h", "data"]
    for sname in ir.sets:
        parts.append(f"n_{sname.lower()}")
    return ", ".join(parts)


def _parse_comparison(expr_str: str) -> tuple[str, str, str]:
    """Split 'lhs <= rhs' into (lhs, op, rhs). Also handles >=, ==."""
    for op in ("<=", ">=", "=="):
        if op in expr_str:
            parts = expr_str.split(op, 1)
            return parts[0].strip(), op, parts[1].strip()
    raise ValueError(f"No comparison operator found in: {expr_str}")


# ---------------------------------------------------------------------------
# Objective code generation
# ---------------------------------------------------------------------------


def _generate_objective(expr: str, ir: IR) -> list[str]:
    """Generate code to set objective coefficients."""
    lines = []
    lines.append("obj_coeffs = np.zeros(n_vars)")

    # Parse the objective expression to find sum(coeff * var for ...)
    # We handle: sum(p[i]*x[i] ...) and sum(f[i]*y[i] ...) + sum(c[i,j]*x[i,j] ...)
    terms = _split_top_level_sums(expr)
    for term in terms:
        lines.extend(_gen_obj_term(term, ir))

    lines.append("h.changeColsCost(n_vars, np.arange(n_vars, dtype=np.int32), obj_coeffs)")
    return lines


def _split_top_level_sums(expr: str) -> list[str]:
    """Split 'sum(...) + sum(...)' into individual sum expressions."""
    # Find top-level + between sum() calls
    results = []
    depth = 0
    current = []
    i = 0
    while i < len(expr):
        ch = expr[i]
        if ch == "(":
            depth += 1
            current.append(ch)
        elif ch == ")":
            depth -= 1
            current.append(ch)
        elif ch == "+" and depth == 0:
            results.append("".join(current).strip())
            current = []
        else:
            current.append(ch)
        i += 1
    if current:
        results.append("".join(current).strip())
    return [r for r in results if r]


def _gen_obj_term(term: str, ir: IR) -> list[str]:
    """Generate code for one sum(...) term in the objective."""
    lines = []

    # Constant objective (e.g. "0") — coefficients stay zero
    try:
        float(term.strip())
        return []
    except ValueError:
        pass

    # Parse: sum(body for idx1 in Set1 for idx2 in Set2)
    m = re.match(r"sum\((.+)\)", term, re.DOTALL)
    if not m:
        # Scalar term: bare variable reference like "x" or "c * x[i]"
        coeff, var_name, var_indices = _parse_product(term, ir)
        if var_name in ir.variables:
            offset = _offset_expr(var_name, ir)
            flat_idx = _flat_index_expr(var_name, var_indices, ir)
            coeff_c = _coeff_code(coeff, ir)
            return [f"obj_coeffs[{offset} + {flat_idx}] += {coeff_c}"]
        return [f"# TODO: non-sum objective term: {term}"]

    inner = m.group(1).strip()

    # Split off 'for ...' clauses
    body, iterators = _split_body_and_iterators(inner)

    # Parse iterators: "i in I", "j in J"
    idx_vars = []
    idx_sets = []
    for it in iterators:
        parts = it.strip().split(" in ")
        idx_vars.append(parts[0].strip())
        idx_sets.append(parts[1].strip())

    # The body might have multiple additive terms: p * x[t] + h * Inv[t] + K * z[t]
    # Split into individual terms and handle each
    body_terms = _split_additive_terms(body)

    # Generate nested loop
    indent = ""
    for iv, iset in zip(idx_vars, idx_sets, strict=True):
        size = f"n_{iset.lower()}"
        lines.append(f"{indent}for {iv} in range({size}):")
        indent += "    "

    for sign, bterm in body_terms:
        bterm = bterm.strip()
        coeff_expr, var_name, var_indices = _parse_product(bterm, ir)
        if var_name not in ir.variables:
            lines.append(f"{indent}# TODO: non-variable term: {bterm}")
            continue

        offset = _offset_expr(var_name, ir)
        flat_idx = _flat_index_expr(var_name, var_indices, ir)
        effective_coeff = coeff_expr if sign == "+" else _negate(coeff_expr)
        lines.append(
            f"{indent}obj_coeffs[{offset} + {flat_idx}] += {_coeff_code(effective_coeff, ir)}"
        )

    return lines


def _split_body_and_iterators(inner: str) -> tuple[str, list[str]]:
    """Split 'coeff * var for i in I for j in J' into body and iterator list."""
    # Find first ' for ' that's not inside brackets
    parts = re.split(r"\s+for\s+", inner)
    body = parts[0]
    iterators = parts[1:]
    return body, iterators


def _parse_product(body: str, ir: IR) -> tuple[str, str, list[str]]:
    """Parse 'c[i,j] * x[i,j]' into (coeff_expr, var_name, var_indices).

    Also handles: 'x[i,j]' (coeff=1), 'p * x[t]' (scalar param * var).
    """
    body = body.strip()

    # Try to parse as AST
    tree = ast.parse(body, mode="eval")
    node = tree.body

    return _extract_product(node, ir)


def _extract_product(node: ast.expr, ir: IR) -> tuple[str, str, list[str]]:
    """Extract (coeff_source, var_name, var_indices) from an AST node."""
    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mult):
        # A * B — figure out which side is the variable
        left_var = _find_var_ref(node.left, ir)
        right_var = _find_var_ref(node.right, ir)

        if right_var:
            vname, vindices = right_var
            coeff = ast.unparse(node.left)
            return coeff, vname, vindices
        elif left_var:
            vname, vindices = left_var
            coeff = ast.unparse(node.right)
            return coeff, vname, vindices
    elif isinstance(node, ast.Subscript):
        vref = _find_var_ref(node, ir)
        if vref:
            return "1", vref[0], vref[1]
    elif isinstance(node, ast.Name) and node.id in ir.variables:
        return "1", node.id, []

    # Fallback
    return ast.unparse(node), "_unknown", []


def _find_var_ref(node: ast.expr, ir: IR) -> tuple[str, list[str]] | None:
    """If node is var[i,j], return (var_name, [i, j])."""
    if isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name):
        name = node.value.id
        if name in ir.variables:
            indices = _extract_subscript_indices(node.slice)
            return name, indices
    return None


def _extract_subscript_indices(node: ast.expr) -> list[str]:
    """Extract index names from a subscript: x[i,j] -> ['i', 'j'], x[i] -> ['i']."""
    if isinstance(node, ast.Tuple):
        return [ast.unparse(e) for e in node.elts]
    return [ast.unparse(node)]


def _flat_index_expr(var_name: str, indices: list[str], ir: IR) -> str:
    """Compute flat index expression for var[i,j,...] given the variable's index sets."""
    var = ir.variables[var_name]
    if len(indices) == 0:
        return "0"
    if len(indices) == 1:
        return indices[0]
    # Row-major: i * n_j + j for indices [i, j] with sets [I, J]
    # General: idx[0] * prod(sizes[1:]) + idx[1] * prod(sizes[2:]) + ... + idx[-1]
    sizes = [f"n_{s.lower()}" for s in var.indices]
    parts = []
    for k, idx in enumerate(indices):
        if k < len(indices) - 1:
            suffix = " * ".join(sizes[k + 1 :])
            parts.append(f"{idx} * {suffix}")
        else:
            parts.append(idx)
    return " + ".join(parts)


def _coeff_code(coeff_expr: str, ir: IR) -> str:
    """Convert a coefficient expression to Python code referencing data dict."""
    if coeff_expr == "1":
        return "1.0"
    # Replace param references: c[i,j] -> data["c"][i,j], p -> data["p"]
    tree = ast.parse(coeff_expr, mode="eval")
    return _coeff_node_to_code(tree.body, ir)


def _coeff_node_to_code(node: ast.expr, ir: IR) -> str:
    """Convert an AST node for a coefficient to Python code."""
    if isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name):
        name = node.value.id
        if name in ir.parameters:
            indices = _extract_subscript_indices(node.slice)
            idx_str = ", ".join(indices)
            return f'data["{name}"][{idx_str}]'
    if isinstance(node, ast.Name):
        name = node.id
        if name in ir.parameters:
            return f'data["{name}"]'
    if isinstance(node, ast.BinOp):
        left = _coeff_node_to_code(node.left, ir)
        right = _coeff_node_to_code(node.right, ir)
        op_map = {ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/"}
        op_str = op_map.get(type(node.op), "+")
        return f"({left} {op_str} {right})"
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        inner = _coeff_node_to_code(node.operand, ir)
        return f"-({inner})"
    if isinstance(node, ast.Constant):
        return repr(node.value)
    # Fallback
    return ast.unparse(node)


# ---------------------------------------------------------------------------
# Constraint code generation
# ---------------------------------------------------------------------------


def _constraint_functions(ir: IR) -> list[str]:
    results = []
    for cname, constraint in ir.constraints.items():
        results.append(_gen_constraint_fn(cname, constraint, ir))
    return results


def _gen_constraint_fn(cname: str, constraint: Constraint, ir: IR) -> str:
    params = _constraint_fn_args(ir)
    lines = [
        f"def setup_{cname}_constraints({params}):",
        f'    """{constraint.expression}"""',
    ]

    # Parse the constraint expression
    lhs_str, op, rhs_str = _parse_comparison(constraint.expression)

    # Determine iteration space
    idx_vars, idx_sets = _parse_for_all(constraint.for_all)

    # Check for special case: lag references like x[t-1] or p[g,t-1]
    has_lag = bool(re.search(r"\[[^\]]*\w+\s*-\s*\d+", constraint.expression))

    if not constraint.for_all:
        lines.extend(_gen_single_constraint(lhs_str, op, rhs_str, ir, has_lag=has_lag))
    else:
        lines.extend(
            _gen_loop_constraint(
                lhs_str,
                op,
                rhs_str,
                idx_vars,
                idx_sets,
                ir,
                has_lag=has_lag,
            )
        )

    return "\n".join(lines)


def _parse_for_all(for_all: list[str]) -> tuple[list[str], list[str]]:
    idx_vars = []
    idx_sets = []
    for fa in for_all:
        parts = fa.split(" in ")
        idx_vars.append(parts[0].strip())
        idx_sets.append(parts[1].strip())
    return idx_vars, idx_sets


def _collect_constraint_terms(
    lhs_str: str,
    rhs_str: str,
    ir: IR,
    mode: Literal["normal", "with_lags", "literal"] = "normal",
) -> tuple[list[LinearTerm], list[str]]:
    """Extract and rearrange terms: all variable terms on LHS, constants on RHS."""
    lhs_terms = _extract_linear_terms(lhs_str, ir, mode=mode)
    rhs_terms = _extract_linear_terms(rhs_str, ir, mode=mode)
    all_var_terms = lhs_terms.var_terms + [t.negated() for t in rhs_terms.var_terms]
    rhs_const = rhs_terms.const_terms + [f"-({c})" for c in lhs_terms.const_terms]
    return all_var_terms, rhs_const


def _compute_rhs_code(rhs_const: list[str], ir: IR) -> str:
    """Join constant terms and substitute parameter references."""
    rhs_code = " + ".join(rhs_const) if rhs_const else "0"
    return _substitute_params(rhs_code, ir)


def _emit_single_addRows(lines: list[str], indent: str, lb: str, ub: str) -> None:
    """Append merge + addRows(1, ...) for a single-row constraint."""
    lines.append(f"{indent}row_indices, row_values = _merge_row(row_indices, row_values)")
    lines.append(f"{indent}row_lower = np.array([{lb}])")
    lines.append(f"{indent}row_upper = np.array([{ub}])")
    lines.append(f"{indent}starts = np.array([0, len(row_indices)], dtype=np.int32)")
    lines.append(f"{indent}h.addRows(1, row_lower, row_upper, len(row_indices),")
    lines.append(f"{indent}          starts, np.array(row_indices, dtype=np.int32),")
    lines.append(f"{indent}          np.array(row_values, dtype=np.float64))")


def _emit_batch_addRows(lines: list[str], *, guard: bool = False) -> None:
    """Append the final n_cons + addRows block for batch constraints."""
    lines.append("    n_cons = len(all_row_lower)")
    if guard:
        lines.append("    if n_cons > 0:")
        pad = "        "
    else:
        pad = "    "
    lines.append(f"{pad}h.addRows(n_cons, np.array(all_row_lower), np.array(all_row_upper),")
    lines.append(f"{pad}          len(all_indices), np.array(starts, dtype=np.int32),")
    lines.append(f"{pad}          np.array(all_indices, dtype=np.int32),")
    lines.append(f"{pad}          np.array(all_values, dtype=np.float64))")


def _gen_single_constraint(
    lhs_str: str, op: str, rhs_str: str, ir: IR, *, has_lag: bool = False
) -> list[str]:
    """Generate a single constraint (no for_all).

    When *has_lag* is True, uses "literal" extraction mode for constraints
    that reference specific indices (e.g. Inv[0]).
    """
    mode: Literal["normal", "literal"] = "literal" if has_lag else "normal"
    all_var_terms, rhs_const = _collect_constraint_terms(lhs_str, rhs_str, ir, mode)

    lines: list[str] = []
    lines.append("    row_indices = []")
    lines.append("    row_values = []")

    for t in all_var_terms:
        offset = _offset_expr(t.var_name, ir)
        inner_indent = "    "
        if t.iterators:
            for iv, iset in t.iterators:
                lines.append(f"{inner_indent}for {iv} in range(n_{iset.lower()}):")
                inner_indent += "    "
        flat = _flat_index_expr(t.var_name, t.var_indices, ir)
        coeff_c = _coeff_code(t.coeff, ir)
        lines.append(f"{inner_indent}row_indices.append({offset} + {flat})")
        lines.append(f"{inner_indent}row_values.append({coeff_c})")

    rhs_code = _compute_rhs_code(rhs_const, ir)
    lb, ub = _bounds_from_op(op, rhs_code)
    _emit_single_addRows(lines, "    ", lb, ub)

    return lines


def _gen_loop_constraint(
    lhs_str: str,
    op: str,
    rhs_str: str,
    idx_vars: list[str],
    idx_sets: list[str],
    ir: IR,
    *,
    has_lag: bool = False,
) -> list[str]:
    """Generate a constraint group with for_all.

    When *has_lag* is True, uses "with_lags" extraction mode, adjusts loop
    start indices, applies lag offsets to variable indices, and guards the
    final addRows with ``if n_cons > 0``.
    """
    mode: Literal["normal", "with_lags"] = "with_lags" if has_lag else "normal"
    all_var_terms, rhs_const = _collect_constraint_terms(lhs_str, rhs_str, ir, mode)

    lines: list[str] = []
    lines.append("    all_row_lower = []")
    lines.append("    all_row_upper = []")
    lines.append("    all_indices = []")
    lines.append("    all_values = []")
    lines.append("    starts = [0]")

    # Determine loop start adjustments for lag constraints
    lag_indices: set[str] = set()
    has_first = False
    if has_lag:
        lag_indices = {t.lag_index for t in all_var_terms if t.lag and t.lag_index}
        has_first = any("_first" in cn for cn in ir.constraints)

    indent = "    "
    for iv, iset in zip(idx_vars, idx_sets, strict=True):
        if has_lag:
            start = "1" if iv in lag_indices and has_first else "0"
            lines.append(f"{indent}for {iv} in range({start}, n_{iset.lower()}):")
        else:
            lines.append(f"{indent}for {iv} in range(n_{iset.lower()}):")
        indent += "    "

    lines.append(f"{indent}row_idx = []")
    lines.append(f"{indent}row_val = []")

    for t in all_var_terms:
        offset = _offset_expr(t.var_name, ir)
        inner_indent = indent
        if not has_lag and t.iterators:
            remaining = [(iv, iset) for iv, iset in t.iterators if iv not in idx_vars]
            for iv, iset in remaining:
                lines.append(f"{inner_indent}for {iv} in range(n_{iset.lower()}):")
                inner_indent += "    "

        if has_lag:
            adjusted_indices = []
            for idx in t.var_indices:
                if t.lag and idx == t.lag_index:
                    adjusted_indices.append(f"{idx} - {t.lag}")
                else:
                    adjusted_indices.append(idx)
            flat = _flat_index_expr(t.var_name, adjusted_indices, ir)
        else:
            flat = _flat_index_expr(t.var_name, t.var_indices, ir)

        coeff_c = _coeff_code(t.coeff, ir)
        lines.append(f"{inner_indent}row_idx.append({offset} + {flat})")
        lines.append(f"{inner_indent}row_val.append({coeff_c})")

    rhs_code = _compute_rhs_code(rhs_const, ir)
    lb, ub = _bounds_from_op(op, rhs_code)

    if not has_lag:
        lines.append(f"{indent}row_idx, row_val = _merge_row(row_idx, row_val)")
    lines.append(f"{indent}all_row_lower.append({lb})")
    lines.append(f"{indent}all_row_upper.append({ub})")
    lines.append(f"{indent}all_indices.extend(row_idx)")
    lines.append(f"{indent}all_values.extend(row_val)")
    lines.append(f"{indent}starts.append(len(all_indices))")

    _emit_batch_addRows(lines, guard=has_lag)

    return lines


# ---------------------------------------------------------------------------
# Linear term extraction
# ---------------------------------------------------------------------------


class LinearTerm:
    """A single variable term extracted from a linear expression."""

    __slots__ = ("coeff", "var_name", "var_indices", "iterators", "lag", "lag_index")

    def __init__(
        self,
        coeff: str,
        var_name: str,
        var_indices: list[str],
        iterators: list[tuple[str, str]] | None = None,
        lag: int = 0,
        lag_index: str = "",
    ):
        self.coeff = coeff
        self.var_name = var_name
        self.var_indices = var_indices
        self.iterators = iterators or []
        self.lag = lag
        self.lag_index = lag_index

    def negated(self) -> LinearTerm:
        return LinearTerm(
            _negate(self.coeff),
            self.var_name,
            self.var_indices,
            self.iterators,
            self.lag,
            self.lag_index,
        )


class ExtractedTerms:
    """Result of extracting linear terms from an expression."""

    __slots__ = ("var_terms", "const_terms")

    def __init__(self) -> None:
        self.var_terms: list[LinearTerm] = []
        self.const_terms: list[str] = []


def _extract_linear_terms(
    expr_str: str, ir: IR, *, mode: Literal["normal", "with_lags", "literal"] = "normal"
) -> ExtractedTerms:
    """Extract variable terms and constant terms from an expression.

    Modes:
      "normal"    — handles sum() expressions, parses iterators, lag=0.
      "with_lags" — no sum() handling, parses lag from indices (e.g. t-1).
      "literal"   — no sum() handling, no lag parsing, iterators=[].
    """
    expr_str = expr_str.strip()
    result = ExtractedTerms()

    sums = _split_additive_terms(expr_str)
    for sign, term in sums:
        term = term.strip()

        # Only "normal" mode handles sum() expressions
        if mode == "normal":
            m = re.match(r"sum\((.+)\)", term, re.DOTALL)
            if m:
                inner = m.group(1)
                body, iterators = _split_body_and_iterators(inner)
                parsed_iters = []
                for it in iterators:
                    parts = it.strip().split(" in ")
                    parsed_iters.append((parts[0].strip(), parts[1].strip()))

                coeff, vname, vindices = _parse_product(body, ir)
                if vname in ir.variables:
                    effective_coeff = coeff if sign == "+" else _negate(coeff)
                    result.var_terms.append(
                        LinearTerm(effective_coeff, vname, vindices, parsed_iters)
                    )
                else:
                    effective = term if sign == "+" else f"-({term})"
                    result.const_terms.append(effective)
                continue

        if _has_var_reference(term, ir):
            if mode == "with_lags":
                coeff, vname, vindices = _parse_product_with_lag(term, ir)
            else:
                coeff, vname, vindices = _parse_product(term, ir)

            if vname in ir.variables:
                effective_coeff = coeff if sign == "+" else _negate(coeff)

                if mode == "with_lags":
                    lag = 0
                    lag_index = ""
                    new_indices = []
                    for idx in vindices:
                        lag_m = re.match(r"(\w+)\s*-\s*(\d+)", idx)
                        if lag_m:
                            new_indices.append(lag_m.group(1))
                            lag = int(lag_m.group(2))
                            lag_index = lag_m.group(1)
                        else:
                            new_indices.append(idx)
                    result.var_terms.append(
                        LinearTerm(effective_coeff, vname, new_indices, [], lag, lag_index)
                    )
                else:
                    iters: list[tuple[str, str]] = []
                    result.var_terms.append(LinearTerm(effective_coeff, vname, vindices, iters))
            else:
                effective = term if sign == "+" else f"-({term})"
                result.const_terms.append(effective)
        else:
            effective = term if sign == "+" else f"-({term})"
            result.const_terms.append(effective)

    return result


def _split_additive_terms(expr: str) -> list[tuple[str, str]]:
    """Split expression into [(sign, term), ...] respecting parentheses and brackets."""
    terms = []
    depth = 0
    current = []
    sign = "+"
    i = 0
    while i < len(expr):
        ch = expr[i]
        if ch in "([":
            depth += 1
            current.append(ch)
        elif ch in ")]":
            depth -= 1
            current.append(ch)
        elif (ch == "+" or ch == "-") and depth == 0 and current:
            terms.append((sign, "".join(current).strip()))
            sign = ch if ch == "-" else "+"
            current = []
        else:
            current.append(ch)
        i += 1
    if current:
        terms.append((sign, "".join(current).strip()))
    return [(s, t) for s, t in terms if t]


@functools.lru_cache(maxsize=8)
def _var_reference_re(var_names: tuple[str, ...]) -> re.Pattern[str]:
    """Compile and cache a regex that matches any of the given variable names."""
    return re.compile(r"\b(?:" + "|".join(re.escape(v) for v in var_names) + r")\b")


def _has_var_reference(term: str, ir: IR) -> bool:
    """Check if term references any variable."""
    if not ir.variables:
        return False
    return bool(_var_reference_re(tuple(ir.variables)).search(term))


def _parse_product_with_lag(body: str, ir: IR) -> tuple[str, str, list[str]]:
    """Like _parse_product but handles lag expressions like Inv[t-1]."""
    body = body.strip()
    # Replace t-1 temporarily for AST parsing: Inv[t-1] isn't valid Python subscript
    # We handle it by regex
    m = re.match(r"(\w+)\[([^\]]+)\]", body)
    if m:
        name = m.group(1)
        idx_str = m.group(2)
        indices = [x.strip() for x in idx_str.split(",")]
        if name in ir.variables:
            return "1", name, indices
    # Try AST-based parsing for products like M * z[t]
    try:
        return _parse_product(body, ir)
    except SyntaxError:
        return body, "_unknown", []


def _negate(coeff: str) -> str:
    """Negate a coefficient expression using AST for correctness."""
    if coeff == "1":
        return "-1"
    if coeff == "-1":
        return "1"
    try:
        tree = ast.parse(coeff, mode="eval")
        node = tree.body
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
            return ast.unparse(node.operand)
        negated = ast.UnaryOp(op=ast.USub(), operand=node)
        return ast.unparse(negated)
    except SyntaxError:
        return f"-({coeff})"


def _bounds_from_op(op: str, rhs_code: str) -> tuple[str, str]:
    """Return (lower_bound, upper_bound) expressions for a comparison."""
    if op == "<=":
        return "-np.inf", rhs_code
    elif op == ">=":
        return rhs_code, "np.inf"
    elif op == "==":
        return rhs_code, rhs_code
    raise ValueError(f"Unknown operator: {op}")


def _substitute_params(code: str, ir: IR) -> str:
    """Replace parameter names with data dict lookups."""
    # Indexed params need individual replacements (capture group references the name)
    for pname in sorted(ir.parameters.keys(), key=len, reverse=True):
        param = ir.parameters[pname]
        if param.indices:
            code = re.sub(
                rf"\b{re.escape(pname)}\[([^\]]+)\]",
                f'data["{pname}"][\\1]',
                code,
            )
    # Consolidate all scalar params into a single regex
    scalar_params = sorted(
        (p for p, param in ir.parameters.items() if not param.indices),
        key=len,
        reverse=True,
    )
    if scalar_params:
        pattern = r"\b(" + "|".join(re.escape(p) for p in scalar_params) + r")\b"
        code = re.sub(pattern, lambda m: f'data["{m.group(1)}"]', code)
    return code


# ---------------------------------------------------------------------------
# Solve + report + main
# ---------------------------------------------------------------------------


def _parse_highs_value_fn() -> str:
    return '''def _parse_highs_value(s):
    """Convert a string option value to the appropriate Python type."""
    if s.lower() in ("true", "false"):
        return s.lower() == "true"
    try:
        return int(s)
    except ValueError:
        pass
    try:
        return float(s)
    except ValueError:
        pass
    return s'''


def _solve_and_report(ir: IR) -> str:
    lines = [
        "def solve_and_report(h: highspy.Highs, data: dict, highs_opts=None):",
        "    if highs_opts:",
        "        for opt in highs_opts:",
        '            name, _, value = opt.partition("=")',
        "            if not name or not value:",
        '                sys.exit(f"Invalid --opt format: {opt!r} (expected NAME=VALUE)")',
        "            h.setOptionValue(name, _parse_highs_value(value))",
        "    h.run()",
        "    model_status = h.getModelStatus()",
        "",
        "    if model_status == highspy.HighsModelStatus.kInfeasible:",
        '        print("Status: infeasible")',
        "        _status, iis_data = h.getIis()",
        "        if _status == highspy.HighsStatus.kOk:",
        '            print("IIS (conflicting constraints):")',
        "            for row in iis_data.row_index_:",
        '                name = CONSTRAINT_NAMES.get(row, f"row_{row}")',
        '                print(f"  - {name} (row {row})")',
        "            for col in iis_data.col_index_:",
        '                name = VARIABLE_NAMES.get(col, f"col_{col}")',
        '                print(f"  - {name} bound (col {col})")',
        "        sys.exit(1)",
        "",
        "    if model_status != highspy.HighsModelStatus.kOptimal:",
        '        print(f"Status: {model_status.name}")',
        "        sys.exit(1)",
        "",
        '    obj = h.getInfoValue("objective_function_value")[1]',
        "    sol = np.array(h.getSolution().col_value)",
        "",
    ]

    # Define set sizes in this function too
    for sname in ir.sets:
        lines.append(f'    n_{sname.lower()} = len(data["{sname}"])')
    lines.append("")

    lines.append('    print("Status: optimal")')
    lines.append('    print(f"Objective: {obj:.4f}")')
    lines.append('    print("Solution:")')

    # Print solution for each variable group
    for vname, var in ir.variables.items():
        size_expr = _var_size_expr(var, ir)
        offset = _offset_expr(vname, ir)
        if var.indices:
            set_sizes = [f"n_{s.lower()}" for s in var.indices]
            idx_names = [s.lower() for s in var.indices]
            shape = ", ".join(set_sizes)
            lines.append(f"    {vname}_sol = sol[{offset}:{offset} + {size_expr}].reshape({shape})")

            # Nested loop to print nonzero values
            indent = "    "
            for idx, sname in zip(idx_names, var.indices, strict=True):
                lines.append(f"{indent}for {idx} in range(n_{sname.lower()}):")
                indent += "    "

            idx_str = ", ".join(idx_names)
            # Build label with string concatenation to avoid nested f-string issues
            label_parts = []
            for idx, sname in zip(idx_names, var.indices, strict=True):
                label_parts.append(f'str(data["{sname}"][{idx}])')
            label_expr = ' + "," + '.join(label_parts)
            lines.append(f"{indent}val = {vname}_sol[{idx_str}]")
            lines.append(f"{indent}if abs(val) > 1e-6:")
            lines.append(f"{indent}    label = {label_expr}")
            lines.append(f'{indent}    print(f"  {vname}[{{label}}] = {{val:.4f}}")')
        else:
            lines.append(f"    val = sol[{offset}]")
            lines.append(f'    print(f"  {vname} = {{val:.4f}}")')

    return "\n".join(lines)


def _main(ir: IR) -> str:
    return f'''if __name__ == "__main__":
    import textwrap

    parser = argparse.ArgumentParser(
        description="{ir.name} solver — generated by md2mip",
        epilog=textwrap.dedent("""\
            Data file format (YAML or JSON):
{_indent_epilog_entries(ir, 14)}
        """),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("data", help="Path to data file (YAML or JSON)")
    parser.add_argument("--opt", action="append", default=[],
                        metavar="NAME=VALUE", help="HiGHS option (repeatable)")
    args = parser.parse_args()

    data = load_data(args.data)
    h = build_model(data)
    solve_and_report(h, data, highs_opts=args.opt)'''


def _build_epilog(ir: IR) -> list[str]:
    """Build epilog lines describing the data file format."""
    lines = []
    for sname, sdef in ir.sets.items():
        desc = sdef.description or "set of labels"
        lines.append(f"{sname}: list of {desc}")
    for pname, param in ir.parameters.items():
        if param.indices:
            idx_desc = " x ".join(param.indices)
            desc = f"array indexed by {idx_desc}"
        else:
            desc = "scalar"
        if param.description:
            desc += f" ({param.description})"
        lines.append(f"{pname}: {desc}")
    return lines


def _indent_epilog_entries(ir: IR, base_indent: int) -> str:
    """Return epilog entries indented for embedding in the f-string."""
    lines = _build_epilog(ir)
    pad = " " * base_indent
    return "\n".join(f"{pad}{line}" for line in lines)
