#!/usr/bin/env python
from __future__ import annotations

import argparse
import glob
import json
from collections import defaultdict
from pathlib import Path
from typing import Any


def _mean(values: list[float]) -> float | None:
    if not values:
        return None
    return float(sum(values) / len(values))


def _clip(value: float, limit: float | None) -> float:
    if limit is None or limit <= 0:
        return value
    return max(-float(limit), min(float(limit), value))


def _branch_values(row: dict[str, Any], objective: str) -> list[float]:
    if objective == "score":
        values = row.get("candidate_oracle_branch_scores")
    elif objective == "progress":
        values = row.get("candidate_oracle_branch_progress")
    elif objective == "success":
        values = [
            1.0 if bool(value) else 0.0
            for value in row.get("candidate_oracle_branch_successes", [])
        ]
    else:
        raise ValueError("objective must be 'score', 'progress', or 'success'")
    if not isinstance(values, list):
        return []
    parsed: list[float] = []
    for value in values:
        try:
            parsed.append(float(value))
        except (TypeError, ValueError):
            parsed.append(0.0)
    return parsed


def _rank_biases(
    values_by_rank: list[list[float]],
    *,
    scale: float,
    min_count: int,
    max_abs_bias: float | None,
) -> tuple[list[float], list[float | None], list[int]]:
    means = [_mean(values) for values in values_by_rank]
    counts = [len(values) for values in values_by_rank]
    baseline = means[0] if means and means[0] is not None and counts[0] >= min_count else None
    biases: list[float] = []
    for rank, mean in enumerate(means):
        if rank == 0 or baseline is None or mean is None or counts[rank] < min_count:
            biases.append(0.0)
        else:
            biases.append(_clip(float(scale) * (float(mean) - float(baseline)), max_abs_bias))
    return biases, means, counts


def _type_bonuses(
    values_by_type: dict[str, list[float]],
    *,
    scale: float,
    min_count: int,
    max_abs_bonus: float | None,
) -> tuple[dict[str, float], dict[str, float | None], dict[str, int], float | None]:
    pooled = [value for values in values_by_type.values() for value in values]
    baseline = _mean(pooled)
    means = {
        candidate_type: _mean(values)
        for candidate_type, values in sorted(values_by_type.items())
    }
    counts = {
        candidate_type: len(values)
        for candidate_type, values in sorted(values_by_type.items())
    }
    if baseline is None or len(pooled) < min_count:
        return {}, means, counts, baseline
    bonuses: dict[str, float] = {}
    for candidate_type, mean in means.items():
        if mean is None or counts[candidate_type] < min_count:
            continue
        bonuses[candidate_type] = _clip(
            float(scale) * (float(mean) - float(baseline)),
            max_abs_bonus,
        )
    return bonuses, means, counts, baseline


def _scale_bonuses(
    values_by_scale: dict[str, list[float]],
    *,
    scale: float,
    min_count: int,
    max_abs_bonus: float | None,
) -> tuple[dict[str, float], dict[str, float | None], dict[str, int], float | None]:
    pooled = [value for values in values_by_scale.values() for value in values]
    baseline = _mean(pooled)
    means = {
        residual_scale: _mean(values)
        for residual_scale, values in sorted(values_by_scale.items(), key=lambda item: float(item[0]))
    }
    counts = {
        residual_scale: len(values)
        for residual_scale, values in sorted(values_by_scale.items(), key=lambda item: float(item[0]))
    }
    if baseline is None or len(pooled) < min_count:
        return {}, means, counts, baseline
    bonuses: dict[str, float] = {}
    for residual_scale, mean in means.items():
        if mean is None or counts[residual_scale] < min_count:
            continue
        bonuses[residual_scale] = _clip(
            float(scale) * (float(mean) - float(baseline)),
            max_abs_bonus,
        )
    return bonuses, means, counts, baseline


def _iter_rollout_paths(patterns: list[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        matches = [Path(path) for path in glob.glob(pattern)]
        paths.extend(matches or [Path(pattern)])
    unique = sorted({path.resolve() for path in paths})
    missing = [path for path in unique if not path.exists()]
    if missing:
        raise FileNotFoundError(f"Missing rollout file(s): {missing}")
    return unique


def build_oracle_selector_calibration(
    rollout_paths: list[Path],
    *,
    objective: str,
    max_rank: int,
    rank_scale: float,
    type_scale: float,
    scale_scale: float,
    min_count: int,
    max_abs_rank_bias: float | None,
    max_abs_type_bonus: float | None,
    max_abs_scale_bonus: float | None,
) -> dict[str, Any]:
    if max_rank <= 0:
        raise ValueError("max_rank must be positive")
    if min_count <= 0:
        raise ValueError("min_count must be positive")

    rank_by_task: dict[str, list[list[float]]] = defaultdict(
        lambda: [[] for _ in range(max_rank)]
    )
    type_by_task: dict[str, dict[str, list[float]]] = defaultdict(
        lambda: defaultdict(list)
    )
    scale_by_task: dict[str, dict[str, list[float]]] = defaultdict(
        lambda: defaultdict(list)
    )
    global_ranks: list[list[float]] = [[] for _ in range(max_rank)]
    global_types: dict[str, list[float]] = defaultdict(list)
    global_scales: dict[str, list[float]] = defaultdict(list)
    rows_seen = 0
    rows_used = 0
    skipped_branches = 0

    for path in rollout_paths:
        payload = json.loads(path.read_text())
        for row in payload.get("rows", []):
            if not isinstance(row, dict):
                continue
            rows_seen += 1
            task_id = str(row.get("task_id") or "")
            values = _branch_values(row, objective)
            candidate_types = row.get("candidate_oracle_types")
            residual_scales = row.get("candidate_oracle_residual_scales")
            valid_mask = row.get("candidate_oracle_valid_mask")
            if (
                not task_id
                or not values
                or not isinstance(candidate_types, list)
                or not isinstance(residual_scales, list)
            ):
                continue
            if not isinstance(valid_mask, list):
                valid_mask = [True] * len(values)
            branch_count = min(
                max_rank,
                len(values),
                len(candidate_types),
                len(residual_scales),
                len(valid_mask),
            )
            if branch_count <= 0:
                continue
            rows_used += 1
            for rank in range(branch_count):
                if not bool(valid_mask[rank]):
                    skipped_branches += 1
                    continue
                value = float(values[rank])
                candidate_type = str(candidate_types[rank])
                residual_scale = f"{float(residual_scales[rank]):g}"
                rank_by_task[task_id][rank].append(value)
                global_ranks[rank].append(value)
                type_by_task[task_id][candidate_type].append(value)
                global_types[candidate_type].append(value)
                scale_by_task[task_id][residual_scale].append(value)
                global_scales[residual_scale].append(value)

    field_rank_biases_by_task: dict[str, list[float]] = {}
    rank_utility_means_by_task: dict[str, list[float | None]] = {}
    rank_counts_by_task: dict[str, list[int]] = {}
    for task_id in sorted(rank_by_task):
        biases, means, counts = _rank_biases(
            rank_by_task[task_id],
            scale=rank_scale,
            min_count=min_count,
            max_abs_bias=max_abs_rank_bias,
        )
        field_rank_biases_by_task[task_id] = biases
        rank_utility_means_by_task[task_id] = means
        rank_counts_by_task[task_id] = counts
    global_rank_biases, global_rank_means, global_rank_counts = _rank_biases(
        global_ranks,
        scale=rank_scale,
        min_count=min_count,
        max_abs_bias=max_abs_rank_bias,
    )
    field_rank_biases_by_task["*"] = global_rank_biases
    rank_utility_means_by_task["*"] = global_rank_means
    rank_counts_by_task["*"] = global_rank_counts

    candidate_type_bonuses_by_task: dict[str, dict[str, float]] = {}
    type_utility_means_by_task: dict[str, dict[str, float | None]] = {}
    type_counts_by_task: dict[str, dict[str, int]] = {}
    type_baselines_by_task: dict[str, float | None] = {}
    for task_id in sorted(type_by_task):
        bonuses, means, counts, baseline = _type_bonuses(
            type_by_task[task_id],
            scale=type_scale,
            min_count=min_count,
            max_abs_bonus=max_abs_type_bonus,
        )
        candidate_type_bonuses_by_task[task_id] = bonuses
        type_utility_means_by_task[task_id] = means
        type_counts_by_task[task_id] = counts
        type_baselines_by_task[task_id] = baseline
    global_type_bonuses, global_type_means, global_type_counts, global_type_baseline = (
        _type_bonuses(
            global_types,
            scale=type_scale,
            min_count=min_count,
            max_abs_bonus=max_abs_type_bonus,
        )
    )
    candidate_type_bonuses_by_task["*"] = global_type_bonuses
    type_utility_means_by_task["*"] = global_type_means
    type_counts_by_task["*"] = global_type_counts
    type_baselines_by_task["*"] = global_type_baseline

    residual_scale_bonuses_by_task: dict[str, dict[str, float]] = {}
    scale_utility_means_by_task: dict[str, dict[str, float | None]] = {}
    scale_counts_by_task: dict[str, dict[str, int]] = {}
    scale_baselines_by_task: dict[str, float | None] = {}
    for task_id in sorted(scale_by_task):
        bonuses, means, counts, baseline = _scale_bonuses(
            scale_by_task[task_id],
            scale=scale_scale,
            min_count=min_count,
            max_abs_bonus=max_abs_scale_bonus,
        )
        residual_scale_bonuses_by_task[task_id] = bonuses
        scale_utility_means_by_task[task_id] = means
        scale_counts_by_task[task_id] = counts
        scale_baselines_by_task[task_id] = baseline
    global_scale_bonuses, global_scale_means, global_scale_counts, global_scale_baseline = (
        _scale_bonuses(
            global_scales,
            scale=scale_scale,
            min_count=min_count,
            max_abs_bonus=max_abs_scale_bonus,
        )
    )
    residual_scale_bonuses_by_task["*"] = global_scale_bonuses
    scale_utility_means_by_task["*"] = global_scale_means
    scale_counts_by_task["*"] = global_scale_counts
    scale_baselines_by_task["*"] = global_scale_baseline

    return {
        "source_rollouts": [str(path) for path in rollout_paths],
        "calibration_source": "candidate_oracle_rollout",
        "objective": objective,
        "max_rank": int(max_rank),
        "rank_scale": float(rank_scale),
        "type_scale": float(type_scale),
        "scale_scale": float(scale_scale),
        "min_count": int(min_count),
        "max_abs_rank_bias": max_abs_rank_bias,
        "max_abs_type_bonus": max_abs_type_bonus,
        "max_abs_scale_bonus": max_abs_scale_bonus,
        "rows_seen": rows_seen,
        "rows_used": rows_used,
        "skipped_branches": skipped_branches,
        "field_rank_biases_by_task": field_rank_biases_by_task,
        "rank_utility_means_by_task": rank_utility_means_by_task,
        "rank_counts_by_task": rank_counts_by_task,
        "candidate_type_bonuses_by_task": candidate_type_bonuses_by_task,
        "type_utility_means_by_task": type_utility_means_by_task,
        "type_counts_by_task": type_counts_by_task,
        "type_baselines_by_task": type_baselines_by_task,
        "residual_scale_bonuses_by_task": residual_scale_bonuses_by_task,
        "scale_utility_means_by_task": scale_utility_means_by_task,
        "scale_counts_by_task": scale_counts_by_task,
        "scale_baselines_by_task": scale_baselines_by_task,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Build train-split selector calibration from candidate-oracle rollout traces."
        )
    )
    parser.add_argument("--rollout", action="append", required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--objective", choices=("score", "progress", "success"), default="score")
    parser.add_argument("--max-rank", type=int, default=4)
    parser.add_argument("--rank-scale", type=float, default=0.05)
    parser.add_argument("--type-scale", type=float, default=0.05)
    parser.add_argument("--scale-scale", type=float, default=0.0)
    parser.add_argument("--min-count", type=int, default=20)
    parser.add_argument("--max-abs-rank-bias", type=float, default=0.02)
    parser.add_argument("--max-abs-type-bonus", type=float, default=0.02)
    parser.add_argument("--max-abs-scale-bonus", type=float, default=0.02)
    args = parser.parse_args(argv)

    max_abs_rank_bias = args.max_abs_rank_bias if args.max_abs_rank_bias > 0 else None
    max_abs_type_bonus = args.max_abs_type_bonus if args.max_abs_type_bonus > 0 else None
    max_abs_scale_bonus = args.max_abs_scale_bonus if args.max_abs_scale_bonus > 0 else None
    rollout_paths = _iter_rollout_paths(args.rollout)
    result = build_oracle_selector_calibration(
        rollout_paths,
        objective=args.objective,
        max_rank=args.max_rank,
        rank_scale=args.rank_scale,
        type_scale=args.type_scale,
        scale_scale=args.scale_scale,
        min_count=args.min_count,
        max_abs_rank_bias=max_abs_rank_bias,
        max_abs_type_bonus=max_abs_type_bonus,
        max_abs_scale_bonus=max_abs_scale_bonus,
    )
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2) + "\n")
    print(
        json.dumps(
            {
                key: value
                for key, value in result.items()
                if key
                not in {
                    "field_rank_biases_by_task",
                    "rank_utility_means_by_task",
                    "rank_counts_by_task",
                    "candidate_type_bonuses_by_task",
                    "type_utility_means_by_task",
                    "type_counts_by_task",
                    "type_baselines_by_task",
                    "residual_scale_bonuses_by_task",
                    "scale_utility_means_by_task",
                    "scale_counts_by_task",
                    "scale_baselines_by_task",
                }
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())