"""Stratified, family-aware subset sampling for the ensemble PGD step.

Plain random.sample over the encoder pool lets a single step be dominated by
near-clone backbones (e.g. several OpenAI-style CLIP ViTs), which biases the
gradient toward one architecture and hurts transfer. This sampler:

  - caps how many encoders of the same architectural `family` appear per step,
  - guarantees at least one `feature` tower (raw-patch VLM backbones) is present
    when any exist, so text-less towers actually shape the perturbation.

Falls back gracefully when the pool is small or a constraint can't be met.
"""

from __future__ import annotations

import random


def stratified_sample(encoders: list, k: int, rng: random.Random,
                      max_per_family: int = 2, min_feature: int = 1) -> list:
    """Pick <=k encoders with per-family caps and a feature-tower floor.

    encoders: list of Encoder (each has .family and .kind).
    """
    if len(encoders) <= k:
        return list(encoders)

    feature = [e for e in encoders if e.kind == "feature"]
    chosen: list = []
    fam_count: dict[str, int] = {}

    def try_add(e) -> bool:
        fam = e.family or e.name
        if fam_count.get(fam, 0) >= max_per_family:
            return False
        chosen.append(e)
        fam_count[fam] = fam_count.get(fam, 0) + 1
        return True

    # 1) seed the required feature towers first (respecting the family cap)
    want_feature = min(min_feature, len(feature))
    for e in rng.sample(feature, len(feature)):
        if len([c for c in chosen if c.kind == "feature"]) >= want_feature:
            break
        try_add(e)

    # 2) fill the rest from a shuffled pool, honoring family caps
    pool = [e for e in encoders if e not in chosen]
    for e in rng.sample(pool, len(pool)):
        if len(chosen) >= k:
            break
        try_add(e)

    # 3) if family caps left us short of k, relax the cap to fill remaining slots
    if len(chosen) < k:
        remaining = [e for e in encoders if e not in chosen]
        for e in rng.sample(remaining, len(remaining)):
            if len(chosen) >= k:
                break
            chosen.append(e)

    return chosen