module4_select.py

import atexit
import math
import os
import pickle
from collections import Counter
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
from drfp import DrfpEncoder
from sklearn.decomposition import PCA
from rdkit import Chem
from rdkit.SimDivFilters.rdSimDivPickers import MaxMinPicker

_CACHE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), ".drfp_cache.pkl"
)
_cache: Dict[str, np.ndarray] = {}
_cache_loaded = False
_cache_dirty = False


def _load_cache() -> None:
    global _cache, _cache_loaded
    if _cache_loaded:
        return
    _cache_loaded = True
    if os.path.exists(_CACHE_PATH):
        try:
            with open(_CACHE_PATH, "rb") as f:
                _cache = pickle.load(f)
        except Exception as e:
            print(f"  DRFP cache load failed ({e}); starting empty.")
            _cache = {}


def _save_cache() -> None:
    global _cache_dirty
    if not _cache_dirty:
        return
    try:
        tmp = _CACHE_PATH + ".tmp"
        with open(tmp, "wb") as f:
            pickle.dump(_cache, f, protocol=pickle.HIGHEST_PROTOCOL)
        os.replace(tmp, _CACHE_PATH)
        _cache_dirty = False
    except Exception as e:
        print(f"  DRFP cache save failed: {e}")


atexit.register(_save_cache)


def encode_with_cache(smiles_list: List[str]) -> np.ndarray:
    """
    Return fingerprints aligned with smiles_list, using the module-level cache
    for known entries and batch-encoding only cache misses via DrfpEncoder.
    The cache is updated in place; persistence happens at process exit.
    """
    global _cache_dirty
    _load_cache()

    n = len(smiles_list)
    result: List[np.ndarray] = [None] * n  # type: ignore[list-item]
    miss_indices: List[int] = []
    miss_smiles: List[str] = []

    for i, s in enumerate(smiles_list):
        hit = _cache.get(s)
        if hit is not None:
            result[i] = hit
        else:
            miss_indices.append(i)
            miss_smiles.append(s)

    if miss_smiles:
        encoded = DrfpEncoder.encode(miss_smiles, mapping=False)
        for idx, smi, fp in zip(miss_indices, miss_smiles, encoded):
            arr = np.asarray(fp)
            _cache[smi] = arr
            result[idx] = arr
        _cache_dirty = True

    return np.stack(result)


def select_diverse_reactions(
    current_dataset_smiles: List[str],
    candidate_smiles: List[str],
    target_count: int,
    pca_dims: int = 100,
) -> List[int]:
    """
    Selects target_count candidates that are most diverse with respect to current_dataset.
    Returns indices of selected candidates.
    """
    if not candidate_smiles:
        return []

    all_smiles = current_dataset_smiles + candidate_smiles
    fps_array = encode_with_cache(all_smiles)

    # Optional PCA
    if pca_dims < fps_array.shape[1]:
        n_comp = min(pca_dims, fps_array.shape[0])
        if n_comp > 1:
            pca = PCA(n_components=n_comp)
            fps_array = pca.fit_transform(fps_array)

    num_current = len(current_dataset_smiles)

    # Distance function between any two indices in the combined array
    def dist_ij(i, j, data=fps_array):
        return float(np.linalg.norm(data[i] - data[j]))

    picker = MaxMinPicker()

    # We want to pick target_count more points starting from the current_dataset.
    # The picker.LazyPick takes indices [0...N-1].
    # But it also takes firstPicks: indices already in the dataset.

    first_picks = list(range(num_current))

    # If target_count + len(first_picks) > total_points, cap it
    total_needed = min(num_current + target_count, len(all_smiles))

    picks = picker.LazyPick(
        dist_ij, len(all_smiles), total_needed, firstPicks=first_picks
    )

    # The results include first_picks, so we filter them out to get only new picks
    new_picks = [p - num_current for p in picks if p >= num_current]

    return new_picks[:target_count]


def _family_key(item: Dict) -> str:
    family = item.get("family")
    return family if family else "unknown"


SIZE_BIN_FRACTIONS = {
    "1-3": 0.25,
    "4-5": 0.30,
    "6-8": 0.30,
    "9-12": 0.15,
}
MIN_HYDROCARBON_FRACTION = 0.10
MIN_CHO_ONLY_FRACTION = 0.35
MAX_N_FRACTION = 0.25
MAX_S_FRACTION = 0.20

# Halogens allowed by the benchmark are restricted to Cl, but match the full
# set so a stray F/Br/I never silently slips past the "Others" accounting.
HALOGENS = frozenset({"F", "Cl", "Br", "I"})

# Per-species element marginal targets for the flat balanced selector. These are
# overlapping fractions over the *species* appearing in the selected reactions:
# a C/O species counts toward both C and O. C and O are treated as floors,
# N / S / halogen as caps (they were previously ~47% / ~48% of reactions).
# C=0.80 -> ~20% carbon-free species.
ELEMENT_MARGINAL_TARGETS = {
    "C": 0.80,
    "O": 0.30,
    "N": 0.15,
    "S": 0.05,
    "halogen": 0.05,
}

# Soft size preference applied to whole reactions during the flat selection:
# weight = exp(-SIZE_WEIGHT_LAMBDA * total_reaction_heavy_atoms). lambda=0.20
# gives a target median total of ~3.5 heavy atoms per reaction.
SIZE_WEIGHT_LAMBDA = 0.20


def reaction_selection_descriptors(item: Dict) -> Dict[str, object]:
    """Return topology, size, and elemental strata for a dataset entry."""
    smiles = item["smiles"]
    left, right = smiles.split(">>", 1)
    reactants = left.split(".") if left else []
    products = right.split(".") if right else []
    heavy_counts: List[int] = []
    side_heavy_counts: List[int] = []
    elements = set()
    num_species = 0
    # Per-species element presence, used for the per-species marginal targets
    # (90% C, 30% O, 15% N, 5% S, 5% halogen). A reaction's union "has_*" flag
    # would over-count, since a 3-species reaction trips "has_N" if any one
    # species contains N.
    species_has = {"C": 0, "O": 0, "N": 0, "S": 0, "halogen": 0}
    for side in (reactants, products):
        side_total = 0
        for molecule_smiles in side:
            molecule = Chem.MolFromSmiles(molecule_smiles)
            if molecule is None:
                continue
            num_species += 1
            heavy = molecule.GetNumHeavyAtoms()
            heavy_counts.append(heavy)
            side_total += heavy
            species_elements = {
                atom.GetSymbol() for atom in molecule.GetAtoms()
                if atom.GetSymbol() != "H"
            }
            elements.update(species_elements)
            if "C" in species_elements:
                species_has["C"] += 1
            if "O" in species_elements:
                species_has["O"] += 1
            if "N" in species_elements:
                species_has["N"] += 1
            if "S" in species_elements:
                species_has["S"] += 1
            if species_elements & HALOGENS:
                species_has["halogen"] += 1
        side_heavy_counts.append(side_total)
    maximum = max(heavy_counts, default=0)
    if maximum <= 3:
        size_bin = "1-3"
    elif maximum <= 5:
        size_bin = "4-5"
    elif maximum <= 8:
        size_bin = "6-8"
    else:
        size_bin = "9-12"
    hydrocarbon = bool(elements) and elements <= {"C"}
    cho_only = bool(elements) and elements <= {"C", "O"}
    return {
        "arity": f"{len(reactants)}->{len(products)}",
        "size_bin": size_bin,
        "elements": elements,
        "hydrocarbon": hydrocarbon,
        "cho_only": cho_only,
        "has_c": "C" in elements,
        "has_o": "O" in elements,
        "has_n": "N" in elements,
        "has_s": "S" in elements,
        "has_halogen": bool(elements & HALOGENS),
        "num_species": num_species,
        "species_has": species_has,
        "total_heavy_atoms": sum(heavy_counts),
        "max_side_heavy_atoms": max(side_heavy_counts, default=0),
        "max_species_heavy_atoms": maximum,
    }


def _prepare_embedding(smiles_list: List[str], pca_dims: int) -> np.ndarray:
    fps_array = encode_with_cache(smiles_list)
    if pca_dims < fps_array.shape[1]:
        n_comp = min(pca_dims, fps_array.shape[0])
        if n_comp > 1:
            pca = PCA(n_components=n_comp)
            fps_array = pca.fit_transform(fps_array)
    return np.asarray(fps_array, dtype=float)


def _initial_min_distances(
    data: np.ndarray, selected_global: Sequence[int]
) -> np.ndarray:
    """
    Distance from each point to its nearest selected point. If there are no
    selected points, use distance from the global centroid as a deterministic
    cold-start proxy for chemical spread.
    """
    if selected_global:
        selected = data[list(selected_global)]
        dists = np.linalg.norm(data[:, None, :] - selected[None, :, :], axis=2)
        return np.min(dists, axis=1)

    centroid = np.mean(data, axis=0)
    return np.linalg.norm(data - centroid, axis=1)


def select_family_aware_diverse_reactions(
    current_dataset: List[Dict],
    candidate_pool: List[Dict],
    target_count: int,
    *,
    min_per_family: int = 3,
    max_family_fraction: float = 0.15,
    family_penalty_alpha: float = 0.75,
    pca_dims: int = 100,
    audit_output: Optional[Dict] = None,
) -> List[int]:
    """
    Select candidates by balancing reaction-family representation with chemical
    diversity.

    Policy:
      1. Treat current_dataset as already selected.
      2. Repair underrepresented families up to min_per_family when candidates
         exist.
      3. Fill remaining slots greedily by nearest-selected DRFP/PCA distance,
         multiplied by a dynamic penalty for already-common families.
      4. Avoid adding more candidates from a family once it reaches the hard
         max_family_fraction cap, unless no uncapped candidates remain.

    Returns indices into candidate_pool.
    """
    if target_count <= 0 or not candidate_pool:
        return []

    current_smiles = [d["smiles"] for d in current_dataset]
    candidate_smiles = [c["smiles"] for c in candidate_pool]
    all_smiles = current_smiles + candidate_smiles
    data = _prepare_embedding(all_smiles, pca_dims=pca_dims)

    num_current = len(current_dataset)
    selected_global: List[int] = list(range(num_current))
    selected_candidate_indices: List[int] = []
    selected_candidate_set = set()

    family_counts = Counter(_family_key(d) for d in current_dataset)
    candidate_families = [_family_key(c) for c in candidate_pool]
    families_with_candidates = sorted(set(candidate_families))

    final_total = len(current_dataset) + min(target_count, len(candidate_pool))
    hard_cap = max(
        min_per_family, int(math.ceil(max_family_fraction * max(1, final_total)))
    )

    min_dist = _initial_min_distances(data, selected_global)
    stage_a_picks = 0
    stage_b_picks = 0
    relaxed_cap_picks = 0

    def candidate_global_idx(candidate_idx: int) -> int:
        return num_current + candidate_idx

    def add_candidate(candidate_idx: int) -> None:
        selected_candidate_set.add(candidate_idx)
        selected_candidate_indices.append(candidate_idx)
        global_idx = candidate_global_idx(candidate_idx)
        selected_global.append(global_idx)
        family_counts[candidate_families[candidate_idx]] += 1

        # Incrementally update nearest-selected distances.
        new_dist = np.linalg.norm(data - data[global_idx], axis=1)
        np.minimum(min_dist, new_dist, out=min_dist)

    def best_from_family(family: str) -> Optional[int]:
        best_idx = None
        best_score = -1.0
        for idx, fam in enumerate(candidate_families):
            if fam != family or idx in selected_candidate_set:
                continue
            score = float(min_dist[candidate_global_idx(idx)])
            if score > best_score:
                best_score = score
                best_idx = idx
        return best_idx

    # Stage A: guarantee minimum family coverage where candidates exist.
    made_progress = True
    while len(selected_candidate_indices) < target_count and made_progress:
        made_progress = False
        underrepresented = [
            fam
            for fam in families_with_candidates
            if family_counts[fam] < min_per_family
        ]
        if not underrepresented:
            break
        underrepresented.sort(key=lambda fam: (family_counts[fam], fam))
        for fam in underrepresented:
            if len(selected_candidate_indices) >= target_count:
                break
            pick = best_from_family(fam)
            if pick is None:
                continue
            add_candidate(pick)
            stage_a_picks += 1
            made_progress = True

    # Stage B: fill remaining slots with chemical diversity, family penalty,
    # and a hard cap. If all remaining candidates are capped, relax the cap so
    # the dataset can still reach the requested size.
    while len(selected_candidate_indices) < target_count:
        best_idx = None
        best_score = -1.0
        best_relaxed_idx = None
        best_relaxed_score = -1.0

        for idx, fam in enumerate(candidate_families):
            if idx in selected_candidate_set:
                continue

            raw_distance = float(min_dist[candidate_global_idx(idx)])
            penalty = 1.0 / math.sqrt(1.0 + family_penalty_alpha * family_counts[fam])
            if family_counts[fam] < min_per_family:
                penalty *= 2.0
            score = raw_distance * penalty

            if score > best_relaxed_score:
                best_relaxed_score = score
                best_relaxed_idx = idx

            if family_counts[fam] >= hard_cap:
                continue
            if score > best_score:
                best_score = score
                best_idx = idx

        pick = best_idx if best_idx is not None else best_relaxed_idx
        if pick is None:
            break
        if best_idx is None:
            relaxed_cap_picks += 1
        add_candidate(pick)
        stage_b_picks += 1

    if audit_output is not None:
        selected_families = [candidate_families[i] for i in selected_candidate_indices]
        audit_output.update(
            {
                "target_count": target_count,
                "selected_count": len(selected_candidate_indices),
                "min_per_family": min_per_family,
                "max_family_fraction": max_family_fraction,
                "family_penalty_alpha": family_penalty_alpha,
                "hard_family_cap": hard_cap,
                "stage_a_min_coverage_picks": stage_a_picks,
                "stage_b_weighted_diversity_picks": stage_b_picks,
                "relaxed_cap_picks": relaxed_cap_picks,
                "candidate_family_counts": dict(Counter(candidate_families)),
                "selected_candidate_family_counts": dict(Counter(selected_families)),
                "final_family_counts_after_selection": dict(family_counts),
            }
        )

    return selected_candidate_indices[:target_count]


def _fraction_targets(total: int, fractions: Dict[str, float]) -> Dict[str, int]:
    raw = {key: total * fraction for key, fraction in fractions.items()}
    targets = {key: int(math.floor(value)) for key, value in raw.items()}
    remainder = total - sum(targets.values())
    order = sorted(raw, key=lambda key: (raw[key] - targets[key], key), reverse=True)
    for key in order[:remainder]:
        targets[key] += 1
    return targets


def select_stratified_reactions(
    current_dataset: List[Dict],
    candidate_pool: List[Dict],
    target_count: int,
    *,
    mandatory_families: Sequence[str],
    min_per_family: int = 1,
    max_family_fraction: float = 0.15,
    family_penalty_alpha: float = 0.75,
    pca_dims: int = 100,
    audit_output: Optional[Dict] = None,
) -> List[int]:
    """
    Select reactions with family, topology, size, and elemental controls.

    DRFP distances are updated only between reactions with the same arity.
    This prevents 2->2 reactions from appearing intrinsically more diverse
    than unimolecular reactions merely because they contain more fragments.
    """
    if target_count <= 0 or not candidate_pool:
        return []

    current_smiles = [item["smiles"] for item in current_dataset]
    candidate_smiles = [item["smiles"] for item in candidate_pool]
    data = _prepare_embedding(current_smiles + candidate_smiles, pca_dims)
    num_current = len(current_dataset)

    current_desc = [reaction_selection_descriptors(item) for item in current_dataset]
    candidate_desc = [reaction_selection_descriptors(item) for item in candidate_pool]
    candidate_families = [_family_key(item) for item in candidate_pool]
    families_with_candidates = set(candidate_families)
    required = sorted(set(mandatory_families) & families_with_candidates)

    final_total = len(current_dataset) + min(target_count, len(candidate_pool))
    size_targets = _fraction_targets(final_total, SIZE_BIN_FRACTIONS)
    hydrocarbon_target = int(math.ceil(MIN_HYDROCARBON_FRACTION * final_total))
    cho_only_target = int(math.ceil(MIN_CHO_ONLY_FRACTION * final_total))
    n_cap = int(math.floor(MAX_N_FRACTION * final_total))
    s_cap = int(math.floor(MAX_S_FRACTION * final_total))
    family_cap = max(
        min_per_family, int(math.ceil(max_family_fraction * max(1, final_total)))
    )

    family_counts = Counter(_family_key(item) for item in current_dataset)
    size_counts = Counter(desc["size_bin"] for desc in current_desc)
    hydrocarbon_count = sum(bool(desc["hydrocarbon"]) for desc in current_desc)
    cho_only_count = sum(bool(desc["cho_only"]) for desc in current_desc)
    n_count = sum(bool(desc["has_n"]) for desc in current_desc)
    s_count = sum(bool(desc["has_s"]) for desc in current_desc)

    selected: List[int] = []
    selected_set = set()
    selected_global_by_arity: Dict[str, List[int]] = {}
    for index, desc in enumerate(current_desc):
        selected_global_by_arity.setdefault(str(desc["arity"]), []).append(index)

    min_dist = np.zeros(len(candidate_pool), dtype=float)
    for arity in sorted({str(desc["arity"]) for desc in candidate_desc}):
        indices = [
            index for index, desc in enumerate(candidate_desc)
            if desc["arity"] == arity
        ]
        global_indices = [num_current + index for index in indices]
        selected_same_arity = selected_global_by_arity.get(arity, [])
        if selected_same_arity:
            selected_data = data[selected_same_arity]
            distances = np.linalg.norm(
                data[global_indices, None, :] - selected_data[None, :, :],
                axis=2,
            )
            min_dist[indices] = np.min(distances, axis=1)
        else:
            centroid = np.mean(data[global_indices], axis=0)
            min_dist[indices] = np.linalg.norm(data[global_indices] - centroid, axis=1)

    positive = min_dist[min_dist > 0]
    distance_scale = float(np.median(positive)) if len(positive) else 1.0
    stage_counts = Counter()
    relaxed_element_cap_picks = 0
    relaxed_size_cap_picks = 0

    def add(index: int, stage: str) -> None:
        nonlocal hydrocarbon_count, cho_only_count, n_count, s_count
        selected.append(index)
        selected_set.add(index)
        family_counts[candidate_families[index]] += 1
        desc = candidate_desc[index]
        size_counts[str(desc["size_bin"])] += 1
        hydrocarbon_count += int(bool(desc["hydrocarbon"]))
        cho_only_count += int(bool(desc["cho_only"]))
        n_count += int(bool(desc["has_n"]))
        s_count += int(bool(desc["has_s"]))
        stage_counts[stage] += 1

        arity = str(desc["arity"])
        global_index = num_current + index
        selected_global_by_arity.setdefault(arity, []).append(global_index)
        same_arity = [
            i for i, other in enumerate(candidate_desc)
            if other["arity"] == arity and i not in selected_set
        ]
        if same_arity:
            new_distance = np.linalg.norm(
                data[[num_current + i for i in same_arity]] - data[global_index],
                axis=1,
            )
            min_dist[same_arity] = np.minimum(min_dist[same_arity], new_distance)

    def allowed_by_element_caps(index: int) -> bool:
        desc = candidate_desc[index]
        return not (
            (bool(desc["has_n"]) and n_count >= n_cap)
            or (bool(desc["has_s"]) and s_count >= s_cap)
        )

    def allowed_by_size_cap(index: int) -> bool:
        size_bin = str(candidate_desc[index]["size_bin"])
        return size_counts[size_bin] < size_targets[size_bin]

    def score(index: int) -> float:
        family = candidate_families[index]
        desc = candidate_desc[index]
        distance_score = float(min_dist[index]) / max(distance_scale, 1e-12)
        family_weight = 1.0 / math.sqrt(
            1.0 + family_penalty_alpha * family_counts[family]
        )
        complexity_penalty = 0.03 * float(desc["max_side_heavy_atoms"])
        return distance_score * family_weight - complexity_penalty

    def choose(
        eligible,
        *,
        enforce_element_caps: bool = True,
        enforce_size_caps: bool = False,
    ) -> Optional[int]:
        choices = [
            index for index in eligible
            if index not in selected_set
            and family_counts[candidate_families[index]] < family_cap
            and (not enforce_element_caps or allowed_by_element_caps(index))
            and (not enforce_size_caps or allowed_by_size_cap(index))
        ]
        if not choices:
            return None
        return max(
            choices,
            key=lambda index: (
                score(index),
                -float(candidate_desc[index]["total_heavy_atoms"]),
                candidate_pool[index]["smiles"],
            ),
        )

    # Stage A: every mandatory family, then the requested minimum per family.
    family_order = required + sorted(families_with_candidates - set(required))
    for family in family_order:
        while (
            len(selected) < target_count
            and family_counts[family] < min_per_family
        ):
            family_indices = [
                index for index, candidate_family in enumerate(candidate_families)
                if candidate_family == family
            ]
            pick = choose(
                family_indices,
                enforce_element_caps=False,
                enforce_size_caps=False,
            )
            if pick is None:
                break
            add(pick, "family_coverage")

    # Stage B: repair maximum-species-size distribution.
    for size_bin in SIZE_BIN_FRACTIONS:
        while (
            len(selected) < target_count
            and size_counts[size_bin] < size_targets[size_bin]
        ):
            eligible = [
                index for index, desc in enumerate(candidate_desc)
                if desc["size_bin"] == size_bin
            ]
            pick = choose(eligible, enforce_element_caps=True)
            if pick is None:
                pick = choose(eligible, enforce_element_caps=False)
                if pick is None:
                    break
                relaxed_element_cap_picks += 1
            add(pick, "size_balance")

    # Stage C: guarantee simple hydrocarbon and C/H/O-only chemistry.
    while len(selected) < target_count and hydrocarbon_count < hydrocarbon_target:
        eligible = [
            index for index, desc in enumerate(candidate_desc)
            if bool(desc["hydrocarbon"])
        ]
        pick = choose(eligible, enforce_size_caps=True)
        if pick is None:
            pick = choose(eligible)
        if pick is None:
            break
        add(pick, "hydrocarbon_minimum")

    while len(selected) < target_count and cho_only_count < cho_only_target:
        eligible = [
            index for index, desc in enumerate(candidate_desc)
            if bool(desc["cho_only"])
        ]
        pick = choose(eligible, enforce_size_caps=True)
        if pick is None:
            pick = choose(eligible)
        if pick is None:
            break
        add(pick, "cho_only_minimum")

    # Stage D: fill while respecting size and elemental caps. Relax size first,
    # then elemental caps only if needed to reach the requested dataset size.
    all_indices = range(len(candidate_pool))
    while len(selected) < target_count:
        pick = choose(
            all_indices,
            enforce_element_caps=True,
            enforce_size_caps=True,
        )
        if pick is None:
            pick = choose(all_indices, enforce_element_caps=True)
            if pick is not None:
                relaxed_size_cap_picks += 1
        if pick is None:
            pick = choose(all_indices, enforce_element_caps=False)
            if pick is not None:
                relaxed_element_cap_picks += 1
        if pick is None:
            break
        add(pick, "stratified_fill")

    if audit_output is not None:
        selected_desc = [candidate_desc[index] for index in selected]
        audit_output.update(
            {
                "type": "family_size_element_arity_stratified",
                "target_count": target_count,
                "selected_count": len(selected),
                "mandatory_families": list(mandatory_families),
                "stage_counts": dict(stage_counts),
                "size_bin_targets": size_targets,
                "final_size_bin_counts": dict(size_counts),
                "hydrocarbon_target": hydrocarbon_target,
                "final_hydrocarbon_count": hydrocarbon_count,
                "cho_only_target": cho_only_target,
                "final_cho_only_count": cho_only_count,
                "n_cap": n_cap,
                "final_n_count": n_count,
                "s_cap": s_cap,
                "final_s_count": s_count,
                "family_cap": family_cap,
                "relaxed_size_cap_picks": relaxed_size_cap_picks,
                "relaxed_element_cap_picks": relaxed_element_cap_picks,
                "selected_arity_counts": dict(
                    Counter(str(desc["arity"]) for desc in selected_desc)
                ),
                "selected_family_counts": dict(
                    Counter(candidate_families[index] for index in selected)
                ),
            }
        )
    return selected[:target_count]


def _even_family_quotas(total: int, families: Sequence[str]) -> Dict[str, int]:
    """Distribute `total` slots across `families` as evenly as possible.

    With 21 families and total=500 this yields 17 families at 24 and 4 at 23.
    Order is deterministic (sorted by name) so the +1 remainder is reproducible.
    """
    ordered = sorted(families)
    n = len(ordered)
    if n == 0:
        return {}
    base, remainder = divmod(total, n)
    return {fam: base + (1 if i < remainder else 0) for i, fam in enumerate(ordered)}


def select_balanced_flat(
    current_dataset: List[Dict],
    candidate_pool: List[Dict],
    target_count: int,
    *,
    target_families: Sequence[str],
    size_lambda: float = SIZE_WEIGHT_LAMBDA,
    element_targets: Optional[Dict[str, float]] = None,
    family_penalty_alpha: float = 0.75,
    species_cap_headroom: float = 0.20,
    pca_dims: int = 100,
    audit_output: Optional[Dict] = None,
) -> Tuple[List[Dict], List[int]]:
    """
    Select a single flat dataset of `target_count` reactions (anchors included),
    quasi-evenly distributed across `target_families`, with:

      * an even per-family quota (target_count / n_families);
      * a soft exp(-size_lambda * total_heavy_atoms) preference for small
        reactions (counteracts the DRFP bias toward many-molecule reactions);
      * within-arity DRFP/PCA diversity, so 2->2 reactions are not treated as
        intrinsically more diverse than unimolecular ones merely for having
        more fragments;
      * per-*species* element caps on N / S / halogen and floors on C / O,
        matching the requested 90/30/15/5/5 marginals.

    Anchors (`current_dataset`) are treated as already selected and count toward
    both the total and the per-family quotas. Returns (output_entries,
    selected_candidate_indices).
    """
    if target_count <= 0:
        return list(current_dataset), []
    element_targets = element_targets or ELEMENT_MARGINAL_TARGETS

    num_anchors = len(current_dataset)
    need = min(target_count - num_anchors, len(candidate_pool))
    if need <= 0:
        return list(current_dataset)[:target_count], []

    anchor_desc = [reaction_selection_descriptors(item) for item in current_dataset]
    candidate_desc = [reaction_selection_descriptors(item) for item in candidate_pool]
    candidate_families = [_family_key(item) for item in candidate_pool]

    data = _prepare_embedding(
        [item["smiles"] for item in current_dataset]
        + [item["smiles"] for item in candidate_pool],
        pca_dims,
    )

    # --- within-arity nearest-selected distances --------------------------
    selected_global_by_arity: Dict[str, List[int]] = {}
    for index, desc in enumerate(anchor_desc):
        selected_global_by_arity.setdefault(str(desc["arity"]), []).append(index)

    min_dist = np.zeros(len(candidate_pool), dtype=float)
    for arity in sorted({str(desc["arity"]) for desc in candidate_desc}):
        indices = [
            index for index, desc in enumerate(candidate_desc)
            if str(desc["arity"]) == arity
        ]
        global_indices = [num_anchors + index for index in indices]
        same_arity_anchors = selected_global_by_arity.get(arity, [])
        if same_arity_anchors:
            anchor_data = data[same_arity_anchors]
            distances = np.linalg.norm(
                data[global_indices, None, :] - anchor_data[None, :, :], axis=2
            )
            min_dist[indices] = np.min(distances, axis=1)
        else:
            centroid = np.mean(data[global_indices], axis=0)
            min_dist[indices] = np.linalg.norm(data[global_indices] - centroid, axis=1)

    positive = min_dist[min_dist > 0]
    distance_scale = float(np.median(positive)) if len(positive) else 1.0

    # --- quotas and running tallies ---------------------------------------
    quotas = _even_family_quotas(target_count, target_families)
    family_counts = Counter(_family_key(item) for item in current_dataset)

    # Per-species element accounting against an estimate of the final total
    # species count (selection changes the denominator, so a fixed estimate is
    # more stable than a running ratio with a tiny early denominator).
    all_desc = anchor_desc + candidate_desc
    avg_species = (
        float(np.mean([d["num_species"] for d in all_desc])) if all_desc else 3.0
    ) or 3.0
    est_total_species = target_count * avg_species
    species_total = sum(d["num_species"] for d in anchor_desc)
    species_has = {key: 0 for key in element_targets}
    for desc in anchor_desc:
        for key in species_has:
            species_has[key] += desc["species_has"].get(key, 0)

    capped_elements = ("N", "S", "halogen")
    species_caps = {
        el: math.ceil(element_targets[el] * est_total_species * (1.0 + species_cap_headroom))
        for el in capped_elements
    }

    selected: List[int] = []
    selected_set: set = set()
    stage_counts: Counter = Counter()
    relaxed_family_cap_picks = 0
    relaxed_element_cap_picks = 0

    def size_weight(index: int) -> float:
        return math.exp(-size_lambda * float(candidate_desc[index]["total_heavy_atoms"]))

    def score(index: int) -> float:
        family = candidate_families[index]
        distance_score = float(min_dist[index]) / max(distance_scale, 1e-12)
        family_weight = 1.0 / math.sqrt(1.0 + family_penalty_alpha * family_counts[family])
        return distance_score * family_weight * size_weight(index)

    def allowed_by_element_caps(index: int) -> bool:
        sh = candidate_desc[index]["species_has"]
        return all(species_has[el] + sh.get(el, 0) <= species_caps[el] for el in capped_elements)

    def add(index: int, stage: str) -> None:
        selected.append(index)
        selected_set.add(index)
        desc = candidate_desc[index]
        family_counts[candidate_families[index]] += 1
        nonlocal species_total
        species_total += desc["num_species"]
        for key in species_has:
            species_has[key] += desc["species_has"].get(key, 0)
        stage_counts[stage] += 1
        # incremental within-arity distance update
        arity = str(desc["arity"])
        global_index = num_anchors + index
        same_arity = [
            i for i, other in enumerate(candidate_desc)
            if str(other["arity"]) == arity and i not in selected_set
        ]
        if same_arity:
            new_distance = np.linalg.norm(
                data[[num_anchors + i for i in same_arity]] - data[global_index], axis=1
            )
            min_dist[same_arity] = np.minimum(min_dist[same_arity], new_distance)

    def choose(eligible, *, enforce_family_cap=True, enforce_element_caps=True) -> Optional[int]:
        choices = [
            index for index in eligible
            if index not in selected_set
            and (not enforce_family_cap
                 or family_counts[candidate_families[index]] < quotas.get(candidate_families[index], 0))
            and (not enforce_element_caps or allowed_by_element_caps(index))
        ]
        if not choices:
            return None
        return max(
            choices,
            key=lambda index: (
                score(index),
                -float(candidate_desc[index]["total_heavy_atoms"]),
                candidate_pool[index]["smiles"],
            ),
        )

    candidate_by_family: Dict[str, List[int]] = {}
    for index, fam in enumerate(candidate_families):
        candidate_by_family.setdefault(fam, []).append(index)

    # Stage A: fill each family up to its quota (rarest-first), honouring caps.
    made_progress = True
    while len(selected) < need and made_progress:
        made_progress = False
        underfilled = sorted(
            (fam for fam in quotas if family_counts[fam] < quotas[fam]),
            key=lambda fam: (family_counts[fam], fam),
        )
        for fam in underfilled:
            if len(selected) >= need:
                break
            pick = choose(candidate_by_family.get(fam, []), enforce_family_cap=False)
            if pick is None:
                # relax element caps before giving up on this family's quota
                pick = choose(
                    candidate_by_family.get(fam, []),
                    enforce_family_cap=False,
                    enforce_element_caps=False,
                )
                if pick is None:
                    continue
                relaxed_element_cap_picks += 1
            add(pick, "family_quota")
            made_progress = True

    # Stage B: fill any remaining slots. Prefer staying within quota + element
    # caps; relax element caps, then family quotas, only if the pool forces it.
    all_indices = range(len(candidate_pool))
    while len(selected) < need:
        pick = choose(all_indices, enforce_family_cap=True, enforce_element_caps=True)
        if pick is None:
            pick = choose(all_indices, enforce_family_cap=True, enforce_element_caps=False)
            if pick is not None:
                relaxed_element_cap_picks += 1
        if pick is None:
            pick = choose(all_indices, enforce_family_cap=False, enforce_element_caps=True)
            if pick is not None:
                relaxed_family_cap_picks += 1
        if pick is None:
            pick = choose(all_indices, enforce_family_cap=False, enforce_element_caps=False)
            if pick is not None:
                relaxed_family_cap_picks += 1
        if pick is None:
            break
        add(pick, "diversity_fill")

    output = list(current_dataset) + [candidate_pool[index] for index in selected]

    if audit_output is not None:
        selected_desc = [candidate_desc[index] for index in selected]
        final_desc = anchor_desc + selected_desc
        final_species_total = sum(d["num_species"] for d in final_desc) or 1
        species_marginals = {
            el: sum(d["species_has"].get(el, 0) for d in final_desc) / final_species_total
            for el in element_targets
        }
        totals = [float(d["total_heavy_atoms"]) for d in final_desc]
        audit_output.update(
            {
                "type": "flat_family_balanced",
                "target_count": target_count,
                "selected_count": len(output),
                "anchor_count": num_anchors,
                "generated_selected_count": len(selected),
                "size_lambda": size_lambda,
                "family_quotas": quotas,
                "stage_counts": dict(stage_counts),
                "relaxed_family_cap_picks": relaxed_family_cap_picks,
                "relaxed_element_cap_picks": relaxed_element_cap_picks,
                "element_marginal_targets": dict(element_targets),
                "species_element_caps": species_caps,
                "final_species_element_marginals": species_marginals,
                "estimated_avg_species_per_reaction": avg_species,
                "final_family_counts": dict(
                    Counter(_family_key(item) for item in output)
                ),
                "final_arity_counts": dict(
                    Counter(str(d["arity"]) for d in final_desc)
                ),
                "total_heavy_atom_distribution": {
                    "mean": float(np.mean(totals)) if totals else 0.0,
                    "median": float(np.median(totals)) if totals else 0.0,
                    "p90": float(np.percentile(totals, 90)) if totals else 0.0,
                    "max": float(np.max(totals)) if totals else 0.0,
                },
            }
        )

    return output, selected