"""混合 mini-batch 切分：大类打散、小类整锁，供 runner 每 step 处理一个 batch。"""

from __future__ import annotations

import math
import random
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from core.types import GeneratedQuestion


def build_batches(
    items: list[GeneratedQuestion],
    correctness: dict[str, bool],
    batch_size: int,
    min_class_per_batch: int,
    seed: int,
    correct_ratio: float = 0.0,
) -> tuple[list[list[GeneratedQuestion]], int]:
    """把诊断池里的题目切成多个混合 mini-batch。

    当 ``correct_ratio > 0`` 时，按题型为每组错题配比一定数量的正确题，使 batch
    包含正误混合样本（"动量"机制）；``correct_ratio <= 0`` 时退化为纯错题模式。

    参数:
        items: 候选题目全集。
        correctness: question_id -> 基线是否答对。
        batch_size: 单个 batch 的样本数上限（> 0）。
        min_class_per_batch: 小类判定阈值——题目数 ≤ 此值的题型整组锁进单一
            batch（> 0）。
        seed: 随机种子，保证相同输入产出完全一致的切分。
        correct_ratio: 正确题占比（0.0 ~ 1.0）。0.0 = 纯错题；0.5 = 错题:正确题 = 1:1。
    返回:
        (非空 mini-batch 列表, selected_count)；无错题时返回 ([], 0)。
        selected_count 是所有 batch 中题目总数。
    异常:
        ValueError: batch_size 或 min_class_per_batch < 1, 或
            min_class_per_batch >= batch_size（破坏小类整组装箱不超容的前提）。
    关键实现细节:
        装箱顺序为「先小类后大类」。小类整组用 first-fit-decreasing 装箱：按组大小
        降序处理（同大小再按 task_type 排序保证确定性），每组放进第一个剩余容量足够
        的 batch；若现有 batch 都装不下就新开一个空 batch——因小类组大小
        ≤ min_class_per_batch < batch_size，新空 batch 必能容纳，故小类装箱永不抛
        ValueError，且保证整组不拆。再把大类样本（seed 确定性 shuffle 后）round-robin
        分发到所有现存 batch 填充剩余容量。这样小类聚集于单 batch、大类散布多 batch
        且与小类共箱，自然产生多类混合 batch（纯类切片会被 multiclass 断言拒绝）。
        nb = ceil(总题数/batch_size) 是初始 batch 数下界估计而非硬上限：小类装箱可能
        新开 bin 使实际 batch 数超过 nb。每次新开 bin 都意味着总容量随之增加，故总容量
        恒 ≥ 总题数，大类 round-robin 跳过满箱后仍能放下全部样本，不会违反 batch_size
        上限。题型按名称排序处理以保证跨运行确定性，不依赖 dict 遍历顺序。
    """
    _validate_params(batch_size, min_class_per_batch)

    rng = random.Random(seed)
    grouped = _select_mixed_by_task_type(items, correctness, correct_ratio, rng)
    total = sum(len(g) for g in grouped.values())
    if total == 0:
        return [], 0

    nb = max(1, math.ceil(total / batch_size))
    batches: list[list[GeneratedQuestion]] = [[] for _ in range(nb)]

    small, large = _split_by_size(grouped, min_class_per_batch)
    for group in _small_groups_decreasing(small):
        _pack_small_class(batches, group, batch_size)
    _distribute_large_classes(batches, large, batch_size, rng)

    result = [b for b in batches if b]
    selected_count = sum(len(b) for b in result)
    return result, selected_count


def _validate_params(batch_size: int, min_class_per_batch: int) -> None:
    """校验切分参数，非法值直接报错而非用默认值掩盖。

    除各自 >= 1 外，强制 min_class_per_batch < batch_size：小类组大小 ≤
    min_class_per_batch，唯有此前提成立才能保证小类整组放入单一 batch 而不超容；否则
    _pack_small_class 新开的 bin 会装入超 batch_size 的整组，静默违反容量合约。此约束
    与 config._validate_minibatch 一致，是 build_batches 对自身前提的防御性自校验（P5）。
    """
    if batch_size < 1:
        raise ValueError(f"batch_size 必须 >= 1, 实为 {batch_size}")
    if min_class_per_batch < 1:
        raise ValueError(f"min_class_per_batch 必须 >= 1, 实为 {min_class_per_batch}")
    if min_class_per_batch >= batch_size:
        raise ValueError(
            f"min_class_per_batch 必须严格 < batch_size, 否则无法保证小类整组放入单一 "
            f"batch 不超容; 实为 min_class_per_batch={min_class_per_batch}, "
            f"batch_size={batch_size}"
        )


def _split_by_size(
    grouped: dict[str, list[GeneratedQuestion]],
    min_class_per_batch: int,
) -> tuple[dict[str, list[GeneratedQuestion]], dict[str, list[GeneratedQuestion]]]:
    """按错题数把题型分为小类（≤ 阈值）与大类（> 阈值）两组。"""
    small = {t: g for t, g in grouped.items() if len(g) <= min_class_per_batch}
    large = {t: g for t, g in grouped.items() if len(g) > min_class_per_batch}
    return small, large


def _select_mixed_by_task_type(
    items: list[GeneratedQuestion],
    correctness: dict[str, bool],
    correct_ratio: float,
    rng: random.Random,
) -> dict[str, list[GeneratedQuestion]]:
    """按题型分组，为每组错题按比例采样正确题混入。

    只对有错题的题型做混合——无错题的题型不进 batch，即使有正确题。
    ``correct_ratio <= 0`` 时退化为纯错题模式（向后兼容）。

    参数:
        items: 候选题目全集。
        correctness: question_id -> 基线是否答对。
        correct_ratio: 正确题占比（0.0 ~ 1.0）。
        rng: 随机数发生器，用于采样正确题。
    返回:
        task_type -> 该题型的混合题目列表（错题全部 + 按比例采样的正确题）。
    """
    errors_by_type: dict[str, list[GeneratedQuestion]] = {}
    correct_by_type: dict[str, list[GeneratedQuestion]] = {}
    for q in items:
        qid = q.question_id
        if correctness.get(qid) is False:
            errors_by_type.setdefault(q.task_type, []).append(q)
        elif correctness.get(qid, False):
            correct_by_type.setdefault(q.task_type, []).append(q)

    if correct_ratio <= 0:
        return errors_by_type

    # 为每个有错题的 task_type 混入正确题
    grouped: dict[str, list[GeneratedQuestion]] = {}
    for task_type in sorted(errors_by_type):
        errs = errors_by_type[task_type]
        n_correct = round(len(errs) * correct_ratio / (1 - correct_ratio))
        available = correct_by_type.get(task_type, [])
        sampled = (
            list(available) if len(available) <= n_correct else rng.sample(available, n_correct)
        )
        grouped[task_type] = errs + sampled

    return grouped


def _small_groups_decreasing(
    small: dict[str, list[GeneratedQuestion]],
) -> list[list[GeneratedQuestion]]:
    """按组大小降序、同大小按 task_type 升序排出小类组（first-fit-decreasing 顺序）。

    参数:
        small: task_type -> 小类错题列表。
    返回:
        排好序的小类组列表；降序处理可降低碎片，确定性 tie-break 保证跨运行一致。
    """
    return [small[t] for t in sorted(small, key=lambda t: (-len(small[t]), t))]


def _pack_small_class(
    batches: list[list[GeneratedQuestion]],
    group: list[GeneratedQuestion],
    batch_size: int,
) -> None:
    """用 first-fit 把一个小类整组放入首个容得下的 batch，装不下则新开 bin（就地修改）。

    因小类组大小 ≤ min_class_per_batch < batch_size，新开的空 batch 必能容纳整组，
    故此函数永不抛 ValueError，且整组不拆。

    参数:
        batches: 当前各 batch（就地追加，必要时 append 新空 batch）。
        group: 待锁定的小类错题（整组不拆）。
        batch_size: 单 batch 容量上限。
    """
    for b in batches:
        if len(b) + len(group) <= batch_size:
            b.extend(group)
            return
    batches.append(list(group))


def _distribute_large_classes(
    batches: list[list[GeneratedQuestion]],
    large: dict[str, list[GeneratedQuestion]],
    batch_size: int,
    rng: random.Random,
) -> None:
    """将各大类样本 shuffle 后 round-robin 分发到所有现存 batch（就地修改）。

    参数:
        batches: 当前各 batch（含小类装箱可能新开的 bin，就地追加）。
        large: task_type -> 大类错题列表。
        batch_size: 单 batch 容量上限。
        rng: 复用的随机数发生器，保证 shuffle 确定性。
    异常:
        ValueError: 所有 batch 均满仍有样本未放置（总容量估算异常，合法输入不可达）。
    关键实现细节:
        轮转范围是「所有现存 batch」而非固定 nb 个——小类装箱新开的 bin 也参与分发。
        总容量 = 现存 batch 数 × batch_size，每次新开 bin 都同步抬高总容量，故总容量恒
        ≥ 总错题数，防御性 ValueError 在合法输入下不可达。全局指针在所有大类样本间持续
        轮转（不为每类重置），满箱即跳过，使大类充分散布并与已锁定的小类共箱。题型按名称
        排序以保证分发顺序确定。
    """
    nb = len(batches)
    pointer = 0
    for task_type in sorted(large):
        group = list(large[task_type])
        rng.shuffle(group)
        for q in group:
            pointer = _place_round_robin(batches, q, pointer, batch_size, nb)


def _place_round_robin(
    batches: list[list[GeneratedQuestion]],
    q: GeneratedQuestion,
    pointer: int,
    batch_size: int,
    nb: int,
) -> int:
    """从 pointer 起找第一个未满 batch 放入 q，返回下一次起始指针。

    参数:
        batches: 当前各 batch（就地追加）。
        q: 待放置的样本。
        pointer: 本次轮转起始 batch 下标。
        batch_size: 单 batch 容量上限。
        nb: batch 总数。
    返回:
        下一次轮转的起始指针（已前移一位）。
    异常:
        ValueError: 扫描一轮所有 batch 均满（总容量估算异常）。
    """
    for offset in range(nb):
        idx = (pointer + offset) % nb
        if len(batches[idx]) < batch_size:
            batches[idx].append(q)
            return (idx + 1) % nb
    raise ValueError("所有 batch 均满仍有样本待放置, 总容量估算异常")
