"""Pareto-front sampling helpers used by SkyDiscover's ``gepa_native`` database.

This is the small utility copied by SkyDiscover from ``gepa/gepa_utils.py``.  It does not compute a
conventional vector Pareto frontier.  Instead, each metric supplies the set of programs tied at its
best value; dominated programs are removed and the survivors are sampled in proportion to how many
metric fronts they occupy.
"""
# Ported by SkyDiscover from gepa/src/gepa/gepa_utils.py.
# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors.
from __future__ import annotations

import random
from typing import Any, Mapping


def is_dominated(y, programs, program_at_pareto_front_valset):
    y_fronts = [front for front in program_at_pareto_front_valset.values() if y in front]
    for front in y_fronts:
        found_dominator_in_front = False
        for other_prog in front:
            if other_prog in programs:
                found_dominator_in_front = True
                break
        if not found_dominator_in_front:
            return False
    return True


def remove_dominated_programs(program_at_pareto_front_valset, scores=None):
    freq = {}
    for front in program_at_pareto_front_valset.values():
        for program in front:
            freq[program] = freq.get(program, 0) + 1

    dominated = set()
    programs = list(freq)
    if scores is None:
        scores = dict.fromkeys(programs, 1)
    programs = sorted(programs, key=lambda program: scores[program], reverse=False)

    found_to_remove = True
    while found_to_remove:
        found_to_remove = False
        for program in programs:
            if program in dominated:
                continue
            if is_dominated(
                program,
                set(programs).difference({program}).difference(dominated),
                program_at_pareto_front_valset,
            ):
                dominated.add(program)
                found_to_remove = True
                break

    dominators = [program for program in programs if program not in dominated]
    return {
        val_id: {program for program in front if program in dominators}
        for val_id, front in program_at_pareto_front_valset.items()
    }


def select_program_candidate_from_pareto_front(
    pareto_front_programs: Mapping[Any, set],
    scores: Mapping[Any, float],
    rng: random.Random,
):
    new_front = remove_dominated_programs(pareto_front_programs, scores=scores)
    freq = {}
    for front in new_front.values():
        for program_id in front:
            freq[program_id] = freq.get(program_id, 0) + 1
    sampling_list = [
        program_id
        for program_id, frequency in freq.items()
        for _ in range(frequency)
    ]
    assert sampling_list
    return rng.choice(sampling_list)