import pandas as pd
from biotite.structure.sse import annotate_sse

from apb.types import ProteinStructure

FEATURE_NAME = "secondary_structure"

# Secondary structure (3-state): [a]lpha helix, [b]eta sheet, [c]oil.
ss3_dtype = pd.CategoricalDtype(categories=["a", "b", "c"])


def calc_3_state_secondary_structure(structure: ProteinStructure) -> pd.Series:
    sse_array = annotate_sse(structure)
    return pd.Series(sse_array, dtype=ss3_dtype, name=FEATURE_NAME)


__all__ = [
    "ss3_dtype",
    "calc_3_state_secondary_structure",
]
