from collections import Counter
from typing import cast

import numpy as np

from zyplib.utils.ensure import ensure_npy
from zyplib.utils.print import print_info


def summary_labels(y: np.ndarray, just_print: bool = False) -> Counter:
    """统计标签的分布

    Parameters
    ----------
    - `y` : `np.ndarray`
        - 标签
    - `just_print` : `bool`, optional
        - 是否只打印结果

    Returns
    -------
    - `Counter`
        - 标签的分布


    Examples
    --------
    ```python
    summary_labels(np.array([0, 1, 0, 1, 2]))
    # Counter({0: 2, 1: 2, 2: 1})
    ```
    """
    y = cast(np.ndarray, ensure_npy(y))
    y = y.squeeze()
    if y.ndim != 1:
        raise ValueError('y.ndim must be 1')

    counter = Counter(y)

    if not just_print:
        return counter

    print_info(counter)
    return counter
