"""Family / link function inference based on DV type.

Ports inferFamilyLinkFunctions.R.
"""

from __future__ import annotations

from ..variables import (
    AbstractVariable,
    Continuous,
    Counts,
    OrderedCategories,
    UnorderedCategories,
    Categories,
)

# ---------------------------------------------------------------------------
# Link functions per family (matching R code)
# ---------------------------------------------------------------------------

_LINK_FUNCTIONS: dict[str, list[str]] = {
    "Binomial": ["logit", "probit", "cauchit", "log", "cloglog"],
    "Gamma": ["inverse", "identity", "log"],
    "Gaussian": ["identity", "log", "inverse"],
    "Inverse Gaussian": ["1/mu^2", "inverse", "identity", "log"],
    "Negative Binomial": ["log", "sqrt", "identity"],
    "Poisson": ["log", "identity", "sqrt"],
    "Multinomial": [],
}


def _family_candidates(dv: AbstractVariable) -> list[str]:
    """Determine candidate family functions from DV type."""
    if isinstance(dv, Continuous):
        return ["Inverse Gaussian", "Gamma", "Gaussian"]

    if isinstance(dv, Counts):
        return ["Poisson", "Negative Binomial"]

    if isinstance(dv, OrderedCategories):
        if dv.cardinality == 2:
            return ["Binomial"]
        return ["Multinomial", "Inverse Gaussian", "Gamma", "Gaussian"]

    if isinstance(dv, UnorderedCategories):
        if dv.cardinality == 2:
            return ["Binomial"]
        return ["Multinomial"]

    # Fallback for bare Categories
    if isinstance(dv, Categories):
        return ["Multinomial"]

    raise ValueError(f"Cannot infer family for {type(dv).__name__}")


def infer_family_link_functions(dv: AbstractVariable) -> dict[str, list[str]]:
    """Return ``{family: [link, ...]}`` candidates for *dv*."""
    families = _family_candidates(dv)
    return {f: list(_LINK_FUNCTIONS.get(f, [])) for f in families}
