src/constraints.py

"""Constraint summaries and greedy-swap refinement for formed groups."""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np
import pandas as pd

from . import group_eval
from .adapters.base import DatasetSchema
from .config import FAIRNESS_TVD_MAX, GROUP_SIZE, MAX_SWAP_ITERS, SIZE_TOLERANCE


def tv_distance(p: pd.Series, q: pd.Series) -> float:
    keys = set(p.index) | set(q.index)
    return float(0.5 * sum(abs(p.get(key, 0.0) - q.get(key, 0.0)) for key in keys))


def fairness_violation(group_df: pd.DataFrame, class_df: pd.DataFrame, attr: str) -> float:
    if attr not in group_df.columns or attr not in class_df.columns:
        return 0.0
    group_dist = group_df[attr].value_counts(normalize=True)
    class_dist = class_df[attr].value_counts(normalize=True)
    return max(0.0, tv_distance(group_dist, class_dist) - FAIRNESS_TVD_MAX)


def size_violations(
    groups: list[list[int]],
    G: int = GROUP_SIZE,
    tolerance: int = SIZE_TOLERANCE,
) -> int:
    lo, hi = G - tolerance, G + tolerance
    return int(sum(not (lo <= len(group) <= hi) for group in groups))


def cluster_complementarity_violations(
    groups: list[list[int]],
    labels: np.ndarray,
    min_unique_clusters: int = 2,
) -> int:
    labels = np.asarray(labels)
    violations = 0
    for group in groups:
        present = set(labels[group].tolist()) - {-1}
        if len(group) > 1 and len(present) < min_unique_clusters:
            violations += 1
    return int(violations)


def _group_unique_clusters(group: list[int], labels: np.ndarray) -> int:
    return len(set(labels[group].tolist()) - {-1})


def _max_unique_clusters_for_group(group: list[int], n_clusters: int, G: int) -> int:
    return min(len(group), G, n_clusters) if n_clusters > 0 else 0


def _would_reduce_max_diversity(
    old_unique: int,
    new_unique: int,
    group: list[int],
    n_clusters: int,
    G: int,
) -> bool:
    max_unique = _max_unique_clusters_for_group(group, n_clusters, G)
    return max_unique > 0 and old_unique == max_unique and new_unique < old_unique


def greedy_swap_complementarity(
    groups: list[list[int]],
    labels: np.ndarray,
    min_unique_clusters: int = 2,
    max_iters: int = MAX_SWAP_ITERS,
) -> list[list[int]]:
    """Reduce cluster-complementarity violations via greedy pairwise swaps.

    For each violating group, try swapping one of its members with a member
    from a different group such that both groups gain or maintain unique-cluster
    count. Iterates until no violation remains or the iteration budget is
    exhausted.
    """
    labels = np.asarray(labels)
    groups = [list(g) for g in groups]
    n_groups = len(groups)

    for _ in range(max_iters):
        # find the first violating group
        violating = [
            gi for gi in range(n_groups)
            if len(groups[gi]) > 1
            and _group_unique_clusters(groups[gi], labels) < min_unique_clusters
        ]
        if not violating:
            break

        improved = False
        for gi in violating:
            current_score = _group_unique_clusters(groups[gi], labels)
            best_gain = 0
            best_swap = None  # (gi, pos_i, gj, pos_j)

            for gj in range(n_groups):
                if gi == gj:
                    continue
                donor_score = _group_unique_clusters(groups[gj], labels)
                for pos_i, idx_i in enumerate(groups[gi]):
                    for pos_j, idx_j in enumerate(groups[gj]):
                        # simulate swap
                        trial_gi = [idx_j if k == pos_i else v for k, v in enumerate(groups[gi])]
                        trial_gj = [idx_i if k == pos_j else v for k, v in enumerate(groups[gj])]
                        new_gi = _group_unique_clusters(trial_gi, labels)
                        new_gj = _group_unique_clusters(trial_gj, labels)
                        gain = (new_gi - current_score) + (new_gj - donor_score)
                        if gain > best_gain:
                            best_gain = gain
                            best_swap = (gi, pos_i, gj, pos_j)

            if best_swap is not None:
                gi2, pos_i, gj2, pos_j = best_swap
                groups[gi2][pos_i], groups[gj2][pos_j] = groups[gj2][pos_j], groups[gi2][pos_i]
                improved = True
                break

        if not improved:
            break  # stuck — no beneficial swap exists

    return groups


def summarize_demo_constraints(
    groups: list[list[int]],
    labels: np.ndarray,
    G: int = GROUP_SIZE,
    feature_df: pd.DataFrame | None = None,
    schema: DatasetSchema | None = None,
) -> dict[str, int | float | str]:
    summary: dict[str, int | float | str] = {
        "size_violations": size_violations(groups, G),
        "cluster_complementarity_violations": cluster_complementarity_violations(groups, labels),
    }
    if feature_df is not None and schema is not None:
        summary.update(normalized_penalty_terms(groups, labels, feature_df, schema, G))
        summary["fairness_refinement"] = "soft_penalty"
        summary["engagement_refinement"] = "soft_penalty"
    else:
        summary["fairness_refinement"] = "not_configured"
        summary["engagement_refinement"] = "not_configured"
    return summary


def normalized_penalty_terms(
    groups: list[list[int]],
    labels: np.ndarray,
    feature_df: pd.DataFrame,
    schema: DatasetSchema,
    G: int = GROUP_SIZE,
) -> dict[str, float]:
    """Return normalized [0, 1] penalty terms.

    Size and complementarity are fractions of violating groups. Fairness is
    mean total-variation distance, naturally bounded [0, 1]. Engagement is
    mean absolute group z-deviation divided by 3 sigma and clipped to [0, 1].
    """
    n_groups = max(1, len(groups))
    size_penalty = size_violations(groups, G) / n_groups
    complementarity_penalty = cluster_complementarity_violations(groups, labels) / n_groups

    fairness_values = [
        group_eval.demographic_fairness(feature_df[col], groups)
        for col in schema.fairness_cols
        if col in feature_df.columns
    ]
    fairness_penalty = float(np.mean(fairness_values)) if fairness_values else 0.0

    if schema.engagement_col and schema.engagement_col in feature_df.columns:
        engagement_raw = group_eval.engagement_balance(feature_df[schema.engagement_col].to_numpy(), groups)
        engagement_penalty = min(float(engagement_raw) / 3.0, 1.0)
    else:
        engagement_penalty = 0.0

    return {
        "size_penalty": float(np.clip(size_penalty, 0.0, 1.0)),
        "complementarity_penalty": float(np.clip(complementarity_penalty, 0.0, 1.0)),
        "fairness_penalty": float(np.clip(fairness_penalty, 0.0, 1.0)),
        "engagement_penalty": float(np.clip(engagement_penalty, 0.0, 1.0)),
    }


def total_penalty(
    groups: list[list[int]],
    labels: np.ndarray,
    feature_df: pd.DataFrame,
    schema: DatasetSchema,
    G: int = GROUP_SIZE,
) -> float:
    terms = normalized_penalty_terms(groups, labels, feature_df, schema, G)
    return (
        terms["size_penalty"]
        + terms["complementarity_penalty"]
        + 0.5 * terms["fairness_penalty"]
        + 0.5 * terms["engagement_penalty"]
    )


@dataclass
class _PenaltyContext:
    """Precomputed numpy state for fast per-trial penalty evaluation.

    Used inside :func:`greedy_refine_groups` so each candidate swap recomputes
    only the affected groups instead of looping pandas value_counts over the
    full cohort.
    """

    G: int
    fairness_codes: list[np.ndarray]   # per fairness col, int-coded categories
    fairness_n_cats: list[int]          # number of categories for each col
    fairness_global: list[np.ndarray]   # per fairness col, global class distribution
    engagement: np.ndarray | None       # NaNs replaced by global mean
    engagement_mu: float
    engagement_sd: float

    @classmethod
    def build(
        cls, feature_df: pd.DataFrame, schema: DatasetSchema, G: int
    ) -> "_PenaltyContext":
        fairness_codes: list[np.ndarray] = []
        fairness_n_cats: list[int] = []
        fairness_global: list[np.ndarray] = []
        for col in schema.fairness_cols:
            if col not in feature_df.columns:
                continue
            codes, _ = pd.factorize(feature_df[col], sort=False, use_na_sentinel=False)
            n_cats = int(codes.max()) + 1 if codes.size else 0
            fairness_codes.append(np.asarray(codes, dtype=np.int64))
            fairness_n_cats.append(n_cats)
            counts = np.bincount(codes, minlength=n_cats).astype(float)
            total = counts.sum()
            fairness_global.append(counts / total if total > 0 else counts)

        if schema.engagement_col and schema.engagement_col in feature_df.columns:
            engage = feature_df[schema.engagement_col].to_numpy(dtype=float)
            mu = float(np.nanmean(engage)) if engage.size else 0.0
            sd = float(np.nanstd(engage)) if engage.size else 0.0
            engage_filled = np.where(np.isnan(engage), mu, engage)
        else:
            engage_filled = None
            mu = 0.0
            sd = 0.0

        return cls(
            G=G,
            fairness_codes=fairness_codes,
            fairness_n_cats=fairness_n_cats,
            fairness_global=fairness_global,
            engagement=engage_filled,
            engagement_mu=mu,
            engagement_sd=sd,
        )

    def fairness_for_group(self, idx_col: int, group: list[int]) -> float:
        """TV distance between this group's distribution and the global one."""
        if not group:
            return 0.0
        codes = self.fairness_codes[idx_col]
        n_cats = self.fairness_n_cats[idx_col]
        if n_cats == 0:
            return 0.0
        counts = np.bincount(codes[group], minlength=n_cats).astype(float)
        group_dist = counts / len(group)
        return 0.5 * float(np.abs(group_dist - self.fairness_global[idx_col]).sum())

    def engagement_for_group(self, group: list[int]) -> float:
        """Absolute z-deviation of this group's engagement mean from global."""
        if self.engagement is None or not group or self.engagement_sd == 0:
            return 0.0
        return abs(float(self.engagement[group].mean()) - self.engagement_mu) / self.engagement_sd


def _per_group_penalties(
    groups: list[list[int]], ctx: _PenaltyContext
) -> tuple[list[list[float]], list[float]]:
    """Return ``(fairness_per_group_per_col, engagement_per_group)``.

    Both are indexed by group position so swaps can update only the two
    affected entries.
    """
    n_fair = len(ctx.fairness_codes)
    fair = [
        [ctx.fairness_for_group(c, group) for group in groups]
        for c in range(n_fair)
    ]
    engage = [ctx.engagement_for_group(group) for group in groups]
    return fair, engage


def _aggregate(
    fair_per_group: list[list[float]],
    engage_per_group: list[float],
    size_pen: float,
    comp_pen: float,
) -> float:
    if fair_per_group:
        fairness = float(np.mean([np.mean(col) for col in fair_per_group]))
    else:
        fairness = 0.0
    engagement = (
        min(float(np.mean(engage_per_group)) / 3.0, 1.0)
        if engage_per_group else 0.0
    )
    fairness = float(np.clip(fairness, 0.0, 1.0))
    engagement = float(np.clip(engagement, 0.0, 1.0))
    return size_pen + comp_pen + 0.5 * fairness + 0.5 * engagement


def greedy_refine_groups(
    groups: list[list[int]],
    labels: np.ndarray,
    feature_df: pd.DataFrame,
    schema: DatasetSchema,
    G: int = GROUP_SIZE,
    max_iters: int = MAX_SWAP_ITERS,
    refinement_time_budget: float | None = None,
) -> list[list[int]]:
    """Refine groups by greedy pairwise swaps.

    Uses a precomputed numpy penalty context so each trial swap only recomputes
    fairness/engagement for the two affected groups. Worst-case runtime stays
    O(iters x g^2 x G^2), but each trial is now O(F) numpy ops instead of
    O(F x g) pandas calls, giving ~50-200x speed-up on cohorts of a few
    hundred learners. For cohorts above roughly 2,000 learners, pass an
    explicit ``max_iters`` and/or ``refinement_time_budget`` because exhaustive
    pairwise swap search is still quadratic in group count.
    """
    groups = [list(group) for group in groups if group]
    if not groups:
        return groups

    labels = np.asarray(labels)
    ctx = _PenaltyContext.build(feature_df, schema, G)
    n_groups = len(groups)
    n_clusters = len(set(labels.tolist()) - {-1})

    # Precompute per-group fairness + engagement at the current state.
    fair_pg, engage_pg = _per_group_penalties(groups, ctx)
    n_groups_f = max(1, n_groups)
    deadline = (
        time.monotonic() + refinement_time_budget
        if refinement_time_budget is not None and refinement_time_budget > 0
        else None
    )
    plateau_count = 0

    for _ in range(max_iters):
        if deadline is not None and time.monotonic() >= deadline:
            break
        base_size_v = size_violations(groups, G)
        base_comp_v = cluster_complementarity_violations(groups, labels)
        size_pen = float(np.clip(base_size_v / n_groups_f, 0.0, 1.0))
        comp_pen = float(np.clip(base_comp_v / n_groups_f, 0.0, 1.0))
        base_penalty = _aggregate(fair_pg, engage_pg, size_pen, comp_pen)
        best_penalty = base_penalty
        best_swap: tuple[int, int, int, int] | None = None
        best_updates: tuple[list[float], list[float], float, float] | None = None

        budget_exhausted = False
        for gi in range(n_groups):
            if deadline is not None and time.monotonic() >= deadline:
                budget_exhausted = True
                break
            group_i = groups[gi]
            for gj in range(gi + 1, n_groups):
                group_j = groups[gj]
                for pi in range(len(group_i)):
                    a = group_i[pi]
                    for pj in range(len(group_j)):
                        b = group_j[pj]
                        # Construct trial groups for the two affected groups only.
                        trial_i = list(group_i)
                        trial_j = list(group_j)
                        trial_i[pi] = b
                        trial_j[pj] = a

                        # Hard constraints: complementarity for the two groups.
                        new_uniq_i = len(set(labels[trial_i].tolist()) - {-1})
                        new_uniq_j = len(set(labels[trial_j].tolist()) - {-1})
                        old_uniq_i = len(set(labels[group_i].tolist()) - {-1})
                        old_uniq_j = len(set(labels[group_j].tolist()) - {-1})
                        if _would_reduce_max_diversity(old_uniq_i, new_uniq_i, group_i, n_clusters, G):
                            continue
                        if _would_reduce_max_diversity(old_uniq_j, new_uniq_j, group_j, n_clusters, G):
                            continue
                        if (new_uniq_i < 2 and len(trial_i) > 1) and not (old_uniq_i < 2):
                            continue
                        if (new_uniq_j < 2 and len(trial_j) > 1) and not (old_uniq_j < 2):
                            continue
                        # Sizes are preserved by a same-cardinality swap, so size
                        # penalty is unchanged.

                        # Recompute only the two affected entries.
                        new_fair_i = [
                            ctx.fairness_for_group(c, trial_i)
                            for c in range(len(ctx.fairness_codes))
                        ]
                        new_fair_j = [
                            ctx.fairness_for_group(c, trial_j)
                            for c in range(len(ctx.fairness_codes))
                        ]
                        new_engage_i = ctx.engagement_for_group(trial_i)
                        new_engage_j = ctx.engagement_for_group(trial_j)

                        # Build trial per-group arrays cheaply.
                        trial_fair_pg = [col.copy() for col in fair_pg]
                        for c in range(len(trial_fair_pg)):
                            trial_fair_pg[c][gi] = new_fair_i[c]
                            trial_fair_pg[c][gj] = new_fair_j[c]
                        trial_engage_pg = list(engage_pg)
                        trial_engage_pg[gi] = new_engage_i
                        trial_engage_pg[gj] = new_engage_j

                        penalty = _aggregate(
                            trial_fair_pg, trial_engage_pg, size_pen, comp_pen
                        )
                        if penalty + 1e-12 < best_penalty:
                            best_penalty = penalty
                            best_swap = (gi, pi, gj, pj)
                            best_updates = (
                                [new_fair_i[c] for c in range(len(ctx.fairness_codes))],
                                [new_fair_j[c] for c in range(len(ctx.fairness_codes))],
                                new_engage_i,
                                new_engage_j,
                            )

        if budget_exhausted:
            break
        if best_swap is None or best_updates is None:
            break
        if base_penalty > 0 and best_penalty / base_penalty > 0.995:
            plateau_count += 1
        else:
            plateau_count = 0
        gi, pi, gj, pj = best_swap
        groups[gi][pi], groups[gj][pj] = groups[gj][pj], groups[gi][pi]
        new_fair_i, new_fair_j, new_engage_i, new_engage_j = best_updates
        for c in range(len(fair_pg)):
            fair_pg[c][gi] = new_fair_i[c]
            fair_pg[c][gj] = new_fair_j[c]
        engage_pg[gi] = new_engage_i
        engage_pg[gj] = new_engage_j
        if plateau_count >= 3:
            break

    return groups


def refine_demo(
    groups: list[list[int]],
    labels: np.ndarray,
    G: int = GROUP_SIZE,
    feature_df: pd.DataFrame | None = None,
    schema: DatasetSchema | None = None,
    max_swap_iters: int = MAX_SWAP_ITERS,
    refinement_time_budget: float | None = None,
    enforce_complementarity: bool = True,
    soft_refine: bool = True,
) -> tuple[list[list[int]], dict[str, int | float | str]]:
    cleaned = [list(map(int, group)) for group in groups if group]
    refined = (
        greedy_swap_complementarity(cleaned, labels, max_iters=max_swap_iters)
        if enforce_complementarity
        else cleaned
    )
    if soft_refine and feature_df is not None and schema is not None:
        refined = greedy_refine_groups(
            refined,
            labels,
            feature_df.reset_index(drop=True),
            schema,
            G,
            max_iters=max_swap_iters,
            refinement_time_budget=refinement_time_budget,
        )
    summary = summarize_demo_constraints(refined, labels, G, feature_df, schema)
    summary["complementarity_refinement"] = (
        "hard_constraint" if enforce_complementarity else "not_applied_mode_objective"
    )
    if not soft_refine and feature_df is not None and schema is not None:
        summary["fairness_refinement"] = "evaluated_only"
        summary["engagement_refinement"] = "evaluated_only"
    return refined, summary