Association Rules (Apriori / FP-Growth)

"""Association rules by Apriori - support counting and rule generation, literally.

A transaction is a SET of items. support(S) is the fraction of transactions that
contain S, and every other quantity here is a ratio of supports. The search is
level-wise: frequent k-itemsets generate (k+1)-candidates, a candidate with any
infrequent subset is discarded unseen (downward closure), and the survivors are
counted in one pass over the data.
"""

from itertools import combinations


def support(itemset, transactions):
    """Fraction of transactions that contain every item of itemset."""
    containing = 0
    for transaction in transactions:
        if itemset <= transaction:
            containing += 1
    return containing / len(transactions)


def frequent_itemsets(transactions, min_support):
    """Return {frozenset: support} for every itemset with support >= min_support."""
    items = set()
    for transaction in transactions:
        for item in transaction:
            items.add(item)

    frequent = {}
    level = []
    for item in sorted(items):
        candidate = frozenset([item])
        candidate_support = support(candidate, transactions)
        if candidate_support >= min_support:
            frequent[candidate] = candidate_support
            level.append(candidate)

    size = 1
    while level:
        # Join: two frequent k-itemsets whose union has k+1 items form a candidate.
        candidates = set()
        for left in level:
            for right in level:
                union = left | right
                if len(union) == size + 1:
                    candidates.add(union)

        next_level = []
        for candidate in candidates:
            # Prune by downward closure: every k-subset must already be frequent.
            all_subsets_frequent = True
            for item in candidate:
                if (candidate - {item}) not in frequent:
                    all_subsets_frequent = False
                    break
            if not all_subsets_frequent:
                continue

            candidate_support = support(candidate, transactions)
            if candidate_support >= min_support:
                frequent[candidate] = candidate_support
                next_level.append(candidate)

        level = next_level
        size += 1

    return frequent


def generate_rules(frequent, min_confidence):
    """Split each frequent itemset into every antecedent => consequent pair."""
    rules = []
    for itemset, itemset_support in frequent.items():
        if len(itemset) < 2:
            continue

        members = sorted(itemset)
        for size in range(1, len(members)):
            for antecedent_items in combinations(members, size):
                antecedent = frozenset(antecedent_items)
                consequent = itemset - antecedent

                # Downward closure guarantees both subsets were counted.
                confidence = itemset_support / frequent[antecedent]
                if confidence < min_confidence:
                    continue

                lift = confidence / frequent[consequent]
                if confidence == 1.0:
                    conviction = float("inf")
                else:
                    conviction = (1.0 - frequent[consequent]) / (1.0 - confidence)

                # Jaccard of the two transaction-id sets: the symmetric relative
                # of confidence, normalized by the union rather than by X alone.
                jaccard = itemset_support / (
                    frequent[antecedent] + frequent[consequent] - itemset_support
                )

                rules.append(
                    {
                        "antecedent": antecedent,
                        "consequent": consequent,
                        "support": itemset_support,
                        "confidence": confidence,
                        "lift": lift,
                        "conviction": conviction,
                        "jaccard": jaccard,
                    }
                )

    return rules

"""Association rules - typed Apriori with prefix join, subset prune, immutable rules."""

from __future__ import annotations

import math
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from itertools import combinations

Itemset = tuple[int, ...]   # sorted ascending, no duplicates


@dataclass(frozen=True)
class Rule:
    antecedent: Itemset
    consequent: Itemset
    support: float
    confidence: float
    lift: float
    conviction: float
    jaccard: float


@dataclass(frozen=True)
class MiningResult:
    counts: dict[Itemset, int]
    n_transactions: int
    rules: tuple[Rule, ...]


def _join(level: Sequence[Itemset]) -> list[Itemset]:
    """F_k x F_k join on a shared (k-1)-prefix.

    Items are kept sorted and the level is lexicographically sorted, so every
    candidate is generated exactly once, and the scan can stop at the first
    tuple whose prefix differs rather than comparing all pairs.
    """
    candidates: list[Itemset] = []
    for index, left in enumerate(level):
        for right in level[index + 1 :]:
            if left[:-1] != right[:-1]:
                break
            candidates.append(left + (right[-1],))
    return candidates


def _all_subsets_frequent(candidate: Itemset, previous: set[Itemset]) -> bool:
    """Downward closure: every k-subset of a (k+1)-candidate must be frequent."""
    return all(
        candidate[:skip] + candidate[skip + 1 :] in previous
        for skip in range(len(candidate))
    )


def frequent_itemsets(
    transactions: Sequence[frozenset[int]],
    min_support: float,
    max_length: int = 4,
) -> dict[Itemset, int]:
    """Counts of every itemset with support >= min_support. Raises ValueError."""
    if not transactions:
        raise ValueError("need at least one transaction")
    if not 0.0 < min_support <= 1.0:
        raise ValueError(f"min_support must lie in (0, 1], got {min_support}")
    if max_length < 1:
        raise ValueError(f"max_length must be at least 1, got {max_length}")

    min_count = math.ceil(min_support * len(transactions))

    singles: dict[int, int] = {}
    for transaction in transactions:
        for item in transaction:
            singles[item] = singles.get(item, 0) + 1

    counts: dict[Itemset, int] = {
        (item,): count for item, count in sorted(singles.items()) if count >= min_count
    }
    level: list[Itemset] = list(counts)

    # Bounded by max_length: the lattice is exponential, so depth is a guard.
    for _ in range(max_length - 1):
        if not level:
            break

        previous = set(level)
        next_level: list[Itemset] = []
        for candidate in _join(level):
            if not _all_subsets_frequent(candidate, previous):
                continue
            candidate_set = frozenset(candidate)
            count = sum(1 for transaction in transactions if candidate_set <= transaction)
            if count >= min_count:
                counts[candidate] = count
                next_level.append(candidate)
        level = next_level

    return counts


def generate_rules(
    counts: dict[Itemset, int],
    n_transactions: int,
    min_confidence: float,
    min_lift: float = 1.0,
) -> tuple[Rule, ...]:
    """Every rule clearing both floors, ranked by lift then confidence."""
    if not 0.0 <= min_confidence <= 1.0:
        raise ValueError(f"min_confidence must lie in [0, 1], got {min_confidence}")

    rules: list[Rule] = []
    for itemset, itemset_count in counts.items():
        if len(itemset) < 2:
            continue

        for size in range(1, len(itemset)):
            for antecedent in combinations(itemset, size):
                consequent = tuple(item for item in itemset if item not in antecedent)
                antecedent_count = counts.get(antecedent)
                consequent_count = counts.get(consequent)
                if antecedent_count is None or consequent_count is None:
                    continue   # unreachable under downward closure; fail soft

                confidence = itemset_count / antecedent_count
                consequent_support = consequent_count / n_transactions
                lift = confidence / consequent_support
                if confidence < min_confidence or lift < min_lift:
                    continue

                conviction = (
                    math.inf
                    if confidence >= 1.0
                    else (1.0 - consequent_support) / (1.0 - confidence)
                )
                rules.append(
                    Rule(
                        antecedent=antecedent,
                        consequent=consequent,
                        support=itemset_count / n_transactions,
                        confidence=confidence,
                        lift=lift,
                        conviction=conviction,
                        jaccard=itemset_count
                        / (antecedent_count + consequent_count - itemset_count),
                    )
                )

    rules.sort(key=lambda rule: (rule.lift, rule.confidence), reverse=True)
    return tuple(rules)


def mine(
    baskets: Iterable[Iterable[int]],
    min_support: float,
    min_confidence: float,
    max_length: int = 4,
) -> MiningResult:
    transactions = [frozenset(basket) for basket in baskets]
    counts = frequent_itemsets(transactions, min_support, max_length)
    rules = generate_rules(counts, len(transactions), min_confidence)
    return MiningResult(counts=counts, n_transactions=len(transactions), rules=rules)

"""Association rules - vertical bitsets, depth-first Eclat, preallocated scratch.

Horizontal Apriori rescans every transaction for every candidate. Storing the
data VERTICALLY inverts that: item i keeps one bit per transaction saying
whether it is present, so the support of an itemset is the popcount of the AND
of its items' bitsets. No candidate is counted by scanning the data again.

Needs NumPy >= 2.0 for np.bitwise_count.
"""

from itertools import combinations

import numpy as np
from numpy.typing import NDArray

Bitsets = NDArray[np.uint64]


def pack_vertical(transactions: list[list[int]], n_items: int) -> Bitsets:
    """(n_items, words): row i is item i's transaction-id bitset."""
    n = len(transactions)
    lengths = np.fromiter((len(t) for t in transactions), dtype=np.int64, count=n)
    rows = np.concatenate([np.asarray(t, dtype=np.int64) for t in transactions])
    columns = np.repeat(np.arange(n, dtype=np.int64), lengths)

    words = (n + 63) // 64
    dense = np.zeros((n_items, words * 64), dtype=np.uint8)
    dense[rows, columns] = 1          # scatter every (item, transaction) at once
    packed = np.packbits(dense, axis=1, bitorder="little")
    return np.ascontiguousarray(packed).view(np.uint64)


class EclatMiner:
    """Depth-first frequent-itemset mining over a vertical bitset index."""

    def __init__(self, vertical: Bitsets, min_count: int, max_length: int) -> None:
        if max_length < 1:
            raise ValueError(f"max_length must be at least 1, got {max_length}")
        self._vertical = np.ascontiguousarray(vertical, dtype=np.uint64)
        self._min_count = min_count
        self._max_length = max_length
        # One scratch bitset per depth, allocated once: the search never
        # allocates, and a child never overwrites its parent's buffer.
        self._scratch = np.empty(
            (max_length + 1, self._vertical.shape[1]), dtype=np.uint64
        )
        self.found: dict[tuple[int, ...], int] = {}

    def mine(self) -> dict[tuple[int, ...], int]:
        counts = np.bitwise_count(self._vertical).sum(axis=1)
        frequent = [int(item) for item in np.flatnonzero(counts >= self._min_count)]

        for position, item in enumerate(frequent):
            self.found[(item,)] = int(counts[item])
            np.copyto(self._scratch[1], self._vertical[item])
            self._extend((item,), 1, frequent[position + 1 :])
        return self.found

    def _extend(self, prefix: tuple[int, ...], depth: int, tail: list[int]) -> None:
        if depth >= self._max_length or not tail:
            return

        parent = self._scratch[depth]
        child = self._scratch[depth + 1]

        # Pass 1: which tail items survive? Each AND writes into the SAME child
        # buffer and only the count is kept, so no per-item array is retained.
        survivors: list[int] = []
        for item in tail:
            np.bitwise_and(parent, self._vertical[item], out=child)
            if int(np.bitwise_count(child).sum()) >= self._min_count:
                survivors.append(item)

        # Pass 2: recompute each survivor's bitset and recurse. Siblings share
        # the buffer, so the intersection is redone rather than stored.
        for position, item in enumerate(survivors):
            np.bitwise_and(parent, self._vertical[item], out=child)
            extended = prefix + (item,)
            self.found[extended] = int(np.bitwise_count(child).sum())
            self._extend(extended, depth + 1, survivors[position + 1 :])


def generate_rules(
    counts: dict[tuple[int, ...], int],
    n_transactions: int,
    min_confidence: float,
    min_lift: float = 1.0,
) -> list[tuple[tuple[int, ...], tuple[int, ...], float, float]]:
    """(antecedent, consequent, confidence, lift); output-bound, so left in Python."""
    rules = []
    for itemset, itemset_count in counts.items():
        for size in range(1, len(itemset)):
            for antecedent in combinations(itemset, size):
                consequent = tuple(item for item in itemset if item not in antecedent)
                confidence = itemset_count / counts[antecedent]
                lift = confidence * n_transactions / counts[consequent]
                if confidence >= min_confidence and lift >= min_lift:
                    rules.append((antecedent, consequent, confidence, lift))
    rules.sort(key=lambda rule: rule[3], reverse=True)
    return rules