import numpy as np
from plaster.tools.schema import check
from plaster.tools.c_common.c_common_tools import DytPepType

"""
I considered a subclass of ndarray but decided against it
owing to the complications of numpy array overloading.
"""

DYT_COL = 0
PEP_COL = 1
CNT_COL = 2


def validate(dytpeps: np.ndarray):
    check.array_t(dytpeps, shape=(None, 3), dtype=DytPepType, c_contiguous=True)


def n_dyts(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.max(dytpeps[:, DYT_COL]) + 1)


def n_peps(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.max(dytpeps[:, PEP_COL]) + 1)


def n_samples(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.sum(dytpeps[:, CNT_COL]))


def argsort_dytpeps(dytpeps: np.ndarray):
    # SORT dytpeps by peptide (col [1]) first then by count (col [2])
    # Note that np.lexsort puts the primary sort key LAST in the argument
    validate(dytpeps)
    return np.lexsort((-dytpeps[:, CNT_COL], dytpeps[:, PEP_COL]))


def sort_dytpeps(dytpeps: np.ndarray):
    return dytpeps[argsort_dytpeps(dytpeps)]

import numpy as np
from plaster.tools.schema import check
from plaster.tools.c_common.c_common_tools import DytPepType

"""
I considered a subclass of ndarray but decided against it
owing to the complications of numpy array overloading.
"""

DYT_COL = 0
PEP_COL = 1
CNT_COL = 2


def validate(dytpeps: np.ndarray):
    check.array_t(dytpeps, shape=(None, 3), dtype=DytPepType, c_contiguous=True)


def n_dyts(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.max(dytpeps[:, DYT_COL]) + 1)


def n_peps(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.max(dytpeps[:, PEP_COL]) + 1)


def n_samples(dytpeps: np.ndarray):
    validate(dytpeps)
    return int(np.sum(dytpeps[:, CNT_COL]))


def argsort_dytpeps(dytpeps: np.ndarray):
    # SORT dytpeps by peptide (col [1]) first then by count (col [2])
    # Note that np.lexsort puts the primary sort key LAST in the argument
    validate(dytpeps)
    return np.lexsort((-dytpeps[:, CNT_COL], dytpeps[:, PEP_COL]))


def sort_dytpeps(dytpeps: np.ndarray):
    return dytpeps[argsort_dytpeps(dytpeps)]

I considered a subclass of ndarray but decided against it owing to the complications of numpy array overloading. # SORT dytpeps by peptide (col [1]) first then by count (col [2]) # Note that np.lexsort puts the primary sort key LAST in the argument