from time import perf_counter

import numba as nb
import numpy as np
from numba.extending import overload


def _prepare_fixed_effects(ary):
    pass


@overload(_prepare_fixed_effects)
def _ol_preproc_fixed_effects(ary):
    # If array is already an F-array we tolerate
    # any dtype because it saves us a copy
    if ary.layout == "F":
        return lambda ary: ary

    if not isinstance(ary.dtype, nb.types.Integer):
        raise nb.TypingError("Fixed effects must be integers")

    max_nbits = 32
    nbits = min(max_nbits, ary.dtype.bitwidth)
    dtype = nb.types.Integer.from_bitwidth(nbits, signed=False)

    def impl(ary):
        n, m = ary.shape
        out = np.empty((m, n), dtype=dtype).T
        out[:] = ary[:]
        return out

    return impl


@nb.njit
def new_detect_singletons(ids):
    """
    Detect singleton fixed effects in a dataset.

    Args:
        ids (np.ndarray): A 2D numpy array representing fixed effects, with a
                          shape of (n_samples, n_features).
                          Elements should be non-negative integers representing
                          fixed effect identifiers.

    Returns:
        np.ndarray: A boolean array of shape (n_samples,), indicating which
                    observations have a singleton fixed effect.

    Note:
        The algorithm iterates over columns to identify fixed effects.
        After completing each column traversal, it updates the record of
        non-singleton rows. This approach is preferred as removing an
        observation in column 'i' may lead to the emergence of new singletons
        in subsequent columns (i.e., columns > i).

        Since we are operating on columns, we enforce column-major order for
        the input array. Working on a row-major array leads to considerable
        performance losses.
    """
    ids = _prepare_fixed_effects(ids)
    n_samples, n_features = ids.shape

    max_fixef = np.max(ids)
    counts = np.empty(max_fixef + 1, dtype=np.uint32)

    n_non_singletons = n_samples
    non_singletons = np.arange(n_non_singletons, dtype=np.uint32)

    while True:
        n_non_singletons_curr = n_non_singletons

        for j in range(n_features):
            col = ids[:, j]

            counts[:] = 0
            n_singletons = 0
            for i in range(n_non_singletons):
                e = col[non_singletons[i]]
                c = counts[e]
                # Branchless version of:
                #
                # if counts[e] == 1:
                #     n_singletons -= 1
                # elif counts[e] == 0:
                #     n_singletons += 1
                #
                n_singletons += (c == 0) - (c == 1)
                counts[e] += 1

            if not n_singletons:
                continue

            cnt = 0
            for i in range(n_non_singletons):
                e = col[non_singletons[i]]
                if counts[e] != 1:
                    non_singletons[cnt] = non_singletons[i]
                    cnt += 1

            n_non_singletons = cnt

        if n_non_singletons_curr == n_non_singletons:
            break

    is_singleton = np.ones(n_samples, dtype=np.bool_)
    for i in range(n_non_singletons):
        is_singleton[non_singletons[i]] = False

    return is_singleton


def old_detect_singletons(ids):
    """
    Detect singleton fixed effects
    Args:have a
        ids (np.ndarray): A numpy array of fixed effects.
    Returns:
        An array of booleans indicating which observations have a singleton fixed effect.
    """

    N, k = ids.shape

    singleton_idx = np.full(N, False, dtype=bool)
    singleton_idx_tmp = np.full(N, False, dtype=bool)
    singleton_idx_tmp_old = singleton_idx.copy()

    ids_tmp = ids.copy()

    while True:
        for x in range(k):
            col = ids[:, x]
            col_tmp = ids_tmp[:, x]

            # note that this only "works" as fixed effects are integers from 0, 1, ..., n_fixef -1
            # and np.bincount orders results in ascending (integer) order
            counts = np.bincount(col_tmp)

            if np.any(counts == 1):
                idx = np.where(counts == 1)

                singleton_idx[np.isin(col, idx)] = True
                singleton_idx_tmp[np.isin(col_tmp, idx)] = True

                ids_tmp = ids_tmp[~singleton_idx_tmp, :].astype(ids_tmp.dtype)
                singleton_idx_tmp = singleton_idx_tmp[~singleton_idx_tmp]

            # break

        if np.array_equal(singleton_idx_tmp, singleton_idx_tmp_old):
            break

        singleton_idx_tmp_old = singleton_idx_tmp.copy()

    return singleton_idx


@nb.njit
def make_deterministic(shape, n_singletons, dtype=np.int64):
    out = np.zeros(shape, dtype=dtype)
    assert n_singletons <= shape[0]
    for i in range(n_singletons):
        j = i % shape[1]
        out[i, j] = i + 1
    return out


class SimpleTimer:
    def __init__(self):
        self.tic = None
        self.toc = None

    def __enter__(self):
        self.tic = perf_counter()
        return self

    def __exit__(self, *args):
        self.toc = perf_counter()

    def __str__(self):
        return f"Elapsed [s]: {self.elapsed:.2f}"

    def __truediv__(self, other):
        return self.elapsed / other.elapsed

    @property
    def elapsed(self):
        return self.toc - self.tic


if __name__ == "__main__":
    # warmup
    new_detect_singletons(np.zeros((1, 1), dtype=np.int64))

    N, M = 100_000, 100

    for i, a in enumerate(
        (
            make_deterministic((N, M), n_singletons=0),
            make_deterministic((N, M), n_singletons=N // 100_000),
            make_deterministic((N, M), n_singletons=N // 10_000),
            make_deterministic((N, M), n_singletons=N // 1_000),
            make_deterministic((N, M), n_singletons=N // 100),
            make_deterministic((N, M), n_singletons=N // 10),
            make_deterministic((N, M), n_singletons=N),
        )
    ):
        print(f"\nRunning case {i+1}")

        with SimpleTimer() as t1:
            r1 = old_detect_singletons(a)

        with SimpleTimer() as t2:
            r2 = new_detect_singletons(a)

        assert np.array_equal(r1, r2)

        print(t1)
        print(t2)
        print(f"Speedup: {t1 / t2:.1f} (#singletons {r1.sum()}/{r1.shape[0]})")