"""Basket analysis — association rules from order data.

Finds "frequently bought together" patterns using co-occurrence counting:
    - Support: How often items appear together relative to all transactions
    - Confidence: P(B | A) — how often B appears when A is present
    - Lift: confidence / P(B) — strength of association above random chance

Backends:
    - builtin: Pure-Python pairwise co-occurrence (no dependencies)
    - mlxtend: Apriori + association rules via mlxtend (production-grade)

Usage via provider:
    from platoon.learning.providers import get_provider
    ba = get_provider("basket")
    result = ba.analyze(transactions, min_support=0.01, min_confidence=0.3)

Direct usage:
    from platoon.learning.basket import basket_analysis
    result = basket_analysis(transactions, min_support=0.01)
"""

from dataclasses import dataclass, field
from typing import Dict, List, Set

from platoon.learning.providers import ModelProvider, register_provider


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------


@dataclass
class AssociationRule:
    """A single association rule A → B."""

    antecedent: str  # item A
    consequent: str  # item B
    support: float  # P(A ∩ B)
    confidence: float  # P(B | A)
    lift: float  # confidence / P(B)


@dataclass
class BasketResult:
    """Output of basket analysis."""

    rules: List[AssociationRule]
    total_transactions: int
    method: str


# ---------------------------------------------------------------------------
# Pure-Python implementation
# ---------------------------------------------------------------------------


def basket_analysis(
    transactions: List[Set[str]],
    min_support: float = 0.01,
    min_confidence: float = 0.3,
) -> BasketResult:
    """Compute association rules from transaction data.

    Uses single-pass pairwise co-occurrence counting. Generates both
    A→B and B→A rules for qualifying pairs.

    Args:
        transactions: List of sets, each set containing item IDs in one transaction.
        min_support: Minimum support threshold (fraction of transactions).
        min_confidence: Minimum confidence threshold.

    Returns:
        BasketResult with sorted rules (by lift descending).
    """
    n = len(transactions)
    if n == 0:
        return BasketResult(rules=[], total_transactions=0, method="builtin_cooccurrence")

    # Count item frequencies
    item_counts: Dict[str, int] = {}
    for txn in transactions:
        for item in txn:
            item_counts[item] = item_counts.get(item, 0) + 1

    # Count pairwise co-occurrences
    pair_counts: Dict[tuple, int] = {}
    for txn in transactions:
        items = sorted(txn)
        for i in range(len(items)):
            for j in range(i + 1, len(items)):
                pair = (items[i], items[j])
                pair_counts[pair] = pair_counts.get(pair, 0) + 1

    # Generate rules
    rules: List[AssociationRule] = []
    min_count = max(1, int(min_support * n))

    for (a, b), count in pair_counts.items():
        support = count / n
        if count < min_count:
            continue

        # A → B
        conf_ab = count / item_counts[a]
        if conf_ab >= min_confidence:
            lift_ab = conf_ab / (item_counts[b] / n)
            rules.append(AssociationRule(
                antecedent=a,
                consequent=b,
                support=round(support, 4),
                confidence=round(conf_ab, 4),
                lift=round(lift_ab, 4),
            ))

        # B → A
        conf_ba = count / item_counts[b]
        if conf_ba >= min_confidence:
            lift_ba = conf_ba / (item_counts[a] / n)
            rules.append(AssociationRule(
                antecedent=b,
                consequent=a,
                support=round(support, 4),
                confidence=round(conf_ba, 4),
                lift=round(lift_ba, 4),
            ))

    rules.sort(key=lambda r: r.lift, reverse=True)

    return BasketResult(
        rules=rules,
        total_transactions=n,
        method="builtin_cooccurrence",
    )


# ---------------------------------------------------------------------------
# Built-in provider
# ---------------------------------------------------------------------------


@register_provider
class BuiltinBasketProvider(ModelProvider):
    """Pure-Python basket analysis — pairwise co-occurrence counting."""

    name = "builtin_basket"
    domain = "basket"
    backend = "builtin"

    def analyze(
        self,
        transactions: List[Set[str]],
        min_support: float = 0.01,
        min_confidence: float = 0.3,
    ) -> BasketResult:
        """Compute association rules from transaction data."""
        return basket_analysis(transactions, min_support=min_support, min_confidence=min_confidence)


# ---------------------------------------------------------------------------
# Library stub
# ---------------------------------------------------------------------------


@register_provider
class MlxtendBasketProvider(ModelProvider):
    """Apriori + association rules via mlxtend.

    TODO: Implement analyze() using mlxtend.frequent_patterns.apriori
    and mlxtend.frequent_patterns.association_rules.
    Requires: mlxtend
    """

    name = "mlxtend_basket"
    domain = "basket"
    backend = "mlxtend"

    @classmethod
    def is_available(cls) -> bool:
        try:
            import mlxtend  # noqa: F401
            return True
        except ImportError:
            return False

    def analyze(
        self,
        transactions: List[Set[str]],
        min_support: float = 0.01,
        min_confidence: float = 0.3,
    ) -> BasketResult:
        """Compute association rules using mlxtend Apriori."""
        raise NotImplementedError("mlxtend basket provider not yet implemented")
