"""
Tafra: a minimalist dataframe

Copyright (c) 2020 Derrick W. Turk and David S. Fulford

Author
------
Derrick W. Turk
David S. Fulford

Notes
-----
Created on April 25, 2020
"""
__all__ = ['GroupBy', 'Transform', 'IterateBy', 'InnerJoin', 'LeftJoin',
           'percentile', 'geomean', 'harmean']

import operator
from itertools import chain
import dataclasses as dc

import numpy as np

from typing import (Any, Callable, Dict, Mapping, List, Tuple, Optional, Union as _Union, Sequence,
                    Iterable, Iterator)
from typing import cast

try:
    from ._accel import (  # type: ignore[import-not-found]
        groupby_sum as _c_sum,
        groupby_mean as _c_mean,
        groupby_var as _c_var,
        groupby_min as _c_min,
        groupby_max as _c_max,
        groupby_count as _c_count,
        inner_join as _c_inner_join,
        left_join as _c_left_join,
    )
    _HAS_ACCEL = True
except ImportError:
    _HAS_ACCEL = False


# Vectorized aggregation functions that can bypass per-group Python loops.
# Maps function identity to a callable(data, labels, n_groups) -> result_array.
_VECTORIZED_AGGS: Dict[int, Callable[..., np.ndarray[Any, Any]]] = {}


def _sorted_segments(
    data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any], n: int
) -> Tuple[np.ndarray[Any, Any], np.ndarray[Any, Any], np.ndarray[Any, Any]]:
    """Sort data by group labels, return (sorted_data, offsets, counts)."""
    order = np.argsort(labels, kind='stable')
    sorted_data = data[order]
    counts = np.bincount(labels, minlength=n)
    offsets = np.zeros(n + 1, dtype=np.intp)
    np.cumsum(counts, out=offsets[1:])
    return sorted_data, offsets, counts


def _vec_sum(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    return np.bincount(labels, weights=data.astype(float), minlength=n)


def _vec_mean(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
              n: int) -> np.ndarray[Any, Any]:
    sums = np.bincount(labels, weights=data.astype(float), minlength=n)
    counts = np.bincount(labels, minlength=n)
    return sums / counts


def _vec_var(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    fdata = data.astype(float)
    counts = np.bincount(labels, minlength=n).astype(float)
    sums = np.bincount(labels, weights=fdata, minlength=n)
    sum_sq = np.bincount(labels, weights=fdata * fdata, minlength=n)
    mean = sums / counts
    return sum_sq / counts - mean * mean


def _vec_std(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    return np.sqrt(_vec_var(data, labels, n))


def _vec_min(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data, labels, n)
    return np.minimum.reduceat(sorted_data, offsets[:n])


def _vec_max(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data, labels, n)
    return np.maximum.reduceat(sorted_data, offsets[:n])


def _vec_ptp(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data, labels, n)
    return (np.maximum.reduceat(sorted_data, offsets[:n])
            - np.minimum.reduceat(sorted_data, offsets[:n]))


def _vec_prod(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
              n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data, labels, n)
    return np.multiply.reduceat(sorted_data.astype(float), offsets[:n])


def _vec_any(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data.astype(bool), labels, n)
    return np.logical_or.reduceat(sorted_data, offsets[:n])


def _vec_all(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, _ = _sorted_segments(data.astype(bool), labels, n)
    return np.logical_and.reduceat(sorted_data, offsets[:n])


def _vec_len(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
             n: int) -> np.ndarray[Any, Any]:
    return np.bincount(labels, minlength=n)


def _vec_count_nonzero(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                       n: int) -> np.ndarray[Any, Any]:
    return np.bincount(labels, weights=(data != 0).astype(float), minlength=n)


def _vec_median(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                n: int) -> np.ndarray[Any, Any]:
    sorted_data, offsets, counts = _sorted_segments(data, labels, n)
    result = np.empty(n, dtype=float)
    for g in range(n):
        lo = offsets[g]
        hi = offsets[g + 1]
        seg = np.sort(sorted_data[lo:hi])
        c = int(counts[g])
        mid = c // 2
        if c % 2 == 1:
            result[g] = seg[mid]
        else:
            result[g] = (seg[mid - 1] + seg[mid]) / 2.0
    return result


def _vec_geomean(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                 n: int) -> np.ndarray[Any, Any]:
    log_sums = np.bincount(labels, weights=np.log(data.astype(float)), minlength=n)
    counts = np.bincount(labels, minlength=n).astype(float)
    return np.exp(log_sums / counts)


def _vec_harmean(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                 n: int) -> np.ndarray[Any, Any]:
    recip_sums = np.bincount(labels, weights=1.0 / data.astype(float), minlength=n)
    counts = np.bincount(labels, minlength=n).astype(float)
    return counts / recip_sums


def _make_vec_percentile(
    q: float
) -> Callable[[np.ndarray[Any, Any], np.ndarray[Any, Any], int], np.ndarray[Any, Any]]:
    """Create a vectorized percentile function for a given quantile."""
    def _vec_pct(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                 n: int) -> np.ndarray[Any, Any]:
        sorted_data, offsets, counts = _sorted_segments(data, labels, n)
        result = np.empty(n, dtype=float)
        for g in range(n):
            lo = offsets[g]
            hi = offsets[g + 1]
            result[g] = np.percentile(sorted_data[lo:hi], q)
        return result
    return _vec_pct


class _PercentileAgg:
    """Callable aggregation for percentile(q). Registered for vectorized fast path."""
    __slots__ = ('q', '_vec_fn')

    def __init__(self, q: float) -> None:
        self.q = q
        self._vec_fn = _make_vec_percentile(q)
        _VECTORIZED_AGGS[id(self)] = self._vec_fn

    def __call__(self, data: np.ndarray[Any, Any]) -> Any:
        return np.percentile(data, self.q)

    def __repr__(self) -> str:
        return f'percentile({self.q})'


def percentile(q: float) -> _PercentileAgg:
    """
    Create a percentile aggregation function for use in ``group_by``.

    Parameters
    ----------
        q: float
            Percentile in range [0, 100].

    Returns
    -------
        agg: callable
            A callable suitable for ``group_by`` aggregation that also
            hits the vectorized fast path.

    Example
    -------
        >>> tf.group_by(['g'], {'p90': (percentile(90), 'value')})
    """
    return _PercentileAgg(q)


def geomean(data: np.ndarray[Any, Any]) -> Any:
    """Geometric mean aggregation for use in ``group_by``."""
    return np.exp(np.mean(np.log(data.astype(float))))


def harmean(data: np.ndarray[Any, Any]) -> Any:
    """Harmonic mean aggregation for use in ``group_by``."""
    return len(data) / np.sum(1.0 / data.astype(float))


def _register_vectorized() -> None:
    """Register known numpy reducers for vectorized GroupBy."""
    # Use C accelerated versions when available, fall back to numpy
    if _HAS_ACCEL:
        def _wrap_c(c_fn: Callable[..., Any]) -> Callable[..., np.ndarray[Any, Any]]:
            def _wrapped(data: np.ndarray[Any, Any], labels: np.ndarray[Any, Any],
                         n: int) -> np.ndarray[Any, Any]:
                return c_fn(labels.astype(np.int64, copy=False),
                            data.astype(np.float64, copy=False), n)
            return _wrapped
        _fast_sum = _wrap_c(_c_sum)
        _fast_mean = _wrap_c(_c_mean)
        _fast_var = _wrap_c(_c_var)
        _fast_min = _wrap_c(_c_min)
        _fast_max = _wrap_c(_c_max)
    else:
        _fast_sum = _vec_sum
        _fast_mean = _vec_mean
        _fast_var = _vec_var
        _fast_min = _vec_min
        _fast_max = _vec_max

    _fast_std = (lambda d, l, n: np.sqrt(_fast_var(d, l, n))) if _HAS_ACCEL else _vec_std

    _VECTORIZED_AGGS[id(np.sum)] = _fast_sum
    _VECTORIZED_AGGS[id(np.nansum)] = _fast_sum
    _VECTORIZED_AGGS[id(np.mean)] = _fast_mean
    _VECTORIZED_AGGS[id(np.nanmean)] = _fast_mean
    _VECTORIZED_AGGS[id(np.std)] = _fast_std
    _VECTORIZED_AGGS[id(np.nanstd)] = _fast_std
    _VECTORIZED_AGGS[id(np.var)] = _fast_var
    _VECTORIZED_AGGS[id(np.nanvar)] = _fast_var
    _VECTORIZED_AGGS[id(np.min)] = _fast_min
    _VECTORIZED_AGGS[id(np.nanmin)] = _fast_min
    _VECTORIZED_AGGS[id(np.amin)] = _fast_min
    _VECTORIZED_AGGS[id(np.max)] = _fast_max
    _VECTORIZED_AGGS[id(np.nanmax)] = _fast_max
    _VECTORIZED_AGGS[id(np.amax)] = _fast_max
    _VECTORIZED_AGGS[id(np.ptp)] = _vec_ptp
    _VECTORIZED_AGGS[id(np.prod)] = _vec_prod
    _VECTORIZED_AGGS[id(np.nanprod)] = _vec_prod
    _VECTORIZED_AGGS[id(np.any)] = _vec_any
    _VECTORIZED_AGGS[id(np.all)] = _vec_all
    _VECTORIZED_AGGS[id(np.median)] = _vec_median
    _VECTORIZED_AGGS[id(np.nanmedian)] = _vec_median
    _VECTORIZED_AGGS[id(np.count_nonzero)] = _vec_count_nonzero
    _VECTORIZED_AGGS[id(len)] = _vec_len
    _VECTORIZED_AGGS[id(sum)] = _vec_sum
    _VECTORIZED_AGGS[id(geomean)] = _vec_geomean
    _VECTORIZED_AGGS[id(harmean)] = _vec_harmean


_register_vectorized()


JOIN_OPS: Dict[str, Callable[[Any, Any], Any]] = {
    '==': operator.eq,
    '!=': operator.ne,
    '<': operator.lt,
    '<=': operator.le,
    '>': operator.gt,
    '>=': operator.ge
}

# for the passed argument to an aggregation
InitAggregation = Mapping[
    str,
    _Union[
        Callable[[np.ndarray[Any, Any]], Any],
        Tuple[Callable[[np.ndarray[Any, Any]], Any], str]
    ]
]


# for the result type of IterateBy
GroupDescription = Tuple[
    Tuple[Any, ...],  # tuple of unique values from group-by columns
    np.ndarray[Any, Any],  # int array of row indices into original tafra for this group
    'Tafra'  # sub-tafra for the group
]


class Union:
    """
    Union two :class:`Tafra` together. Analogy to SQL UNION or
    `pandas.append`. All column names and dtypes must match.
    """
    @staticmethod
    def _validate(left: 'Tafra', right: 'Tafra') -> None:
        """
        Validate the :class:`Tafra` before applying.
        """
        # These should be unreachable unless attributes were directly modified
        if len(left._data) != len(left._dtypes):
            raise ValueError('This `Tafra` length of data and dtypes do not match')
        if len(right._data) != len(right._dtypes):
            raise ValueError('right `Tafra` length of data and dtypes do not match')

        # ensure same number of columns
        if len(left._data) != len(right._data) or len(left._dtypes) != len(right._dtypes):
            raise ValueError(
                'This `Tafra` column count does not match right `Tafra` column count.')

        # ensure all columns in this `Tafra` exist in right `Tafra`
        # if len() is same AND all columns in this exist in right,
        # do not need to check right `Tafra` columns in this `Tafra`.
        for (data_column, value), (dtype_column, dtype) \
                in zip(left._data.items(), left._dtypes.items()):

            if data_column not in right._data or dtype_column not in right._dtypes:
                raise TypeError(
                    f'This `Tafra` column `{data_column}` does not exist in right `Tafra`.')

            elif value.dtype != right._data[data_column].dtype:
                raise TypeError(
                    f'This `Tafra` column `{data_column}` dtype `{value.dtype}` '
                    f'does not match right `Tafra` dtype `{right._data[data_column].dtype}`.')

            # should not happen unless dtypes manually changed, but let's check it
            elif dtype != right._dtypes[dtype_column]:
                raise TypeError(
                    f'This `Tafra` column `{data_column}` dtype `{dtype}` '
                    f'does not match right `Tafra` dtype `{right._dtypes[dtype_column]}`.')

    def apply(self, left: 'Tafra', right: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`Union_` to the :class:`Tafra`.

        Parameters
        ----------
            left: Tafra
                The left :class:`Tafra` to union.

            right: Tafra
                The right :class:`Tafra` to union.

        Returns
        -------
            tafra: Tafra
                The unioned :class`Tafra`.
        """
        self._validate(left, right)

        return Tafra(
            {column: np.append(value, right._data[column]) for column, value in left._data.items()},
            left._dtypes.copy()
        )

    def apply_inplace(self, left: 'Tafra', right: 'Tafra') -> None:
        """
        In-place version.

        Apply the :class:`Union_` to the :class:`Tafra`.

        Parameters
        ----------
            left: Tafra
                The left :class:`Tafra` to union.

            right: Tafra
                The right :class:`Tafra` to union.

        Returns
        -------
            tafra: Tafra
                The unioned :class`Tafra`.
        """
        self._validate(left, right)

        for column, value in left._data.items():
            left._data[column] = np.append(value, right._data[column])
        left._update_rows()

@dc.dataclass
class GroupSet:
    """
    A `GroupSet` is the set of columns by which we construct our groups.
    """

    @staticmethod
    def _encode_columns(
        col_arrays: List[np.ndarray[Any, Any]]
    ) -> Tuple[List[np.ndarray[Any, Any]], List[Optional[np.ndarray[Any, Any]]]]:
        """
        Encode columns for structured array compatibility.
        StringDType columns are converted to integer codes via np.unique.

        Returns (encoded_arrays, codebooks) where codebooks[i] is the
        unique values array for encoded columns, or None if no encoding
        was needed.
        """
        encoded = []
        codebooks: List[Optional[np.ndarray[Any, Any]]] = []
        for c in col_arrays:
            if c.dtype.kind in ('T', 'U', 'S', 'O'):
                uniq, codes = np.unique(c, return_inverse=True)
                encoded.append(codes)
                codebooks.append(uniq)
            else:
                encoded.append(c)
                codebooks.append(None)
        return encoded, codebooks

    @staticmethod
    def _encode_columns_paired(
        left_cols: List[np.ndarray[Any, Any]],
        right_cols: List[np.ndarray[Any, Any]],
    ) -> Tuple[List[np.ndarray[Any, Any]], List[np.ndarray[Any, Any]]]:
        """
        Encode left and right column pairs with a shared codebook.
        Ensures the same string value gets the same integer code on both sides.
        """
        left_enc = []
        right_enc = []
        for lc, rc in zip(left_cols, right_cols):
            if lc.dtype.kind in ('T', 'U', 'S', 'O'):
                combined = np.concatenate([lc, rc])
                _, codes = np.unique(combined, return_inverse=True)
                left_enc.append(codes[:len(lc)])
                right_enc.append(codes[len(lc):])
            else:
                left_enc.append(lc)
                right_enc.append(rc)
        return left_enc, right_enc

    @staticmethod
    def _build_composite_key(
        encoded: List[np.ndarray[Any, Any]]
    ) -> np.ndarray[Any, Any]:
        """
        Combine multiple integer-coded columns into a single flat key.
        Uses positional encoding: key = c0 * N1*N2*... + c1 * N2*... + c2 * ...
        Falls back to structured array if values would overflow int64.
        """
        if len(encoded) == 1:
            return encoded[0]

        # compute cardinality of each column
        cards = [int(c.max()) + 1 for c in encoded]

        # check for overflow: product of all cardinalities must fit int64
        product = 1
        overflow = False
        for card in cards:
            product *= card
            if product > 2**62:
                overflow = True
                break

        if not overflow:
            # flat integer key via positional encoding
            key = np.zeros(len(encoded[0]), dtype=np.int64)
            multiplier = 1
            for c, card in zip(reversed(encoded), reversed(cards)):
                key += c.astype(np.int64) * multiplier
                multiplier *= card
            return key
        else:
            # fallback: structured array (slow but handles any cardinality)
            n_rows = len(encoded[0])
            dt = np.dtype([(f'f{i}', c.dtype) for i, c in enumerate(encoded)])
            data = np.empty(n_rows, dtype=dt)
            for i, c in enumerate(encoded):
                data[f'f{i}'] = c
            return data

    @staticmethod
    def _decode_unique(
        first_seen_indices: np.ndarray[Any, Any],
        col_arrays: List[np.ndarray[Any, Any]],
    ) -> List[Tuple[Any, ...]]:
        """Build unique group tuples by indexing original columns at first-seen rows."""
        if len(col_arrays) == 1:
            return [(v,) for v in col_arrays[0][first_seen_indices]]
        else:
            return [
                tuple(c[idx] for c in col_arrays)
                for idx in first_seen_indices
            ]

    @staticmethod
    def _direct_labels_firstseen(
        data: np.ndarray[Any, Any], n_rows: int
    ) -> Tuple[np.ndarray[Any, Any], np.ndarray[Any, Any], int]:
        """
        Assign first-seen-order group labels using direct array mapping.
        O(n + max_key) — no sort. Requires non-negative integer keys.

        Returns (labels, first_seen_indices, n_groups).
        """
        max_key = int(data.max())
        # sorted label map
        seen = np.zeros(max_key + 1, dtype=bool)
        seen[data] = True
        n_groups = int(seen.sum())
        sorted_map = np.empty(max_key + 1, dtype=np.intp)
        sorted_map[seen] = np.arange(n_groups)
        sorted_labels = sorted_map[data]
        # reorder to first-seen
        first_pos = np.empty(n_groups, dtype=np.intp)
        first_pos[sorted_labels[::-1]] = np.arange(n_rows - 1, -1, -1)
        order = np.argsort(first_pos)
        rank = np.empty(n_groups, dtype=np.intp)
        rank[order] = np.arange(n_groups, dtype=np.intp)
        labels = rank[sorted_labels]
        return labels, first_pos[order], n_groups

    @staticmethod
    def _direct_labels_sorted(
        data: np.ndarray[Any, Any]
    ) -> Tuple[np.ndarray[Any, Any], int]:
        """
        Assign sorted-order group labels using direct array mapping.
        O(n + max_key) — no sort. For Transform (order doesn't matter).

        Returns (labels, n_groups).
        """
        max_key = int(data.max())
        seen = np.zeros(max_key + 1, dtype=bool)
        seen[data] = True
        n_groups = int(seen.sum())
        label_map = np.empty(max_key + 1, dtype=np.intp)
        label_map[seen] = np.arange(n_groups, dtype=np.intp)
        return label_map[data], n_groups

    @staticmethod
    def _prepare_keys(
        tafra: 'Tafra', columns: Iterable[str]
    ) -> Tuple[np.ndarray[Any, Any], List[np.ndarray[Any, Any]]]:
        """Encode columns and build composite integer key."""
        cols = list(columns)
        col_arrays = [tafra._data[c] for c in cols]
        encoded, _ = GroupSet._encode_columns(col_arrays)
        if len(cols) == 1:
            data = encoded[0]
        else:
            data = GroupSet._build_composite_key(encoded)
        return data, col_arrays

    @staticmethod
    def _build_group_indices(
        tafra: 'Tafra', columns: Iterable[str]
    ) -> Tuple[List[Tuple[Any, ...]], List[np.ndarray[Any, Any]]]:
        """
        Build per-group row index arrays in a single pass.

        Returns (unique, group_indices) where:
        - unique: list of tuples of group-by values, in first-seen order
        - group_indices: list of int arrays (row positions for each group)
        """
        data, col_arrays = GroupSet._prepare_keys(tafra, columns)

        labels, first_seen_idx, n_groups = GroupSet._direct_labels_firstseen(
            data, tafra._rows)

        sorted_row_indices = np.argsort(labels, kind='stable')
        counts = np.bincount(labels, minlength=n_groups)
        splits = np.cumsum(counts[:-1])
        group_indices: List[np.ndarray[Any, Any]] = list(
            np.split(sorted_row_indices, splits))

        unique: List[Tuple[Any, ...]] = GroupSet._decode_unique(
            first_seen_idx, col_arrays)

        return unique, group_indices

    @staticmethod
    def _build_group_labels(
        tafra: 'Tafra', columns: Iterable[str]
    ) -> Tuple[List[Tuple[Any, ...]], np.ndarray[Any, Any], int]:
        """
        Build per-row group labels (integers 0..n_groups-1) in first-seen order.

        Returns (unique, labels, n_groups) where labels[i] is the group index
        for row i.
        """
        data, col_arrays = GroupSet._prepare_keys(tafra, columns)

        labels, first_seen_idx, n_groups = GroupSet._direct_labels_firstseen(
            data, tafra._rows)

        unique: List[Tuple[Any, ...]] = GroupSet._decode_unique(
            first_seen_idx, col_arrays)

        return unique, labels, n_groups

    @staticmethod
    def _validate(tafra: 'Tafra', columns: Iterable[str]) -> None:  # pragma: no cover
        """
        Validate the :class:`Tafra` before applying.
        """
        if tafra._rows < 1:
            raise ValueError('No rows exist in `tafra`.')
        tafra._validate_columns(columns)


@dc.dataclass
class AggMethod(GroupSet):
    """
    Basic methods for aggregations over a data table.
    """
    group_by_cols: Iterable[str]
    aggregation: dc.InitVar[InitAggregation]
    _aggregation: Mapping[
        str, Tuple[Callable[[np.ndarray[Any, Any]], Any], str]
    ] = dc.field(init=False)
    iter_fn: Mapping[str, Callable[[np.ndarray[Any, Any]], Any]]

    def __post_init__(self, aggregation: InitAggregation) -> None:
        self._aggregation = dict()
        for rename, agg in aggregation.items():
            if callable(agg):
                self._aggregation[rename] = (agg, rename)
            elif (isinstance(agg, Sequence) and len(agg) == 2
                  and callable(agg[0])):
                self._aggregation[rename] = agg
            else:
                raise ValueError(f'{rename}: {agg} is not a valid aggregation argument')

        for rename, agg in self.iter_fn.items():
            if not callable(agg):
                raise ValueError(f'{rename}: {agg} is not a valid aggregation argument')

    def result_factory(
        self, fn: Callable[[str, str], np.ndarray[Any, Any]]
    ) -> Dict[str, np.ndarray[Any, Any]]:
        """
        Factory function to generate the dict for the results set.
        A function to take the new column name and source column name
        and return an empty `np.ndarray` should be given.
        """
        return {
            rename: fn(rename, col) for rename, col in (
                *((col, col) for col in self.group_by_cols),
                *((rename, agg[1]) for rename, agg in self._aggregation.items())
            )
        }

    def iter_fn_factory(
        self, fn: Callable[[], np.ndarray[Any, Any]]
    ) -> Dict[str, np.ndarray[Any, Any]]:
        return {rename: fn() for rename in self.iter_fn.keys()}

    def apply(self, tafra: 'Tafra') -> 'Tafra':
        raise NotImplementedError


class GroupBy(AggMethod):
    """
    Aggregation by a set of unique values.

    Analogy to SQL ``GROUP BY``, not :meth:`pandas.DataFrame.groupby()`.

    Parameters
    ----------
        columns: Iterable[str]
            The column names to group by.

        aggregation: Mapping[str, Union[Callable[[np.ndarray], Any], \
            Optional. Tuple[Callable[[np.ndarray], Any], str]]]
            A mapping for columns and aggregation functions. Should be
            given as {'column': fn} or {'new_column': (fn, 'column')}.

        iter_fn: Mapping[str, Callable[[np.ndarray], Any]]
            Optional. A mapping for new columns names to the function to apply to
            the enumeration. Should be given as {'new_column': fn}.
    """

    def apply(self, tafra: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`GroupBy` to the :class:`Tafra`.

        Parameters
        ----------
            tafra: Tafra
                The tafra to apply the operation to.

        Returns
        -------
            tafra: Tafra
                The aggregated :class:`Tafra`.
        """
        self._validate(tafra, (
            *self.group_by_cols,
            *(col for (_, col) in self._aggregation.values())
        ))

        # check if all aggregations can be vectorized
        all_vectorized = (
            not self.iter_fn
            and all(id(fn) in _VECTORIZED_AGGS
                    for fn, _ in self._aggregation.values())
        )

        if all_vectorized and self._aggregation:
            # fast path: vectorized aggregation via labels, no per-group loop
            unique, labels, n_groups = self._build_group_labels(
                tafra, self.group_by_cols)

            result: Dict[str, np.ndarray[Any, Any]] = {}
            for i, col in enumerate(self.group_by_cols):
                vals = tafra._data[col]
                # pick one representative value per group
                first_occurrence = np.empty(n_groups, dtype=np.intp)
                first_occurrence[labels] = np.arange(len(labels))
                # overwrite gives last, we want first — reverse
                first_occurrence[labels[::-1]] = np.arange(
                    len(labels) - 1, -1, -1)
                result[col] = vals[first_occurrence]

            for rename, (fn, col) in self._aggregation.items():
                vec_fn = _VECTORIZED_AGGS[id(fn)]
                result[rename] = vec_fn(tafra._data[col], labels, n_groups)

            return Tafra(result)

        # standard path: per-group loop
        unique, group_indices = self._build_group_indices(
            tafra, self.group_by_cols)
        n_groups = len(unique)

        result = self.result_factory(
            lambda rename, col: np.empty(n_groups, dtype=tafra._data[col].dtype))
        iter_fn = self.iter_fn_factory(lambda: np.ones(n_groups, dtype=int))

        for i, (u, rows) in enumerate(zip(unique, group_indices)):
            for val, col in zip(u, self.group_by_cols):
                result[col][i] = val

            for rename, (fn, col) in self._aggregation.items():
                result[rename][i] = fn(tafra._data[col][rows])

            for rename, fn in self.iter_fn.items():
                iter_fn[rename][i] = fn(np.full(len(rows), i, dtype=int))

        result.update(iter_fn)
        return Tafra(result)


class Transform(AggMethod):
    """
    Apply a function to each unique set of values and join to the original table.

    Analogy to :meth:`pandas.DataFrame.groupby().transform()`,
    i.e. a SQL ``GROUP BY`` and ``LEFT JOIN`` back to the original table.

    Parameters
    ----------
        group_by: Iterable[str]
            The column names to group by.

        aggregation: Mapping[str, Union[Callable[[np.ndarray], Any], \
        Tuple[Callable[[np.ndarray], Any], str]]]
            Optional. A mapping for columns and aggregation functions. Should be
            given as {'column': fn} or {'new_column': (fn, 'column')}.

        iter_fn: Mapping[str, Callable[[np.ndarray], Any]]
            Optional. A mapping for new columns names to the function to apply to
            the enumeration. Should be given as {'new_column': fn}.
    """

    def apply(self, tafra: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`Transform` to the :class:`Tafra`.

        Parameters
        ----------
            tafra: Tafra
                The tafra to apply the operation to.

        Returns
        -------
            tafra: Tafra
                The transformed :class:`Tafra`.
        """
        self._validate(tafra, (
            *self.group_by_cols,
            *(col for (_, col) in self._aggregation.values())
        ))

        # check if all aggregations can be vectorized
        all_vectorized = (
            not self.iter_fn
            and all(id(fn) in _VECTORIZED_AGGS
                    for fn, _ in self._aggregation.values())
        )

        if all_vectorized and self._aggregation:
            # fast path: compute per-group aggregates, broadcast via labels
            # Transform doesn't need first-seen order — use direct mapping
            data, _ = self._prepare_keys(tafra, self.group_by_cols)
            labels, n_groups = self._direct_labels_sorted(data)

            result: Dict[str, np.ndarray[Any, Any]] = {}
            for col in self.group_by_cols:
                result[col] = tafra._data[col].copy()

            for rename, (fn, col) in self._aggregation.items():
                vec_fn = _VECTORIZED_AGGS[id(fn)]
                group_values = vec_fn(tafra._data[col], labels, n_groups)
                result[rename] = group_values[labels]

            return Tafra(result)

        # standard path: per-group loop
        unique, group_indices = self._build_group_indices(
            tafra, self.group_by_cols)

        result = self.result_factory(
            lambda rename, col: np.empty_like(tafra._data[col]))
        iter_fn = self.iter_fn_factory(lambda: np.ones(tafra._rows, dtype=int))

        for i, (u, rows) in enumerate(zip(unique, group_indices)):
            for col in self.group_by_cols:
                result[col][rows] = tafra._data[col][rows]

            for rename, (fn, col) in self._aggregation.items():
                result[rename][rows] = fn(tafra._data[col][rows])

            for rename, fn in self.iter_fn.items():
                iter_fn[rename][rows] = fn(np.full(len(rows), i, dtype=int))

        result.update(iter_fn)
        return Tafra(result)


@dc.dataclass
class IterateBy(GroupSet):
    """
    A generator that yields a :class:`Tafra` for each set of unique values.

    Analogy to `pandas.DataFrame.groupby()`, i.e. an Sequence of `Tafra` objects.
    Yields tuples of ((unique grouping values, ...), row indices array, subset tafra)

    Parameters
    ----------
        group_by: Iterable[str]
            The column names to group by.
    """
    group_by_cols: Iterable[str]

    def apply(self, tafra: 'Tafra') -> Iterator[GroupDescription]:
        """
        Apply the :class:`IterateBy` to the :class:`Tafra`.

        Parameters
        ----------
            tafra: Tafra
                The tafra to apply the operation to.

        Returns
        -------
            tafras: Iterator[GroupDescription]
                An iterator over the grouped :class:`Tafra`.
        """
        self._validate(tafra, self.group_by_cols)
        unique, group_indices = self._build_group_indices(tafra, self.group_by_cols)

        for u, rows in zip(unique, group_indices):
            yield (u, rows, tafra._ndindex(rows))


@dc.dataclass
class Join(GroupSet):
    """
    Base class for SQL-like JOINs.
    """
    on: Iterable[Tuple[str, str, str]]
    select: Iterable[str]

    def _validate_dtypes(self, l_table: 'Tafra', r_table: 'Tafra') -> None:
        for l_column, r_column, _ in self.on:
            l_value = l_table._data[l_column]
            r_value = r_table._data[r_column]
            l_dtype = l_table._dtypes[l_column]
            r_dtype = r_table._dtypes[r_column]

            if l_value.dtype != r_value.dtype:
                raise TypeError(
                    f'This `Tafra` column `{l_column}` dtype `{l_value.dtype}` '
                    f'does not match other `Tafra` dtype `{r_value.dtype}`.')

            # should not happen unless dtypes manually changed, but let's check it
            elif l_dtype != r_dtype:
                raise TypeError(
                    f'This `Tafra` column `{l_column}` dtype `{l_dtype}` '
                    f'does not match other `Tafra` dtype `{r_dtype}`.')

    @staticmethod
    def _validate_ops(ops: Iterable[str]) -> None:
        for op in ops:
            _op = JOIN_OPS.get(op, None)
            if _op is None:
                raise TypeError(f'The operator {op} is not valid.')

    @staticmethod
    def _build_composite_key(
        cols: List[np.ndarray[Any, Any]]
    ) -> np.ndarray[Any, Any]:
        """Build a single sortable key array from multiple columns."""
        if len(cols) == 1:
            return cols[0]
        # structured array for multi-column sort
        dt = np.dtype([(f'f{i}', c.dtype) for i, c in enumerate(cols)])
        key = np.empty(len(cols[0]), dtype=dt)
        for i, c in enumerate(cols):
            key[f'f{i}'] = c
        return key

    @staticmethod
    def _sort_merge_indices(
        left_key: np.ndarray[Any, Any],
        right_key: np.ndarray[Any, Any],
    ) -> Tuple[np.ndarray[Any, Any], np.ndarray[Any, Any]]:
        """
        Compute inner-join index pairs using sort-merge.
        Returns (left_indices, right_indices) as intp arrays.
        """
        right_order = np.argsort(right_key, kind='stable')
        right_sorted = right_key[right_order]

        left_lo = np.searchsorted(right_sorted, left_key, side='left')
        left_hi = np.searchsorted(right_sorted, left_key, side='right')
        counts = left_hi - left_lo

        total = int(counts.sum())
        if total == 0:
            return np.array([], dtype=np.intp), np.array([], dtype=np.intp)

        # left indices: repeat each left row index by its match count
        li = np.repeat(np.arange(len(left_key), dtype=np.intp), counts)

        # right indices: for each left row, enumerate its match range
        # build offsets into right_sorted for each output row
        offsets = np.repeat(left_lo, counts)
        group_starts = np.cumsum(counts) - counts
        within = np.arange(total, dtype=np.intp) - np.repeat(group_starts, counts)
        ri = right_order[offsets + within]

        return li, ri

    @staticmethod
    def _left_join_indices(
        left_key: np.ndarray[Any, Any],
        right_key: np.ndarray[Any, Any],
    ) -> Tuple[np.ndarray[Any, Any], np.ndarray[Any, Any], bool]:
        """
        Compute left-join index pairs using sort-merge.
        Unmatched left rows get right index = -1.
        Returns (left_indices, right_indices, has_null).
        """
        right_order = np.argsort(right_key, kind='stable')
        right_sorted = right_key[right_order]

        left_lo = np.searchsorted(right_sorted, left_key, side='left')
        left_hi = np.searchsorted(right_sorted, left_key, side='right')
        counts = left_hi - left_lo

        # unmatched left rows get count=0 → force to 1 (with sentinel)
        unmatched = counts == 0
        has_null = bool(np.any(unmatched))
        out_counts = np.where(unmatched, 1, counts)

        total = int(out_counts.sum())
        li = np.repeat(np.arange(len(left_key), dtype=np.intp), out_counts)

        # right indices
        ri = np.empty(total, dtype=np.intp)
        pos = 0
        if has_null:
            # mixed matched/unmatched — fill per group
            for i in range(len(left_key)):
                c = int(counts[i])
                if c > 0:
                    ri[pos:pos + c] = right_order[left_lo[i]:left_hi[i]]
                    pos += c
                else:
                    ri[pos] = -1
                    pos += 1
        else:
            # all matched — fully vectorized
            offsets = np.repeat(left_lo, counts)
            group_starts = np.cumsum(counts) - counts
            within = np.arange(total, dtype=np.intp) - np.repeat(
                group_starts, counts)
            ri = right_order[offsets + within]

        return li, ri, has_null

    def _resolve_join_cols(
        self, left_t: 'Tafra', right_t: 'Tafra'
    ) -> Tuple[List[str], Dict[str, str]]:
        """Compute deduplicated output columns and dtypes."""
        seen_cols: Dict[str, None] = {}
        for c in chain(left_t._data.keys(), right_t._data.keys()):
            if not self.select or c in self.select:
                seen_cols[c] = None
        join_cols = list(seen_cols.keys())
        dtypes: Dict[str, str] = {
            c: d for c, d in chain(
                right_t._dtypes.items(), left_t._dtypes.items())
            if c in seen_cols
        }
        return join_cols, dtypes

    def apply(self, left_t: 'Tafra', right_t: 'Tafra') -> 'Tafra':
        raise NotImplementedError


class InnerJoin(Join):
    """
    An inner join.

    Analogy to SQL INNER JOIN, or `pandas.merge(..., how='inner')`,

    Parameters
    ----------
        right: Tafra
            The right-side :class:`Tafra` to join.

        on: Iterable[Tuple[str, str, str]]
            The columns and operator to join on. Should be given as
            ('left column', 'right column', 'op') Valid ops are:

            '==' : equal to
            '!=' : not equal to
            '<'  : less than
            '<=' : less than or equal to
            '>'  : greater than
            '>=' : greater than or equal to

        select: Iterable[str] = []
            The columns to return. If not given, all unique columns names
            are returned. If the column exists in both :class`Tafra`,
            prefers the left over the right.
    """

    def apply(self, left_t: 'Tafra', right_t: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`InnerJoin` to the :class:`Tafra`.

        Parameters
        ----------
            left_t: Tafra
                The left tafra to join.

            right_t: Tafra
                The right tafra to join.

        Returns
        -------
            tafra: Tafra
                The joined :class:`Tafra`.
        """
        left_cols, right_cols, ops = list(zip(*self.on))
        self._validate(left_t, left_cols)
        self._validate(right_t, right_cols)
        self._validate_dtypes(left_t, right_t)
        self._validate_ops(ops)

        join_cols, dtypes = self._resolve_join_cols(left_t, right_t)
        all_equi = all(op_str == '==' for _, _, op_str in self.on)

        if all_equi:
            # Encode left+right together for consistent codebooks
            left_cols_data = [left_t._data[lc] for lc, _, _ in self.on]
            right_cols_data = [right_t._data[rc] for _, rc, _ in self.on]
            l_enc, r_enc = GroupSet._encode_columns_paired(
                left_cols_data, right_cols_data)
            left_key = GroupSet._build_composite_key(l_enc)
            right_key = GroupSet._build_composite_key(r_enc)

            if _HAS_ACCEL:
                li, ri = _c_inner_join(
                    np.ascontiguousarray(left_key, dtype=np.int64),
                    np.ascontiguousarray(right_key, dtype=np.int64))
            else:
                li, ri = self._sort_merge_indices(left_key, right_key)

            if len(li) == 0:
                return Tafra(
                    {c: np.array(
                        [], dtype=left_t._data[c].dtype
                        if c in left_t._data
                        else right_t._data[c].dtype)
                     for c in join_cols},
                    dtypes
                )

            result: Dict[str, np.ndarray[Any, Any]] = {}
            for c in join_cols:
                if c in left_t._data:
                    result[c] = left_t._data[c][li]
                else:
                    result[c] = right_t._data[c][ri]

            return Tafra(result, dtypes)

        else:
            _on = tuple(
                (left_col, right_col, JOIN_OPS[op])
                for left_col, right_col, op in self.on
            )
            right_rows = np.empty(right_t._rows, dtype=bool)
            join: Dict[str, List[Any]] = {c: [] for c in join_cols}

            for i in range(left_t._rows):
                right_rows[:] = True
                for left_col, right_col, op in _on:
                    right_rows &= op(
                        left_t._data[left_col][i], right_t._data[right_col])

                right_count = int(np.sum(right_rows))
                if right_count <= 0:
                    continue

                for column in join_cols:
                    if column in left_t._data:
                        join[column].extend(
                            [left_t._data[column][i]] * right_count)
                    elif column in right_t._data:
                        join[column].extend(
                            right_t._data[column][right_rows])

            return Tafra(
                {c: np.array(v) for c, v in join.items()},
                dtypes
            )


class LeftJoin(Join):
    """
    A left join.

    Analogy to SQL LEFT JOIN, or `pandas.merge(..., how='left')`,

    Parameters
    ----------
        right: Tafra
            The right-side :class:`Tafra` to join.

        on: Iterable[Tuple[str, str, str]]
            The columns and operator to join on. Should be given as
            ('left column', 'right column', 'op') Valid ops are:

            '==' : equal to
            '!=' : not equal to
            '<'  : less than
            '<=' : less than or equal to
            '>'  : greater than
            '>=' : greater than or equal to

        select: Iterable[str] = []
            The columns to return. If not given, all unique columns names
            are returned. If the column exists in both :class`Tafra`,
            prefers the left over the right.
    """

    def apply(self, left_t: 'Tafra', right_t: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`LeftJoin` to the :class:`Tafra`.

        Parameters
        ----------
            left_t: Tafra
                The left tafra to join.

            right_t: Tafra
                The right tafra to join.

        Returns
        -------
            tafra: Tafra
                The joined :class:`Tafra`.
        """
        left_cols, right_cols, ops = list(zip(*self.on))
        self._validate(left_t, left_cols)
        self._validate(right_t, right_cols)
        self._validate_dtypes(left_t, right_t)
        self._validate_ops(ops)

        join_cols, dtypes = self._resolve_join_cols(left_t, right_t)
        all_equi = all(op_str == '==' for _, _, op_str in self.on)

        if all_equi:
            left_cols_data = [left_t._data[lc] for lc, _, _ in self.on]
            right_cols_data = [right_t._data[rc] for _, rc, _ in self.on]
            l_enc, r_enc = GroupSet._encode_columns_paired(
                left_cols_data, right_cols_data)
            left_key = GroupSet._build_composite_key(l_enc)
            right_key = GroupSet._build_composite_key(r_enc)

            if _HAS_ACCEL:
                li, ri, has_null = _c_left_join(
                    np.ascontiguousarray(left_key, dtype=np.int64),
                    np.ascontiguousarray(right_key, dtype=np.int64))
            else:
                li, ri, has_null = self._left_join_indices(left_key, right_key)

            if has_null:
                for c in join_cols:
                    if c not in left_t._data and dtypes.get(c) != 'object':
                        dtypes[c] = 'object'

            result: Dict[str, np.ndarray[Any, Any]] = {}
            matched = ri >= 0
            for c in join_cols:
                if c in left_t._data:
                    result[c] = left_t._data[c][li]
                else:
                    # right column: fill matched rows, None for unmatched
                    if has_null:
                        out = np.empty(len(li), dtype=object)
                        out[matched] = right_t._data[c][ri[matched]]
                        out[~matched] = None
                        result[c] = out
                    else:
                        result[c] = right_t._data[c][ri]

            return Tafra(result, dtypes)

        else:
            _on = tuple(
                (left_col, right_col, JOIN_OPS[op])
                for left_col, right_col, op in self.on
            )
            right_rows = np.empty(right_t._rows, dtype=bool)
            join: Dict[str, List[Any]] = {c: [] for c in join_cols}

            for i in range(left_t._rows):
                right_rows[:] = True
                for left_col, right_col, op in _on:
                    right_rows &= op(
                        left_t._data[left_col][i], right_t._data[right_col])

                right_count = int(np.sum(right_rows))

                for column in join_cols:
                    if column in left_t._data:
                        join[column].extend(
                            [left_t._data[column][i]] * max(1, right_count))
                    elif column in right_t._data:
                        if right_count <= 0:
                            join[column].append(None)
                            if dtypes[column] != 'object':
                                dtypes[column] = 'object'
                        else:
                            join[column].extend(
                                right_t._data[column][right_rows])

            return Tafra(
                {c: np.array(v) for c, v in join.items()},
                dtypes
            )


@dc.dataclass
class CrossJoin(Join):
    """
    A cross join.

    Analogy to SQL CROSS JOIN, or `pandas.merge(..., how='outer')
    using temporary columns of static value to intersect all rows`.

    Parameters
    ----------
        right: Tafra
            The right-side :class:`Tafra` to join.

        select: Iterable[str] = []
            The columns to return. If not given, all unique columns names
            are returned. If the column exists in both :class`Tafra`,
            prefers the left over the right.
    """

    def apply(self, left_t: 'Tafra', right_t: 'Tafra') -> 'Tafra':
        """
        Apply the :class:`CrossJoin` to the :class:`Tafra`.

        Parameters
        ----------
            left_t: Tafra
                The left tafra to join.

            right_t: Tafra
                The right tafra to join.

        Returns
        -------
            tafra: Tafra
                The joined :class:`Tafra`.
        """
        self._validate_dtypes(left_t, right_t)

        left_rows = left_t._rows
        right_rows = right_t._rows

        select = set(self.select)
        if len(select) > 0:
            left_cols = list(select.intersection(list(left_t._data.keys())))
            right_cols = list(select.intersection(list(right_t._data.keys())))

            if len(left_cols) == 0:
                raise IndexError('No columns given to select from left `Tafra`.')
            if len(right_cols) == 0:
                raise IndexError('No columns given to select from right `Tafra`.')

        else:
            left_cols = list(left_t._data.keys())
            right_cols = list(right_t._data.keys())

        left_new = Tafra(left_t[left_cols].key_map(np.tile, reps=right_rows))
        right_new = Tafra(right_t[right_cols].key_map(np.tile, reps=left_rows))

        left_new.update_inplace(right_new)

        return left_new


# Import here to resolve circular dependency
from .base import Tafra
