from typing import Literal
from typing import Optional
from typing import Tuple

import numpy


def compute_axis_transform(
    coordinates: numpy.ndarray,
    method: Optional[Literal["range", "std"]] = "range",
) -> Tuple[Optional[numpy.ndarray], Optional[numpy.ndarray]]:
    """
    Compute per-axis scaling factors and offsets.

    Any zero scale values are replaced with 1.0 to avoid division by zero.

    :param coordinates: Array of shape (Npoints, Ndim)
    :param method:
        - "range": scale = max - min, offset = min
        - "std": scale = std, offset = mean
        - None: scale = 1, offset = 0
    :return: (scale, offset), each shape (Ndim,)
    """
    dtype = coordinates.dtype

    if method is None:
        return None, None

    if method == "range":
        mins = numpy.min(coordinates, axis=0)
        maxs = numpy.max(coordinates, axis=0)
        scale = maxs - mins
        offset = mins

    elif method == "std":
        scale = numpy.std(coordinates, axis=0)
        offset = numpy.mean(coordinates, axis=0)

    else:
        raise ValueError(f"Unknown method: {method}")

    scale = scale.astype(dtype, copy=False)
    offset = offset.astype(dtype, copy=False)

    scale[scale == 0] = 1.0

    return scale, offset


def normalize_coordinates(
    coordinates: numpy.ndarray,
    scale: Optional[numpy.ndarray],
    offset: Optional[numpy.ndarray],
) -> numpy.ndarray:
    if scale is None and offset is None:
        return coordinates
    if scale is None:
        scale = 1
    if offset is None:
        offset = 0
    return (coordinates - offset) / scale
