"""
Biochemistry constants and enumerations.

Defines atoms, residues, elements, and nucleotide structures.

Atom enums are accessed via Residue members:
    Residue.A.C3p.value    # atom index for C3' in adenosine
    Residue.ALA.CA.value   # atom index for CA in alanine
"""

import numpy as np

from ._generated_elements import Element, ELEMENT_NAMES
from ._generated_residues import Residue, CIF_RESIDUE_NAMES
from ._generated_atoms import (
    # Combined enums (for bulk operations)
    RibonucleicAcid,
    RibonucleicAcidNoPrefix,
    DeoxyribonucleicAcid,
    ModifiedNucleotides,
    AminoAcids,
    # Reverse lookup
    ATOM_NAMES,
)

# Dihedral ownership arrays - generated by codegen
# These may not exist before first build, so provide fallback empty arrays
try:
    from ._generated_dihedrals import ATOM_DIHEDRAL_TYPE, ATOM_DIHEDRAL_REFS
except ImportError:
    # Fallback for development before codegen runs
    ATOM_DIHEDRAL_TYPE = np.array([], dtype=np.int8)
    ATOM_DIHEDRAL_REFS = np.zeros((0, 3, 2), dtype=np.int8)

# Canonical Z-matrix reference arrays - generated by codegen
# These define canonical references for deterministic Z-matrix construction
try:
    from ._generated_zmatrix import ATOM_CANONICAL_REFS, ATOM_HAS_CANONICAL_REFS
except ImportError:
    # Fallback for development before codegen runs
    ATOM_CANONICAL_REFS = np.zeros((0, 6), dtype=np.int16)
    ATOM_HAS_CANONICAL_REFS = np.array([], dtype=bool)

# DihedralType enum - generated by codegen (IndexEnum for array methods)
from ._generated_dihedraltypes import (
    DihedralType,
    PROTEIN_BACKBONE,
    RNA_BACKBONE,
    RNA_GLYCOSIDIC,
    RNA_BACKBONE_EXTENDED,
    PROTEIN_SIDECHAIN,
    DIHEDRAL_NAME_TO_TYPE,
    DIHEDRAL_ATOMS,
)
from .constants import (
    # Flat IndexEnums (prefixed atom names)
    Backbone,
    Nucleobase,
    Phosphate,
    Sidechain,
    # Hierarchical atom groups (chemical identity)
    Sugar,
    PhosphateGroup,
    PurineImidazole,
    PurinePyrimidine,
    PurineBase,
    PyrimidineBase,
)
from .linking import (
    LinkingDefinition,
    NUCLEIC_ACID_LINK,
    PEPTIDE_LINK,
    LINKING_BY_TYPE,
)
from ._generated_molecule import Molecule, molecule_type

# =============================================================================
# VOCABULARY SIZES (for embedding layers)
# =============================================================================

# Number of element types (max index + 1)
NUM_ELEMENTS: int = max(e.value for e in Element) + 1

# Number of residue types
NUM_RESIDUES: int = len(Residue)

# Number of atom types (max index + 1)
NUM_ATOMS: int = max(ATOM_NAMES.keys()) + 1


__all__ = [
    # Vocabulary sizes
    "NUM_ELEMENTS",
    "NUM_RESIDUES",
    "NUM_ATOMS",
    # Reverse lookups
    "ATOM_NAMES",
    "ELEMENT_NAMES",
    "CIF_RESIDUE_NAMES",
    # Dihedral ownership arrays
    "ATOM_DIHEDRAL_TYPE",
    "ATOM_DIHEDRAL_REFS",
    # Canonical Z-matrix reference arrays
    "ATOM_CANONICAL_REFS",
    "ATOM_HAS_CANONICAL_REFS",
    # Elements
    "Element",
    # Residues (access atoms via Residue.X.atom_name)
    "Residue",
    # Combined atom enums (for bulk operations)
    "RibonucleicAcid",
    "RibonucleicAcidNoPrefix",
    "DeoxyribonucleicAcid",
    "ModifiedNucleotides",
    "AminoAcids",
    # Atom group constants (flat IndexEnums with prefixed names)
    "Backbone",
    "Nucleobase",
    "Phosphate",
    "Sidechain",
    # Hierarchical atom groups (chemical identity, HierarchicalEnum)
    "Sugar",
    "PhosphateGroup",
    "PurineImidazole",
    "PurinePyrimidine",
    "PurineBase",
    "PyrimidineBase",
    # Linking
    "LinkingDefinition",
    "NUCLEIC_ACID_LINK",
    "PEPTIDE_LINK",
    "LINKING_BY_TYPE",
    # Molecule types
    "Molecule",
    "molecule_type",
    # Dihedral types (IndexEnum)
    "DihedralType",
    "PROTEIN_BACKBONE",
    "RNA_BACKBONE",
    "RNA_GLYCOSIDIC",
    "RNA_BACKBONE_EXTENDED",
    "PROTEIN_SIDECHAIN",
    "DIHEDRAL_NAME_TO_TYPE",
    "DIHEDRAL_ATOMS",
]
