"""Degrees-of-freedom counting via Jacobian rank."""

from __future__ import annotations

from typing import List

import numpy as np

from .expr import Expr
from .params import ParamTable


def count_dof(
    residuals: List[Expr],
    params: ParamTable,
    rank_tol: float = 1e-8,
    jac_exprs: "List[List[Expr]] | None" = None,
) -> int:
    """Compute DOF = n_free_params - rank(Jacobian).

    Evaluates the Jacobian numerically at the current parameter values
    and computes its rank via SVD.

    When *jac_exprs* is provided, reuses the pre-built symbolic
    Jacobian instead of re-differentiating every residual.
    """
    free = params.free_names()
    n_free = len(free)
    n_res = len(residuals)

    if n_free == 0:
        return 0
    if n_res == 0:
        return n_free

    env = params.get_env()

    J = np.empty((n_res, n_free))
    if jac_exprs is not None:
        for i in range(n_res):
            for j in range(n_free):
                J[i, j] = jac_exprs[i][j].eval(env)
    else:
        for i, r in enumerate(residuals):
            for j, name in enumerate(free):
                J[i, j] = r.diff(name).simplify().eval(env)

    if J.size == 0:
        return n_free

    sv = np.linalg.svd(J, compute_uv=False)
    rank = int(np.sum(sv > rank_tol))

    return n_free - rank