"""Source-of-truth constants for the supported modified-RNA residues.

The PDB->modXNA base-code map is deliberately NOT a 1:1 name mapping: modXNA's
own `PSU` fragment is a thiouracil (real pseudouridine is `PUU`) and its `2MG`
is an N2,N2-dimethyl (mono-methyl guanosine is `M2G`).
"""

RESIDUES: tuple[str, ...] = (
    "5MC", "PSU", "5MU", "2MG", "1MG", "2MA", "1MA", "4AC", "6IA",
)

TERMINI: tuple[str, ...] = ("internal", "5prime", "3prime")

MODXNA_BASE: dict[str, str] = {
    "5MC": "M5C", "PSU": "PUU", "5MU": "M5U", "2MG": "M2G", "1MG": "M1G",
    "2MA": "M2A", "1MA": "M1A", "4AC": "A4C", "6IA": "I6A",
}

# net formal charge target per form (modXNA convention)
NET_CHARGE: dict[str, float] = {
    "internal": -1.0, "5prime": -0.320348, "3prime": -0.679652,
}

# modXNA generation recipe per form: (backbone fragment, extra flag)
_FORM_RECIPE: dict[str, tuple[str, str | None]] = {
    "internal": ("RPO", None),
    "5prime": ("5PO", "--5cap"),
    "3prime": ("RPO", "--3cap"),
}

# ribose sugar fragment, same for every form
SUGAR: str = "RC3"

# unit-name suffix per form
_SUFFIX: dict[str, str] = {"internal": "", "5prime": "5", "3prime": "3"}

# verified internal-form atom counts; caps derive uniformly (backbone-only change)
_INTERNAL_ATOMS: dict[str, int] = {
    "5MC": 34, "PSU": 30, "5MU": 33, "2MG": 37, "1MG": 37,
    "2MA": 36, "1MA": 37, "4AC": 36, "6IA": 46,
}
# 5' cap removes P, OP1, OP2 and adds HO5' (net -2); 3' cap adds HO3' (net +1)
_ATOM_DELTA: dict[str, int] = {"internal": 0, "5prime": -2, "3prime": 1}


def unit_name(resname: str, terminus: str = "internal") -> str:
    """OFF unit name for a residue form, e.g. '5MC', '5MC5', '5MC3'."""
    if resname not in MODXNA_BASE:
        raise KeyError(f"unsupported residue {resname!r}")
    if terminus not in TERMINI:
        raise KeyError(f"unknown terminus {terminus!r}")
    return resname + _SUFFIX[terminus]


def expected_atom_count(resname: str, terminus: str = "internal") -> int:
    """Expected atom count for a residue form."""
    if resname not in _INTERNAL_ATOMS:
        raise KeyError(f"unsupported residue {resname!r}")
    if terminus not in TERMINI:
        raise KeyError(f"unknown terminus {terminus!r}")
    return _INTERNAL_ATOMS[resname] + _ATOM_DELTA[terminus]
