import numpy as np

from kinematic_tracker.types import FMT, IVT


TYPE_ERR_MSG = 'Time stamp (nanoseconds) should be an integer or np.int64, but got '
SIZE_ERR_MSG = 'Wrong size of the measurement vector '


def check_det_rz(det_rz: FMT, num_z: int) -> None:
    if not isinstance(det_rz, np.ndarray):
        raise TypeError(f'Expected (two-dimensional) NumPy array, got {type(det_rz)}')
    if det_rz.ndim != 2:
        raise ValueError(f'Expected two-dimensional NumPy array, got {det_rz.ndim}')
    if det_rz.shape[1] != num_z:
        raise ValueError(f'{SIZE_ERR_MSG}{det_rz.shape[1]}. Expected {num_z}')


def check_stamp_type(stamp: int | np.int64) -> None:
    if not isinstance(stamp, int) and not isinstance(stamp, np.int64):
        raise TypeError(f'{TYPE_ERR_MSG}{type(stamp)}')


def check_cov_zz(cov_zz: FMT) -> None:
    """Ensure the measurement covariance is fully defined.

    Raises:
        ValueError: if the measurement covariance contains NaN values.
    """
    if np.isnan(cov_zz.sum()):
        raise ValueError('Not all elements of measurement covariance are defined.')


def check_cls_r(cls_r: IVT, num_reports: int, num_cls: int) -> None:
    """Ensure classes matches the reports and are within bounds."""
    if cls_r.shape[0] != num_reports:
        raise ValueError(f'Number of reports should be {num_reports}, but got {cls_r.shape[0]}.')
    if num_reports < 1:
        return
    if cls_r.min() < 0:
        raise ValueError(f'Class IDs should be greater or equal zero, but I see {cls_r.min()}.')
    if cls_r.max() >= num_cls:
        raise ValueError(f'Class IDs should be less than {num_cls}, but I see {cls_r.max()}.')


def check_ids_r(ids_r: IVT, num_reports: int) -> None:
    """Ensure the annotation ids are of correct type and size.

    Raises:
        ValueError: if the dimensionality of the ids vector is incorrect.
        TypeError: if the vector of IDs is not integer-valued.
    """
    if not isinstance(ids_r, np.ndarray):
        raise TypeError(f'I expect NumPy array, but I see {type(ids_r)}.')

    if ids_r.ndim == 1 and ids_r.shape[0] == num_reports and np.issubdtype(ids_r.dtype, np.integer):
        return
    if ids_r.ndim != 1:
        raise ValueError(f'I expect one-dimensional vector, but I see {ids_r.ndim}.')
    elif ids_r.shape[0] != num_reports:
        raise ValueError(f'I expect {num_reports} ids, but I see {ids_r.shape[0]}.')
    elif not np.issubdtype(ids_r.dtype, np.integer):
        raise TypeError(f'I expect integer ids, but I see {ids_r.dtype}.')
