106def calculate_sasa(
107    coords: NDArray[np.float64],
108    radii: NDArray[np.float64],
109    *,
110    algorithm: Literal["sr", "lr"] = "sr",
111    n_points: int = 100,
112    n_slices: int = 20,
113    probe_radius: float = 1.4,
114    n_threads: int = 0,
115    use_bitmask: bool = False,
116    bitmask_correction: bool = False,
117    bitmask_correction_coeff: float | None = None,
118) -> SasaResult:
119    """Calculate Solvent Accessible Surface Area (SASA).
120
121    Args:
122        coords: Atom coordinates as (N, 3) array.
123        radii: Atom radii as (N,) array.
124        algorithm: Algorithm to use: "sr" (Shrake-Rupley) or "lr" (Lee-Richards).
125        n_points: Number of test points per atom (for SR algorithm). Default: 100.
126        n_slices: Number of slices per atom (for LR algorithm). Default: 20.
127        probe_radius: Water probe radius in Angstroms. Default: 1.4.
128        n_threads: Number of threads to use. 0 = auto-detect. Default: 0.
129        use_bitmask: Use bitmask LUT optimization for SR algorithm.
130            Supports n_points 1..1024. Default: False.
131        bitmask_correction: Apply experimental bitmask exposed-fraction
132            correction. Requires use_bitmask=True. Default: False.
133        bitmask_correction_coeff: Optional non-negative correction coefficient.
134            Defaults to the library's experimental coefficient when omitted.
135
136    Returns:
137        SasaResult containing total_area and per-atom atom_areas.
138
139    Raises:
140        ValueError: If input arrays have invalid shapes or calculation fails.
141
142    Example:
143        >>> import numpy as np
144        >>> from zsasa import calculate_sasa
145        >>>
146        >>> # Two atoms
147        >>> coords = np.array([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]])
148        >>> radii = np.array([1.5, 1.5])
149        >>> result = calculate_sasa(coords, radii)
150        >>> print(f"Total: {result.total_area:.2f}")
151    """
152    # Validate parameters before loading the C library so invalid Python inputs
153    # never reach CFFI conversion or native code.
154    n_points = _validate_uint32_param("n_points", n_points)
155    n_slices = _validate_uint32_param("n_slices", n_slices)
156    probe_radius = _validate_positive_finite_float("probe_radius", probe_radius)
157    n_threads = _validate_size_t_param("n_threads", n_threads)
158    if bitmask_correction and not use_bitmask:
159        msg = "bitmask_correction=True requires use_bitmask=True"
160        raise ValueError(msg)
161    if bitmask_correction_coeff is not None and (
162        bitmask_correction_coeff < 0 or not np.isfinite(bitmask_correction_coeff)
163    ):
164        msg = f"bitmask_correction_coeff must be non-negative, got {bitmask_correction_coeff}"
165        raise ValueError(msg)
166
167    # Validate bitmask constraints
168    if use_bitmask:
169        use_bitmask = _validate_bitmask_params(algorithm, n_points, strict=bitmask_correction)
170
171    # Validate and convert inputs
172    coords = np.ascontiguousarray(coords, dtype=np.float64)
173    radii = np.ascontiguousarray(radii, dtype=np.float64)
174
175    if coords.ndim != 2 or coords.shape[1] != 3:
176        msg = f"coords must be (N, 3) array, got shape {coords.shape}"
177        raise ValueError(msg)
178
179    n_atoms = coords.shape[0]
180    if radii.shape != (n_atoms,):
181        msg = f"radii must be ({n_atoms},) array, got shape {radii.shape}"
182        raise ValueError(msg)
183
184    _validate_finite_array("coords", coords)
185    _validate_finite_array("radii", radii)
186
187    if np.any(radii < 0):
188        msg = "All radii must be non-negative"
189        raise ValueError(msg)
190
191    ffi, lib = _get_lib()
192
193    # Extract x, y, z as contiguous arrays
194    x = np.ascontiguousarray(coords[:, 0])
195    y = np.ascontiguousarray(coords[:, 1])
196    z = np.ascontiguousarray(coords[:, 2])
197
198    # Allocate output arrays
199    atom_areas = np.zeros(n_atoms, dtype=np.float64)
200    total_area = ffi.new("double*")
201
202    # Get cffi pointers from numpy arrays
203    x_ptr = ffi.cast("double*", x.ctypes.data)
204    y_ptr = ffi.cast("double*", y.ctypes.data)
205    z_ptr = ffi.cast("double*", z.ctypes.data)
206    radii_ptr = ffi.cast("double*", radii.ctypes.data)
207    areas_ptr = ffi.cast("double*", atom_areas.ctypes.data)
208
209    # Call the appropriate function
210    if use_bitmask:
211        if bitmask_correction:
212            coeff = 0.020 if bitmask_correction_coeff is None else bitmask_correction_coeff
213            result = lib.zsasa_calc_sr_bitmask_corrected(
214                x_ptr,
215                y_ptr,
216                z_ptr,
217                radii_ptr,
218                n_atoms,
219                n_points,
220                probe_radius,
221                n_threads,
222                coeff,
223                areas_ptr,
224                total_area,
225            )
226        else:
227            result = lib.zsasa_calc_sr_bitmask(
228                x_ptr,
229                y_ptr,
230                z_ptr,
231                radii_ptr,
232                n_atoms,
233                n_points,
234                probe_radius,
235                n_threads,
236                areas_ptr,
237                total_area,
238            )
239    elif algorithm == "sr":
240        result = lib.zsasa_calc_sr(
241            x_ptr,
242            y_ptr,
243            z_ptr,
244            radii_ptr,
245            n_atoms,
246            n_points,
247            probe_radius,
248            n_threads,
249            areas_ptr,
250            total_area,
251        )
252    elif algorithm == "lr":
253        result = lib.zsasa_calc_lr(
254            x_ptr,
255            y_ptr,
256            z_ptr,
257            radii_ptr,
258            n_atoms,
259            n_slices,
260            probe_radius,
261            n_threads,
262            areas_ptr,
263            total_area,
264        )
265    else:
266        msg = f"Unknown algorithm: {algorithm}. Use 'sr' or 'lr'."
267        raise ValueError(msg)
268
269    # Check for errors
270    if result == ZSASA_ERROR_INVALID_INPUT:
271        msg = "Invalid input to SASA calculation"
272        raise ValueError(msg)
273    elif result == ZSASA_ERROR_OUT_OF_MEMORY:
274        msg = "Out of memory during SASA calculation"
275        raise MemoryError(msg)
276    elif result == ZSASA_ERROR_CALCULATION:
277        msg = "Error during SASA calculation"
278        raise RuntimeError(msg)
279    elif result == ZSASA_ERROR_UNSUPPORTED_N_POINTS:
280        msg = (
281            f"Unsupported n_points for bitmask: {n_points}. "
282            f"Must be {_BITMASK_MIN_N_POINTS}..{_BITMASK_MAX_N_POINTS}"
283        )
284        raise ValueError(msg)
285    elif result != ZSASA_OK:
286        msg = f"Unknown error code: {result}"
287        raise RuntimeError(msg)
288
289    return SasaResult(
290        total_area=total_area[0],
291        atom_areas=atom_areas,
292    )

325def calculate_sasa_batch(
326    coordinates: NDArray[np.floating],
327    radii: NDArray[np.floating],
328    *,
329    algorithm: Literal["sr", "lr"] = "sr",
330    n_points: int = 100,
331    n_slices: int = 20,
332    probe_radius: float = 1.4,
333    n_threads: int = 0,
334    precision: Literal["f64", "f32"] = "f64",
335    use_bitmask: bool = False,
336    bitmask_correction: bool = False,
337    bitmask_correction_coeff: float | None = None,
338) -> BatchSasaResult:
339    """Calculate SASA for multiple frames (batch processing).
340
341    Optimized for MD trajectory analysis where the same atoms are processed
342    across multiple frames. Parallelizes across frames for maximum performance.
343
344    Args:
345        coordinates: Atom coordinates as (n_frames, n_atoms, 3) array in Angstroms.
346        radii: Atom radii as (n_atoms,) array in Angstroms.
347        algorithm: Algorithm to use: "sr" (Shrake-Rupley) or "lr" (Lee-Richards).
348        n_points: Number of test points per atom (for SR algorithm). Default: 100.
349        n_slices: Number of slices per atom (for LR algorithm). Default: 20.
350        probe_radius: Water probe radius in Angstroms. Default: 1.4.
351        n_threads: Number of threads to use. 0 = auto-detect. Default: 0.
352        precision: Internal calculation precision: "f64" (default, higher precision)
353            or "f32" (matches RustSASA/mdsasa-bolt for comparison). Default: "f64".
354        use_bitmask: Use bitmask LUT optimization for SR algorithm.
355            Supports n_points 1..1024. Default: False.
356        bitmask_correction: Apply experimental bitmask exposed-fraction
357            correction. Requires use_bitmask=True. Default: False.
358        bitmask_correction_coeff: Optional non-negative correction coefficient.
359            Defaults to the library's experimental coefficient when omitted.
360
361    Returns:
362        BatchSasaResult containing per-atom SASA for all frames.
363
364    Raises:
365        ValueError: If input arrays have invalid shapes or calculation fails.
366
367    Example:
368        >>> import numpy as np
369        >>> from zsasa import calculate_sasa_batch
370        >>>
371        >>> # 10 frames, 100 atoms
372        >>> coords = np.random.randn(10, 100, 3).astype(np.float32)
373        >>> radii = np.full(100, 1.5, dtype=np.float32)
374        >>> result = calculate_sasa_batch(coords, radii)
375        >>> print(f"Shape: {result.atom_areas.shape}")  # (10, 100)
376        >>> print(f"Total SASA per frame: {result.total_areas}")
377    """
378    # Validate parameters before loading the C library so invalid Python inputs
379    # never reach CFFI conversion or native code.
380    n_points = _validate_uint32_param("n_points", n_points)
381    n_slices = _validate_uint32_param("n_slices", n_slices)
382    probe_radius = _validate_positive_finite_float("probe_radius", probe_radius)
383    n_threads = _validate_size_t_param("n_threads", n_threads)
384    if bitmask_correction and not use_bitmask:
385        msg = "bitmask_correction=True requires use_bitmask=True"
386        raise ValueError(msg)
387    if bitmask_correction_coeff is not None and (
388        bitmask_correction_coeff < 0 or not np.isfinite(bitmask_correction_coeff)
389    ):
390        msg = f"bitmask_correction_coeff must be non-negative, got {bitmask_correction_coeff}"
391        raise ValueError(msg)
392
393    # Validate bitmask constraints
394    if use_bitmask:
395        use_bitmask = _validate_bitmask_params(algorithm, n_points, strict=bitmask_correction)
396
397    # Validate and convert inputs
398    coordinates = np.ascontiguousarray(coordinates, dtype=np.float32)
399    radii = np.ascontiguousarray(radii, dtype=np.float32)
400
401    if coordinates.ndim != 3 or coordinates.shape[2] != 3:
402        msg = f"coordinates must be (n_frames, n_atoms, 3) array, got shape {coordinates.shape}"
403        raise ValueError(msg)
404
405    n_frames = coordinates.shape[0]
406    n_atoms = coordinates.shape[1]
407
408    if radii.shape != (n_atoms,):
409        msg = f"radii must be ({n_atoms},) array, got shape {radii.shape}"
410        raise ValueError(msg)
411
412    _validate_finite_array("coordinates", coordinates)
413    _validate_finite_array("radii", radii)
414
415    if np.any(radii < 0):
416        msg = "All radii must be non-negative"
417        raise ValueError(msg)
418
419    ffi, lib = _get_lib()
420
421    # Allocate output array
422    atom_areas = np.zeros((n_frames, n_atoms), dtype=np.float32)
423
424    # Get cffi pointers from numpy arrays
425    coords_ptr = ffi.cast("float*", coordinates.ctypes.data)
426    radii_ptr = ffi.cast("float*", radii.ctypes.data)
427    areas_ptr = ffi.cast("float*", atom_areas.ctypes.data)
428
429    # Call the appropriate batch function based on algorithm, precision, and bitmask
430    if use_bitmask:
431        coeff = 0.020 if bitmask_correction_coeff is None else bitmask_correction_coeff
432        if precision == "f32":
433            if bitmask_correction:
434                result = lib.zsasa_calc_sr_batch_bitmask_f32_corrected(
435                    coords_ptr,
436                    n_frames,
437                    n_atoms,
438                    radii_ptr,
439                    n_points,
440                    probe_radius,
441                    n_threads,
442                    coeff,
443                    areas_ptr,
444                )
445            else:
446                result = lib.zsasa_calc_sr_batch_bitmask_f32(
447                    coords_ptr,
448                    n_frames,
449                    n_atoms,
450                    radii_ptr,
451                    n_points,
452                    probe_radius,
453                    n_threads,
454                    areas_ptr,
455                )
456        else:
457            if bitmask_correction:
458                result = lib.zsasa_calc_sr_batch_bitmask_corrected(
459                    coords_ptr,
460                    n_frames,
461                    n_atoms,
462                    radii_ptr,
463                    n_points,
464                    probe_radius,
465                    n_threads,
466                    coeff,
467                    areas_ptr,
468                )
469            else:
470                result = lib.zsasa_calc_sr_batch_bitmask(
471                    coords_ptr,
472                    n_frames,
473                    n_atoms,
474                    radii_ptr,
475                    n_points,
476                    probe_radius,
477                    n_threads,
478                    areas_ptr,
479                )
480    elif precision == "f64":
481        # Default: f32 I/O with f64 internal precision
482        if algorithm == "sr":
483            result = lib.zsasa_calc_sr_batch(
484                coords_ptr,
485                n_frames,
486                n_atoms,
487                radii_ptr,
488                n_points,
489                probe_radius,
490                n_threads,
491                areas_ptr,
492            )
493        elif algorithm == "lr":
494            result = lib.zsasa_calc_lr_batch(
495                coords_ptr,
496                n_frames,
497                n_atoms,
498                radii_ptr,
499                n_slices,
500                probe_radius,
501                n_threads,
502                areas_ptr,
503            )
504        else:
505            msg = f"Unknown algorithm: {algorithm}. Use 'sr' or 'lr'."
506            raise ValueError(msg)
507    elif precision == "f32":
508        # Pure f32 precision (matches RustSASA/mdsasa-bolt)
509        if algorithm == "sr":
510            result = lib.zsasa_calc_sr_batch_f32(
511                coords_ptr,
512                n_frames,
513                n_atoms,
514                radii_ptr,
515                n_points,
516                probe_radius,
517                n_threads,
518                areas_ptr,
519            )
520        elif algorithm == "lr":
521            result = lib.zsasa_calc_lr_batch_f32(
522                coords_ptr,
523                n_frames,
524                n_atoms,
525                radii_ptr,
526                n_slices,
527                probe_radius,
528                n_threads,
529                areas_ptr,
530            )
531        else:
532            msg = f"Unknown algorithm: {algorithm}. Use 'sr' or 'lr'."
533            raise ValueError(msg)
534    else:
535        msg = f"Unknown precision: {precision}. Use 'f64' or 'f32'."
536        raise ValueError(msg)
537
538    # Check for errors
539    if result == ZSASA_ERROR_INVALID_INPUT:
540        msg = "Invalid input to batch SASA calculation"
541        raise ValueError(msg)
542    elif result == ZSASA_ERROR_OUT_OF_MEMORY:
543        msg = "Out of memory during batch SASA calculation"
544        raise MemoryError(msg)
545    elif result == ZSASA_ERROR_CALCULATION:
546        msg = "Error during batch SASA calculation"
547        raise RuntimeError(msg)
548    elif result == ZSASA_ERROR_UNSUPPORTED_N_POINTS:
549        msg = (
550            f"Unsupported n_points for bitmask: {n_points}. "
551            f"Must be {_BITMASK_MIN_N_POINTS}..{_BITMASK_MAX_N_POINTS}"
552        )
553        raise ValueError(msg)
554    elif result != ZSASA_OK:
555        msg = f"Unknown error code: {result}"
556        raise RuntimeError(msg)
557
558    return BatchSasaResult(atom_areas=atom_areas)

 93@dataclass
 94class SasaResult:
 95    """Result of SASA calculation.
 96
 97    Attributes:
 98        total_area: Total solvent accessible surface area in Ų.
 99        atom_areas: Per-atom SASA values in Ų.
100    """
101
102    total_area: float
103    atom_areas: NDArray[np.float64]

295@dataclass
296class BatchSasaResult:
297    """Result of batch SASA calculation for multiple frames.
298
299    Attributes:
300        atom_areas: Per-atom SASA values for all frames, shape (n_frames, n_atoms).
301                    Values are in Angstrom² (Ų).
302    """
303
304    atom_areas: NDArray[np.float32]
305
306    @property
307    def n_frames(self) -> int:
308        """Number of frames."""
309        return self.atom_areas.shape[0]
310
311    @property
312    def n_atoms(self) -> int:
313        """Number of atoms."""
314        return self.atom_areas.shape[1]
315
316    @property
317    def total_areas(self) -> NDArray[np.float32]:
318        """Total SASA per frame, shape (n_frames,)."""
319        return self.atom_areas.sum(axis=1)
320
321    def __repr__(self) -> str:
322        return f"BatchSasaResult(n_frames={self.n_frames}, n_atoms={self.n_atoms})"

 71def process_directory(
 72    input_dir: str | Path,
 73    *,
 74    output_dir: str | Path | None = None,
 75    algorithm: Literal["sr", "lr"] = "sr",
 76    n_points: int = 100,
 77    n_slices: int = 20,
 78    probe_radius: float = 1.4,
 79    n_threads: int = 0,
 80    classifier: ClassifierType | None = ClassifierType.CCD,
 81    include_hydrogens: bool = False,
 82    include_hetatm: bool = False,
 83) -> BatchDirResult:
 84    """Process all supported structure files in a directory for SASA calculation.
 85
 86    Supported formats: PDB (.pdb), mmCIF (.cif, .mmcif), BinaryCIF (.bcif),
 87    PDB/ENT (.ent), JSON (.json), and their gzip- or zstd-compressed variants
 88    (.gz, .zst).
 89
 90    Args:
 91        input_dir: Path to directory containing structure files.
 92        output_dir: Optional path for per-file output. None = no file output.
 93        algorithm: Algorithm to use: "sr" (Shrake-Rupley) or "lr" (Lee-Richards).
 94        n_points: Number of test points per atom (SR only; ignored for LR).
 95            Default: 100.
 96        n_slices: Number of slices per atom (LR only; ignored for SR).
 97            Default: 20.
 98        probe_radius: Water probe radius in Angstroms. Default: 1.4.
 99        n_threads: Number of threads to use. 0 = auto-detect. Default: 0.
100        classifier: Classifier for radius assignment. None = use input radii.
101            Default: ClassifierType.CCD.
102        include_hydrogens: Whether to include hydrogen atoms. Default: False.
103        include_hetatm: Whether to include HETATM records. Default: False.
104
105    Returns:
106        BatchDirResult with per-file details.
107
108    Raises:
109        ValueError: If input parameters are invalid.
110        FileNotFoundError: If the input directory does not exist.
111        MemoryError: If out of memory.
112        RuntimeError: For other processing errors.
113
114    Example:
115        >>> from zsasa import process_directory
116        >>> result = process_directory("path/to/pdbs/")
117        >>> print(f"Processed {result.successful}/{result.total_files} files")
118    """
119    ffi, lib = _get_lib()
120
121    # Validate algorithm
122    if algorithm == "sr":
123        algo_int = ZSASA_ALGORITHM_SR
124        n_points_val = n_points
125    elif algorithm == "lr":
126        algo_int = ZSASA_ALGORITHM_LR
127        n_points_val = n_slices
128    else:
129        msg = f"Unknown algorithm: {algorithm}. Use 'sr' or 'lr'."
130        raise ValueError(msg)
131
132    if n_points <= 0:
133        msg = f"n_points must be positive, got {n_points}"
134        raise ValueError(msg)
135    if n_slices <= 0:
136        msg = f"n_slices must be positive, got {n_slices}"
137        raise ValueError(msg)
138    if probe_radius <= 0:
139        msg = f"probe_radius must be positive, got {probe_radius}"
140        raise ValueError(msg)
141    if n_threads < 0:
142        msg = f"n_threads must be non-negative, got {n_threads}"
143        raise ValueError(msg)
144
145    # Map classifier
146    classifier_int = int(classifier) if classifier is not None else -1
147
148    # Convert paths
149    input_dir_bytes = str(input_dir).encode("utf-8")
150    output_dir_bytes = str(output_dir).encode("utf-8") if output_dir is not None else ffi.NULL
151
152    error_code = ffi.new("int*")
153
154    handle = lib.zsasa_batch_dir_process(
155        input_dir_bytes,
156        output_dir_bytes,
157        algo_int,
158        n_points_val,
159        probe_radius,
160        n_threads,
161        classifier_int,
162        int(include_hydrogens),
163        int(include_hetatm),
164        error_code,
165    )
166
167    if handle == ffi.NULL:
168        ec = error_code[0]
169        if ec == ZSASA_ERROR_INVALID_INPUT:
170            msg = "Invalid input parameters for directory batch processing"
171            raise ValueError(msg)
172        elif ec == ZSASA_ERROR_OUT_OF_MEMORY:
173            msg = "Out of memory during directory batch processing"
174            raise MemoryError(msg)
175        elif ec == ZSASA_ERROR_CALCULATION:
176            msg = "SASA calculation failed during directory batch processing"
177            raise RuntimeError(msg)
178        elif ec == ZSASA_ERROR_FILE_IO:
179            msg = f"Directory not found or not readable: {input_dir}"
180            raise FileNotFoundError(msg)
181        else:
182            msg = f"Directory batch processing failed with error code: {ec}"
183            raise RuntimeError(msg)
184
185    try:
186        total_files = lib.zsasa_batch_dir_get_total_files(handle)
187        successful = lib.zsasa_batch_dir_get_successful(handle)
188        failed = lib.zsasa_batch_dir_get_failed(handle)
189
190        filenames: list[str] = []
191        n_atoms_list: list[int] = []
192        total_sasa_list: list[float] = []
193        status_list: list[int] = []
194
195        for i in range(total_files):
196            fname_ptr = lib.zsasa_batch_dir_get_filename(handle, i)
197            if fname_ptr == ffi.NULL:
198                msg = (
199                    f"Internal error: zsasa_batch_dir_get_filename returned NULL "
200                    f"for index {i} (total_files={total_files})"
201                )
202                raise RuntimeError(msg)
203            filenames.append(ffi.string(fname_ptr).decode("utf-8"))
204            n_atoms_list.append(lib.zsasa_batch_dir_get_n_atoms(handle, i))
205            sasa = lib.zsasa_batch_dir_get_total_sasa(handle, i)
206            total_sasa_list.append(float("nan") if math.isnan(sasa) else sasa)
207            st = lib.zsasa_batch_dir_get_status(handle, i)
208            if st not in (0, 1):
209                msg = f"Internal error: unexpected status {st} for file index {i} (expected 0 or 1)"
210                raise RuntimeError(msg)
211            status_list.append(st)
212
213        return BatchDirResult(
214            total_files=total_files,
215            successful=successful,
216            failed=failed,
217            filenames=filenames,
218            n_atoms=n_atoms_list,
219            total_sasa=total_sasa_list,
220            status=status_list,
221        )
222    finally:
223        lib.zsasa_batch_dir_free(handle)

23@dataclass
24class BatchDirResult:
25    """Result of directory batch processing.
26
27    Attributes:
28        total_files: Number of supported structure files found in the directory.
29        successful: Number of successfully processed files.
30        failed: Number of files that failed processing.
31        filenames: List of filenames (file name only, without directory path).
32        n_atoms: Per-file atom counts.
33        total_sasa: Per-file total SASA in Angstroms² (NaN for failed files).
34        status: Per-file status (1=ok, 0=failed).
35    """
36
37    total_files: int
38    successful: int
39    failed: int
40    filenames: list[str]
41    n_atoms: list[int]
42    total_sasa: list[float]
43    status: list[int]
44
45    def __post_init__(self) -> None:
46        n = len(self.filenames)
47        if len(self.n_atoms) != n or len(self.total_sasa) != n or len(self.status) != n:
48            msg = (
49                f"All per-file lists must have the same length as filenames ({n}), "
50                f"got n_atoms={len(self.n_atoms)}, total_sasa={len(self.total_sasa)}, "
51                f"status={len(self.status)}"
52            )
53            raise ValueError(msg)
54        if self.total_files != n:
55            msg = f"total_files ({self.total_files}) != len(filenames) ({n})"
56            raise ValueError(msg)
57        if self.successful + self.failed != self.total_files:
58            msg = (
59                f"successful ({self.successful}) + failed ({self.failed}) "
60                f"!= total_files ({self.total_files})"
61            )
62            raise ValueError(msg)
63
64    def __repr__(self) -> str:
65        return (
66            f"BatchDirResult(total_files={self.total_files}, "
67            f"successful={self.successful}, failed={self.failed})"
68        )

26class ClassifierType(IntEnum):
27    """Available classifier types for atom radius assignment.
28
29    Attributes:
30        CCD: CCD-based radii (default). Hardcoded ProtOr radii + runtime CCD
31             analysis for non-standard residues via bond topology.
32        PROTOR: Alias for CCD. Kept for backward compatibility.
33        NACCESS: NACCESS-compatible radii.
34        OONS: Ooi, Oobatake, Nemethy, Scheraga radii.
35    """
36
37    NACCESS = ZSASA_CLASSIFIER_NACCESS
38    PROTOR = ZSASA_CLASSIFIER_PROTOR
39    OONS = ZSASA_CLASSIFIER_OONS
40    CCD = ZSASA_CLASSIFIER_CCD

43class AtomClass(IntEnum):
44    """Atom polarity classes.
45
46    Attributes:
47        POLAR: Polar atoms (e.g., N, O).
48        APOLAR: Apolar/hydrophobic atoms (e.g., C).
49        UNKNOWN: Unknown classification.
50    """
51
52    POLAR = ZSASA_ATOM_CLASS_POLAR
53    APOLAR = ZSASA_ATOM_CLASS_APOLAR
54    UNKNOWN = ZSASA_ATOM_CLASS_UNKNOWN

173@dataclass
174class ClassificationResult:
175    """Result of batch atom classification.
176
177    Attributes:
178        radii: Per-atom van der Waals radii in Angstroms.
179               NaN values indicate atoms not found in the classifier.
180               Use np.isnan(result.radii) to find unknown atoms.
181        classes: Per-atom polarity classes (AtomClass constants).
182
183    Example:
184        >>> result = classify_atoms(residues, atoms)
185        >>> unknown_mask = np.isnan(result.radii)
186        >>> if unknown_mask.any():
187        ...     print(f"Found {unknown_mask.sum()} unknown atoms")
188    """
189
190    radii: NDArray[np.float64]
191    classes: NDArray[np.int32]
192
193    def __repr__(self) -> str:
194        return f"ClassificationResult(n_atoms={len(self.radii)})"

57def get_radius(
58    residue: str,
59    atom: str,
60    classifier_type: ClassifierType = ClassifierType.CCD,
61) -> float | None:
62    """Get van der Waals radius for an atom using the specified classifier.
63
64    Args:
65        residue: Residue name (e.g., "ALA", "GLY").
66        atom: Atom name (e.g., "CA", "CB").
67        classifier_type: Classifier to use. Default: ClassifierType.CCD.
68
69    Returns:
70        Radius in Angstroms, or None if atom is not found in classifier.
71
72    Example:
73        >>> from zsasa import get_radius, ClassifierType
74        >>> get_radius("ALA", "CA")
75        1.87
76        >>> get_radius("ALA", "XX")  # Unknown atom
77        None
78    """
79    _, lib = _get_lib()
80    radius = lib.zsasa_classifier_get_radius(
81        classifier_type,
82        residue.encode("utf-8"),
83        atom.encode("utf-8"),
84    )
85    if np.isnan(radius):
86        return None
87    return radius

 90def get_atom_class(
 91    residue: str,
 92    atom: str,
 93    classifier_type: ClassifierType = ClassifierType.CCD,
 94) -> int:
 95    """Get atom polarity class using the specified classifier.
 96
 97    Args:
 98        residue: Residue name (e.g., "ALA", "GLY").
 99        atom: Atom name (e.g., "CA", "CB").
100        classifier_type: Classifier to use. Default: ClassifierType.CCD.
101
102    Returns:
103        AtomClass constant (POLAR, APOLAR, or UNKNOWN).
104
105    Example:
106        >>> from zsasa import get_atom_class, AtomClass
107        >>> get_atom_class("ALA", "CA") == AtomClass.APOLAR
108        True
109        >>> get_atom_class("ALA", "O") == AtomClass.POLAR
110        True
111    """
112    _, lib = _get_lib()
113    return lib.zsasa_classifier_get_class(
114        classifier_type,
115        residue.encode("utf-8"),
116        atom.encode("utf-8"),
117    )

120def guess_radius(element: str) -> float | None:
121    """Guess van der Waals radius from element symbol.
122
123    Args:
124        element: Element symbol (e.g., "C", "N", "FE").
125                 Case-insensitive, whitespace is trimmed.
126
127    Returns:
128        Radius in Angstroms, or None if element is not recognized.
129
130    Example:
131        >>> from zsasa import guess_radius
132        >>> guess_radius("C")
133        1.7
134        >>> guess_radius("FE")
135        1.26
136        >>> guess_radius("XX")  # Unknown
137        None
138    """
139    _, lib = _get_lib()
140    radius = lib.zsasa_guess_radius(element.encode("utf-8"))
141    if np.isnan(radius):
142        return None
143    return radius

146def guess_radius_from_atom_name(atom_name: str) -> float | None:
147    """Guess van der Waals radius from PDB-style atom name.
148
149    Extracts element symbol from atom name following PDB conventions:
150    - Leading space indicates single-char element (e.g., " CA " = Carbon alpha)
151    - No leading space may indicate 2-char element (e.g., "FE  " = Iron)
152
153    Args:
154        atom_name: PDB-style atom name (e.g., " CA ", "FE  ").
155
156    Returns:
157        Radius in Angstroms, or None if element cannot be determined.
158
159    Example:
160        >>> from zsasa import guess_radius_from_atom_name
161        >>> guess_radius_from_atom_name(" CA ")  # Carbon alpha
162        1.7
163        >>> guess_radius_from_atom_name("FE  ")  # Iron
164        1.26
165    """
166    _, lib = _get_lib()
167    radius = lib.zsasa_guess_radius_from_atom_name(atom_name.encode("utf-8"))
168    if np.isnan(radius):
169        return None
170    return radius

197def classify_atoms(
198    residues: list[str],
199    atoms: list[str],
200    classifier_type: ClassifierType = ClassifierType.CCD,
201    *,
202    include_classes: bool = True,
203) -> ClassificationResult:
204    """Classify multiple atoms at once (batch operation).
205
206    This is more efficient than calling get_radius for each atom individually.
207
208    Args:
209        residues: List of residue names.
210        atoms: List of atom names (must be same length as residues).
211        classifier_type: Classifier to use. Default: ClassifierType.CCD.
212        include_classes: Whether to compute atom classes. Default: True.
213
214    Returns:
215        ClassificationResult with radii and classes arrays.
216        Unknown atoms have NaN radius and UNKNOWN class.
217
218    Raises:
219        ValueError: If residues and atoms have different lengths.
220
221    Example:
222        >>> from zsasa import classify_atoms
223        >>> result = classify_atoms(
224        ...     ["ALA", "ALA", "GLY"],
225        ...     ["CA", "O", "N"],
226        ... )
227        >>> result.radii
228        array([1.87, 1.4 , 1.65])
229    """
230    ffi, lib = _get_lib()
231
232    if len(residues) != len(atoms):
233        msg = f"residues and atoms must have same length: {len(residues)} != {len(atoms)}"
234        raise ValueError(msg)
235
236    n_atoms = len(residues)
237    if n_atoms == 0:
238        return ClassificationResult(
239            radii=np.array([], dtype=np.float64),
240            classes=np.array([], dtype=np.int32),
241        )
242
243    # Encode strings and create cffi arrays
244    # Keep references to prevent garbage collection
245    residues_bytes = [ffi.new("char[]", r.encode("utf-8")) for r in residues]
246    atoms_bytes = [ffi.new("char[]", a.encode("utf-8")) for a in atoms]
247
248    residues_arr = ffi.new("char*[]", residues_bytes)
249    atoms_arr = ffi.new("char*[]", atoms_bytes)
250
251    # Allocate output arrays
252    radii = np.zeros(n_atoms, dtype=np.float64)
253    classes = np.zeros(n_atoms, dtype=np.int32) if include_classes else None
254
255    radii_ptr = ffi.cast("double*", radii.ctypes.data)
256    classes_ptr = ffi.cast("int*", classes.ctypes.data) if include_classes else ffi.NULL
257
258    result = lib.zsasa_classify_atoms(
259        classifier_type,
260        residues_arr,
261        atoms_arr,
262        n_atoms,
263        radii_ptr,
264        classes_ptr,
265    )
266
267    if result == ZSASA_ERROR_INVALID_INPUT:
268        msg = f"Invalid classifier type: {classifier_type}"
269        raise ValueError(msg)
270    elif result != ZSASA_OK:
271        msg = f"Classification error: {result}"
272        raise RuntimeError(msg)
273
274    if not include_classes:
275        classes = np.full(n_atoms, AtomClass.UNKNOWN, dtype=np.int32)
276
277    return ClassificationResult(radii=radii, classes=classes)

36def get_max_sasa(residue_name: str) -> float | None:
37    """Get maximum SASA value for a standard amino acid.
38
39    Values from Tien et al. (2013) "Maximum allowed solvent accessibilities
40    of residues in proteins".
41
42    Args:
43        residue_name: 3-letter residue code (e.g., "ALA", "GLY").
44
45    Returns:
46        Maximum SASA in Angstroms², or None if residue is not a standard amino acid.
47
48    Example:
49        >>> from zsasa import get_max_sasa
50        >>> get_max_sasa("ALA")
51        129.0
52        >>> get_max_sasa("TRP")
53        285.0
54        >>> get_max_sasa("HOH")  # Water - not a standard amino acid
55        None
56    """
57    _, lib = _get_lib()
58    max_sasa = lib.zsasa_get_max_sasa(residue_name.encode("utf-8"))
59    if np.isnan(max_sasa):
60        return None
61    return max_sasa

64def calculate_rsa(sasa: float, residue_name: str) -> float | None:
65    """Calculate RSA (Relative Solvent Accessibility) for a single residue.
66
67    RSA = SASA / MaxSASA
68
69    Args:
70        sasa: Observed SASA value in Angstroms².
71        residue_name: 3-letter residue code (e.g., "ALA", "GLY").
72
73    Returns:
74        RSA value (typically 0.0-1.0), or None if residue is not a standard amino acid.
75        Note: RSA > 1.0 is possible for exposed terminal residues.
76
77    Example:
78        >>> from zsasa import calculate_rsa
79        >>> calculate_rsa(64.5, "ALA")  # 64.5 / 129.0 = 0.5
80        0.5
81        >>> calculate_rsa(150.0, "GLY")  # RSA > 1.0 is possible
82        1.4423076923076923
83    """
84    _, lib = _get_lib()
85    rsa = lib.zsasa_calculate_rsa(sasa, residue_name.encode("utf-8"))
86    if np.isnan(rsa):
87        return None
88    return rsa

 91def calculate_rsa_batch(
 92    sasas: NDArray[np.float64] | list[float],
 93    residue_names: list[str],
 94) -> NDArray[np.float64]:
 95    """Calculate RSA for multiple residues at once (batch operation).
 96
 97    This is more efficient than calling calculate_rsa for each residue individually.
 98
 99    Args:
100        sasas: Array of SASA values in Angstroms².
101        residue_names: List of 3-letter residue codes (must be same length as sasas).
102
103    Returns:
104        Array of RSA values. NaN values indicate non-standard amino acids.
105
106    Raises:
107        ValueError: If sasas and residue_names have different lengths.
108
109    Example:
110        >>> import numpy as np
111        >>> from zsasa import calculate_rsa_batch
112        >>> sasas = np.array([64.5, 52.0, 100.0])
113        >>> residues = ["ALA", "GLY", "HOH"]  # HOH is not standard
114        >>> rsa = calculate_rsa_batch(sasas, residues)
115        >>> rsa
116        array([0.5       , 0.5       ,        nan])
117    """
118    ffi, lib = _get_lib()
119
120    sasas = np.ascontiguousarray(sasas, dtype=np.float64)
121    n_residues = len(sasas)
122
123    if len(residue_names) != n_residues:
124        msg = f"sasas and residue_names must have same length: {n_residues} != {len(residue_names)}"
125        raise ValueError(msg)
126
127    if n_residues == 0:
128        return np.array([], dtype=np.float64)
129
130    # Encode strings and create cffi array
131    # Keep references to prevent garbage collection
132    residues_bytes = [ffi.new("char[]", r.encode("utf-8")) for r in residue_names]
133    residues_arr = ffi.new("char*[]", residues_bytes)
134
135    # Allocate output array
136    rsa_out = np.zeros(n_residues, dtype=np.float64)
137
138    sasas_ptr = ffi.cast("double*", sasas.ctypes.data)
139    rsa_ptr = ffi.cast("double*", rsa_out.ctypes.data)
140
141    result = lib.zsasa_calculate_rsa_batch(sasas_ptr, residues_arr, n_residues, rsa_ptr)
142    if result != ZSASA_OK:
143        msg = f"RSA batch calculation failed with error code: {result}"
144        raise RuntimeError(msg)
145
146    return rsa_out