# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility functions for GenRM comparison server.

This module provides:
- Prompt key and conversation history extraction (for cohort grouping)
- Comparison pair generation strategies (circular, all_pairs)
- GenRM output parsing (JSON score extraction)
- Response API object text extraction
- Score aggregation with tiebreaker logic
- Length-based bonus/penalty computation
"""

from __future__ import annotations

import hashlib
import itertools
import json
import logging
import re
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
from pydantic import BaseModel


logger = logging.getLogger(__name__)


# =============================================================================
# Constants
# =============================================================================

SUPPORTED_COMPARISON_STRATEGIES = frozenset({"circular", "all_pairs"})
SUPPORTED_AGGREGATOR_METHODS = frozenset({"simple_tiebreaker"})

# GenRM ranking midpoint - rankings are 1-6, midpoint is 3.5
# Used in tiebreaker: ranking < 3.5 means response_1 is better, > 3.5 means response_2 is better
RANKING_MIDPOINT = 3.5

# Default placeholder for empty/missing output
EMPTY_OUTPUT_PLACEHOLDER = "None"


# =============================================================================
# Exceptions
# =============================================================================


class GenRMOutputParseError(ValueError):
    """Raised when GenRM output cannot be parsed into expected JSON scores.

    Expected format: {"score_1": <1-5>, "score_2": <1-5>, "ranking": <1-6>}
    """

    pass


# =============================================================================
# Prompt key and conversation history (for cohort-based verify)
# =============================================================================


def _to_jsonable_data(value: Any) -> Any:
    """Recursively convert request payloads into JSON-serializable data."""
    if isinstance(value, BaseModel):
        return value.model_dump(mode="json")
    if isinstance(value, dict):
        return {str(k): _to_jsonable_data(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_jsonable_data(item) for item in value]
    return value


def get_prompt_key_from_input(input_messages: List[Any], principle: Optional[str] = None) -> str:
    """Stable key for grouping rollouts by prompt and principle.

    Used by cohort-based verify to group N rollouts per prompt.

    Args:
        input_messages: Conversation input (e.g. responses_create_params.input)
        principle: Optional principle string (principle-based GenRM)

    Returns:
        A stable string key for the prompt + principle combination
    """
    key_data = {"input": _to_jsonable_data(input_messages), "principle": principle}
    return hashlib.sha256(json.dumps(key_data, sort_keys=True).encode()).hexdigest()


def get_prompt_key(example: Dict) -> str:
    """Get stable key for grouping examples by prompt and principle.

    Examples with the same conversation history but different principles
    should be in separate groups. Accepts either prompt_id or hashes input + principle.

    Args:
        example: Dict with optional "prompt_id", or "responses_create_params" with "input"

    Returns:
        A stable string key
    """
    if "prompt_id" in example:
        prompt_id = str(example["prompt_id"])
        principle = example.get("principle")
        if principle is not None:
            return f"{prompt_id}::{principle}"
        return prompt_id
    conv = extract_conversation_history(example)
    principle = example.get("principle")
    return get_prompt_key_from_input(conv, principle)


def extract_conversation_history(example: Dict) -> List[Dict]:
    """Extract conversation history from example (responses_create_params.input).

    Args:
        example: Dict with responses_create_params.input (list of message dicts)

    Returns:
        List of message dicts with 'role' and 'content'

    Raises:
        ValueError: If required fields are missing
    """
    responses_create_params = example.get("responses_create_params")
    if responses_create_params is None:
        raise ValueError(f"Example missing 'responses_create_params': {list(example.keys())}")
    if "input" not in responses_create_params:
        raise ValueError(f"responses_create_params missing 'input': {list(responses_create_params.keys())}")
    return responses_create_params["input"]


# =============================================================================
# Comparison Strategy
# =============================================================================


def generate_comparison_pairs(strategy: str, num_responses: int) -> List[Tuple[int, int]]:
    """Generate pairs of response indices for pairwise comparison.

    Args:
        strategy: Comparison strategy - "circular" or "all_pairs"
            - "circular": Each response compared with next (N comparisons)
            - "all_pairs": Every pair compared (N*(N-1)/2 comparisons)
        num_responses: Number of responses to compare

    Returns:
        List of (i, j) tuples where i < j for all_pairs, or circular neighbors

    Raises:
        ValueError: If strategy is not supported or num_responses < 2
    """
    if strategy not in SUPPORTED_COMPARISON_STRATEGIES:
        raise ValueError(
            f"Unknown comparison strategy: '{strategy}'. Supported: {sorted(SUPPORTED_COMPARISON_STRATEGIES)}"
        )

    if num_responses < 2:
        raise ValueError(f"Need at least 2 responses for comparison, got {num_responses}")

    if strategy == "all_pairs":
        return list(itertools.combinations(range(num_responses), 2))
    else:  # circular
        return [(i, (i + 1) % num_responses) for i in range(num_responses)]


# =============================================================================
# GenRM Output Parsing
# =============================================================================


def parse_genrm_output(
    output: str,
    default_score: float,
    default_ranking: float,
    *,
    raise_on_fail: bool = False,
) -> Tuple[float, float, float]:
    """Parse GenRM output to extract scores from JSON format.

    Searches for JSON in the output text, trying:
    1. Fenced JSON blocks (```json {...} ```)
    2. Any {...} JSON objects, taking the last valid one

    Expected JSON format:
        {"score_1": <1-5>, "score_2": <1-5>, "ranking": <1-6>}

    Args:
        output: Raw text output from GenRM model
        default_score: Default score if parsing fails
        default_ranking: Default ranking if parsing fails
        raise_on_fail: If True, raise GenRMOutputParseError on failure

    Returns:
        Tuple of (score_1, score_2, ranking)

    Raises:
        GenRMOutputParseError: If raise_on_fail=True and parsing fails
    """

    def _try_parse(json_str: str) -> Optional[Tuple[float, float, float]]:
        """Attempt to parse a JSON string into scores."""
        try:
            parsed = json.loads(json_str)
        except json.JSONDecodeError:
            return None

        if not isinstance(parsed, dict):
            return None

        # Handle nested format: {"rubric_evaluations": [...], "overall": {...}}
        if "overall" in parsed and isinstance(parsed["overall"], dict):
            parsed = parsed["overall"]

        # Must have at least one expected key
        if not any(k in parsed for k in ("score_1", "score_2", "ranking")):
            return None

        try:
            score_1 = float(parsed.get("score_1", default_score))
            score_2 = float(parsed.get("score_2", default_score))
            ranking = float(parsed.get("ranking", default_ranking))
            return score_1, score_2, ranking
        except (TypeError, ValueError):
            return None

    def _find_json_objects(text: str) -> List[str]:
        """Find top-level JSON objects in text using brace counting."""
        results = []
        i = 0
        while i < len(text):
            if text[i] == "{":
                depth = 0
                start = i
                in_string = False
                escape_next = False
                for j in range(i, len(text)):
                    ch = text[j]
                    if escape_next:
                        escape_next = False
                        continue
                    if ch == "\\" and in_string:
                        escape_next = True
                        continue
                    if ch == '"' and not escape_next:
                        in_string = not in_string
                        continue
                    if in_string:
                        continue
                    if ch == "{":
                        depth += 1
                    elif ch == "}":
                        depth -= 1
                        if depth == 0:
                            results.append(text[start : j + 1])
                            i = j
                            break
            i += 1
        return results

    try:
        # Strategy 1: Look for fenced JSON blocks (```json ... ```)
        for match in re.finditer(r"```json\s*([\s\S]*?)\s*```", output, flags=re.IGNORECASE):
            for json_str in _find_json_objects(match.group(1)):
                result = _try_parse(json_str)
                if result is not None:
                    return result

        # Strategy 2: Find all top-level {...} and take the last valid one.
        #
        # Note: We keep this intentionally permissive because model outputs can include
        # extra prose around a JSON blob, and the JSON may be pretty-printed across lines.
        last_valid: Optional[Tuple[float, float, float]] = None
        for json_str in _find_json_objects(output):
            result = _try_parse(json_str)
            if result is not None:
                last_valid = result

        if last_valid is not None:
            return last_valid

        # Parsing failed
        preview = output[:200] + "..." if len(output) > 200 else output
        msg = f"No parseable JSON found in GenRM output: {preview}"

        if raise_on_fail:
            raise GenRMOutputParseError(msg)

        logger.warning(msg)
        return default_score, default_score, default_ranking
    except Exception as e:
        preview = output[:200] + "..." if len(output) > 200 else output
        msg = f"Error parsing GenRM output: {e}. Output: {preview}"
        if raise_on_fail:
            raise GenRMOutputParseError(msg) from e
        logger.exception(msg)
        return default_score, default_score, default_ranking


# =============================================================================
# Response API Object Extraction
# =============================================================================


def extract_from_response_obj(response_obj: Dict[str, Any]) -> Tuple[str, str]:
    """Extract reasoning and output text from a Response API object.

    Parses the nested Response API structure to find:
    - Reasoning content from "reasoning" type items
    - Output text from "message" type items with "output_text" content

    Args:
        response_obj: Raw Response API object with "output" field

    Returns:
        Tuple of (reasoning_content, output_text)
    """
    reasoning_content = ""
    output_text = ""

    if not isinstance(response_obj, dict):
        return reasoning_content, output_text

    output = response_obj.get("output", [])
    if not isinstance(output, list):
        return reasoning_content, output_text

    for item in output:
        if not isinstance(item, dict):
            continue

        item_type = item.get("type", "")

        if item_type == "reasoning":
            # Extract from summary field
            summary = item.get("summary", [])
            if isinstance(summary, list):
                for s in summary:
                    if isinstance(s, dict) and "text" in s:
                        reasoning_content += s.get("text", "")

        elif item_type == "message":
            # Extract from content field
            content = item.get("content", [])
            if isinstance(content, list):
                for c in content:
                    if isinstance(c, dict) and c.get("type") == "output_text":
                        output_text += c.get("text", "")

    return reasoning_content, output_text


def extract_output_text(response_obj: Dict[str, Any]) -> str:
    """Extract only the output text (final answer) from Response API object.

    Args:
        response_obj: Raw Response API object

    Returns:
        The output text, or "None" if empty/missing
    """
    _, output = extract_from_response_obj(response_obj)

    if not output or not output.strip():
        return EMPTY_OUTPUT_PLACEHOLDER

    return output


# =============================================================================
# Style Element Counting (Arena-Hard style + emoji)
# =============================================================================

# Regex pattern to match code blocks (excluded from style counting)
# NOTE: Non-greedy, matches across newlines, and tolerates backticks inside code.
_CODE_BLOCK_PATTERN = re.compile(r"```[\s\S]*?```", re.MULTILINE)

# Emoji detection pattern - covers most common emoji ranges
_EMOJI_CHAR_PATTERN = re.compile(
    "["
    "\U0001f600-\U0001f64f"  # emoticons
    "\U0001f300-\U0001f5ff"  # symbols & pictographs
    "\U0001f680-\U0001f6ff"  # transport & map symbols
    "\U0001f1e0-\U0001f1ff"  # flags
    "\U00002702-\U000027b0"  # dingbats
    "\U000024c2-\U0001f251"  # enclosed characters
    "\U0001f900-\U0001f9ff"  # supplemental symbols
    "\U0001fa00-\U0001fa6f"  # chess symbols
    "\U0001fa70-\U0001faff"  # symbols and pictographs extended-A
    "\U00002600-\U000026ff"  # misc symbols
    "\U00002700-\U000027bf"  # dingbats
    "]",
    re.UNICODE,
)


def count_style_elements(text: str) -> Dict[str, int]:
    """Count style elements in text (Arena-Hard style + emoji).

    Counts markdown formatting elements and emojis, excluding code blocks.

    Args:
        text: Text to analyze

    Returns:
        Dictionary with counts for each style element type:
        - header_count: Total headers (h1-h6)
        - list_count: Total list items (ordered + unordered)
        - bold_count: Total bold text occurrences
        - emoji_count: Total emoji occurrences
    """
    # Remove code blocks before counting (same as Arena-Hard)
    text_no_code = _CODE_BLOCK_PATTERN.sub("", text)

    # Count headers (h1-h6)
    header_count = len(re.findall(r"^#{1,6}\s", text_no_code, re.MULTILINE))

    # Count lists (ordered + unordered)
    ordered_list = len(re.findall(r"^\s*\d+\.\s", text_no_code, re.MULTILINE))
    unordered_list = len(re.findall(r"^\s*[-*+]\s", text_no_code, re.MULTILINE))
    list_count = ordered_list + unordered_list

    # Count bold text
    bold_asterisk = len(re.findall(r"\*\*[^*\n]+\*\*", text_no_code))
    bold_underscore = len(re.findall(r"__[^_\n]+__", text_no_code))
    bold_count = bold_asterisk + bold_underscore

    # Count emojis
    # NOTE: this counts individual emoji codepoints in common emoji ranges.
    # It will *not* perfectly count grapheme clusters for ZWJ sequences, but it
    # avoids undercounting when multiple emojis appear back-to-back.
    emoji_count = len(_EMOJI_CHAR_PATTERN.findall(text_no_code))

    return {
        "header_count": header_count,
        "list_count": list_count,
        "bold_count": bold_count,
        "emoji_count": emoji_count,
    }


def compute_style_density(text: str) -> float:
    """Compute style density (total style elements / text length).

    Args:
        text: Text to analyze

    Returns:
        Style density value (elements per character). Returns 0 if text is empty.
    """
    # Keep numerator/denominator consistent: exclude code blocks from both.
    text_no_code = _CODE_BLOCK_PATTERN.sub("", text)
    text_len = len(text_no_code.strip())
    if text_len <= 0:
        return 0.0

    counts = count_style_elements(text_no_code)
    total_elements = sum(counts.values())

    return total_elements / text_len


def _compute_style_density_weights(densities: List[float]) -> List[float]:
    """Compute zero-centered weights where lower density = higher weight.

    Similar to _compute_length_weights but for style density.
    Lower style density (less formatting) gets positive weight (bonus).
    Higher style density (more formatting) gets negative weight (penalty).

    Args:
        densities: List of style density values

    Returns:
        List of weights, zero-centered (sum to ~0)
    """
    max_density = max(densities)
    min_density = min(densities)

    if max_density == min_density:
        return [0.0] * len(densities)

    span = max_density - min_density
    # Lower density = weight closer to 1, higher density = weight closer to 0
    raw_weights = [1.0 - ((density - min_density) / span) for density in densities]
    # Zero-center
    mean_weight = sum(raw_weights) / len(raw_weights)
    return [w - mean_weight for w in raw_weights]


def apply_style_penalties(
    scores: List[float],
    response_objs: List[Dict[str, Any]],
    group_style_penalty_coeff: float,
) -> Tuple[List[float], List[float]]:
    """Apply style-based penalties to scores.

    Penalizes responses with higher style density (more formatting elements
    relative to text length). This controls for "style hacking" where models
    produce verbose formatting to appear more helpful.

    Style elements counted (following Arena-Hard):
    - Headers (h1-h6)
    - Lists (ordered and unordered)
    - Bold text
    - Emojis (extension beyond Arena-Hard)

    Args:
        scores: Base scores to adjust (modified in place)
        response_objs: Response API objects to extract text from
        group_style_penalty_coeff: Coefficient for style density penalty.
            Higher values penalize formatting more heavily.

    Returns:
        Tuple of (adjusted_scores, style_adjustments_per_response)
    """
    num_responses = len(response_objs)

    if num_responses < 2 or group_style_penalty_coeff <= 0:
        return scores, [0.0] * len(scores)

    # Compute style density for each response
    style_adjustments = [0.0] * num_responses
    answer_densities: List[float] = []

    for obj in response_objs:
        _, answer = extract_from_response_obj(obj)
        density = compute_style_density(answer)
        answer_densities.append(density)

    logger.info(f"[StyleControl] Densities: {[f'{d:.6f}' for d in answer_densities]}")
    logger.info(f"[StyleControl] Scores before: {[f'{s:.4f}' for s in scores]}")

    # Compute zero-centered weights (lower density = higher/positive weight)
    style_weights = _compute_style_density_weights(answer_densities)

    for idx in range(num_responses):
        adjustment = style_weights[idx] * group_style_penalty_coeff
        if adjustment != 0:
            scores[idx] += adjustment
            style_adjustments[idx] = adjustment

    logger.info(f"[StyleControl] Scores after: {[f'{s:.4f}' for s in scores]}")
    logger.info(f"[StyleControl] Adjustments: {[f'{a:+.4f}' for a in style_adjustments]}")

    return scores, style_adjustments


# =============================================================================
# Length-Based Bonuses
# =============================================================================


def apply_length_bonuses(
    scores: List[float],
    response_objs: List[Dict[str, Any]],
    reasoning_bonus: float,
    answer_bonus: float,
    top_percentile: float,
    group_reasoning_length_penalty_coeff: float,
    group_answer_length_penalty_coeff: float,
) -> Tuple[List[float], List[float]]:
    """Apply length-based bonuses/penalties to scores.

    Two types of adjustments:
    1. Top-performer bonuses: Shortest reasoning/answer among top scorers gets bonus
    2. Group-relative penalties: Scores adjusted based on relative length within group

    Args:
        scores: Base scores to adjust (modified in place)
        response_objs: Response API objects to extract lengths from
        reasoning_bonus: Bonus for shortest reasoning among top performers
        answer_bonus: Bonus for shortest answer among top performers
        top_percentile: Fraction of top scorers eligible for bonuses (0.0-1.0)
        group_reasoning_length_penalty_coeff: Coefficient for reasoning length penalty
        group_answer_length_penalty_coeff: Coefficient for answer length penalty

    Returns:
        Tuple of (adjusted_scores, bonuses_per_response)
    """
    num_responses = len(response_objs)

    if num_responses < 2:
        return scores, [0.0] * len(scores)

    # Extract lengths from Response API objects
    bonuses = [0.0] * num_responses
    reasoning_lengths: List[int] = []
    answer_lengths: List[int] = []

    for obj in response_objs:
        reasoning, answer = extract_from_response_obj(obj)
        reasoning_lengths.append(len(reasoning.strip()))
        answer_lengths.append(len(answer.strip()))

    logger.debug(f"Reasoning lengths: {reasoning_lengths}")
    logger.debug(f"Answer lengths: {answer_lengths}")

    # Determine top percentile threshold
    sorted_scores = sorted(scores, reverse=True)
    threshold_idx = max(0, int(len(scores) * top_percentile) - 1)
    top_threshold = sorted_scores[threshold_idx]

    # Bonus for shortest non-empty reasoning among top performers
    if reasoning_bonus > 0:
        valid = [(i, length) for i, length in enumerate(reasoning_lengths) if length > 0]
        if valid:
            idx, length = min(valid, key=lambda x: x[1])
            if scores[idx] >= top_threshold:
                scores[idx] += reasoning_bonus
                bonuses[idx] += reasoning_bonus
                logger.debug(f"Reasoning bonus +{reasoning_bonus} to response {idx} (len={length})")

    # Bonus for shortest non-empty answer among top performers
    if answer_bonus > 0:
        valid = [(i, length) for i, length in enumerate(answer_lengths) if length > 0]
        if valid:
            idx, length = min(valid, key=lambda x: x[1])
            if scores[idx] >= top_threshold:
                scores[idx] += answer_bonus
                bonuses[idx] += answer_bonus
                logger.debug(f"Answer bonus +{answer_bonus} to response {idx} (len={length})")

    # Group-relative length adjustment (shorter = higher weight, zero-centered)
    if group_reasoning_length_penalty_coeff > 0 or group_answer_length_penalty_coeff > 0:
        reasoning_weights = _compute_length_weights(reasoning_lengths)
        answer_weights = _compute_length_weights(answer_lengths)

        for idx in range(num_responses):
            reasoning_adj = reasoning_weights[idx] * group_reasoning_length_penalty_coeff
            answer_adj = answer_weights[idx] * group_answer_length_penalty_coeff
            total_adj = reasoning_adj + answer_adj

            if total_adj != 0:
                scores[idx] += total_adj
                bonuses[idx] += total_adj
                logger.debug(
                    f"Length adjustment {total_adj:+.4f} to response {idx} "
                    f"(reasoning={reasoning_adj:+.4f}, answer={answer_adj:+.4f})"
                )

    return scores, bonuses


def _compute_length_weights(lengths: List[int]) -> List[float]:
    """Compute zero-centered weights where shorter = higher weight.

    Args:
        lengths: List of lengths

    Returns:
        List of weights, zero-centered (sum to ~0)
    """
    max_len, min_len = max(lengths), min(lengths)

    if max_len == min_len:
        return [0.0] * len(lengths)

    span = max_len - min_len
    # Shorter = weight closer to 1, longer = weight closer to 0
    raw_weights = [1.0 - ((length - min_len) / span) for length in lengths]
    # Zero-center
    mean_weight = sum(raw_weights) / len(raw_weights)
    return [w - mean_weight for w in raw_weights]


# =============================================================================
# Score Aggregation
# =============================================================================


def aggregate_scores(
    comparison_results: List[Tuple[float, float, float]],
    comparison_metadata: List[Tuple[int, int, int]],
    response_objs: List[Dict[str, Any]],
    aggregator_method: str,
    default_score: float,
    reasoning_bonus: float,
    answer_bonus: float,
    top_percentile: float,
    group_reasoning_length_penalty_coeff: float,
    group_answer_length_penalty_coeff: float,
    group_style_penalty_coeff: float = 0.0,
) -> Tuple[List[float], Dict[str, float], List[float], List[float]]:
    """Aggregate pairwise comparison results into per-response rewards.

    For "simple_tiebreaker" method:
    - When score_1 == score_2, use ranking to break the tie
    - ranking < 3.5 means response_1 is better → boost score_1, penalize score_2
    - ranking > 3.5 means response_2 is better → boost score_2, penalize score_1

    Args:
        comparison_results: List of (score_1, score_2, ranking) from pairwise comparisons
        comparison_metadata: List of (response_i, response_j, judge_idx) for each comparison
        response_objs: Raw Response API objects for length bonus computation
        aggregator_method: Only "simple_tiebreaker" is supported
        default_score: Default score when no comparisons exist for a response
        reasoning_bonus: Bonus for shortest reasoning among top performers
        answer_bonus: Bonus for shortest answer among top performers
        top_percentile: Percentile threshold for length bonuses
        group_reasoning_length_penalty_coeff: Coefficient for reasoning length penalty
        group_answer_length_penalty_coeff: Coefficient for answer length penalty
        group_style_penalty_coeff: Coefficient for style density penalty.
            Penalizes responses with higher density of formatting elements
            (headers, lists, bold, emojis) relative to text length.

    Returns:
        Tuple of:
        - final_scores: Per-response rewards after all adjustments
        - metrics: Aggregation statistics (mean, std, tiebreak rate)
        - base_scores: Scores before length/style bonuses
        - bonuses: Length and style adjustments applied to each response

    Raises:
        ValueError: If aggregator_method is not supported
    """
    if aggregator_method not in SUPPORTED_AGGREGATOR_METHODS:
        raise ValueError(
            f"Unsupported aggregator_method: '{aggregator_method}'. Supported: {sorted(SUPPORTED_AGGREGATOR_METHODS)}"
        )

    num_responses = len(response_objs)

    # Initialize accumulators
    accumulated_scores = [0.0] * num_responses
    comparison_counts = [0] * num_responses

    # Track metrics
    all_individual_scores: List[float] = []
    tiebreak_count = 0

    # Process each comparison
    for (score_1, score_2, ranking), (i, j, _judge_idx) in zip(comparison_results, comparison_metadata):
        all_individual_scores.extend([score_1, score_2])

        # Apply tiebreaker when scores are equal
        if score_1 == score_2:
            tiebreak_count += 1
            # ranking < 3.5 → response_1 better, ranking > 3.5 → response_2 better
            adjustment = RANKING_MIDPOINT - ranking
            score_1 = score_1 + adjustment
            score_2 = score_2 - adjustment

        # Accumulate
        accumulated_scores[i] += score_1
        accumulated_scores[j] += score_2
        comparison_counts[i] += 1
        comparison_counts[j] += 1

    # Compute average scores
    final_scores = [
        accumulated_scores[idx] / comparison_counts[idx] if comparison_counts[idx] > 0 else default_score
        for idx in range(num_responses)
    ]

    # Store base scores before length/style adjustments
    base_scores = list(final_scores)
    bonuses = [0.0] * num_responses

    # Apply length bonuses if any are configured
    if any(
        [
            reasoning_bonus > 0,
            answer_bonus > 0,
            group_reasoning_length_penalty_coeff > 0,
            group_answer_length_penalty_coeff > 0,
        ]
    ):
        final_scores, length_bonuses = apply_length_bonuses(
            scores=final_scores,
            response_objs=response_objs,
            reasoning_bonus=reasoning_bonus,
            answer_bonus=answer_bonus,
            top_percentile=top_percentile,
            group_reasoning_length_penalty_coeff=group_reasoning_length_penalty_coeff,
            group_answer_length_penalty_coeff=group_answer_length_penalty_coeff,
        )
        bonuses = [b + lb for b, lb in zip(bonuses, length_bonuses)]

    # Apply style penalties if configured
    if group_style_penalty_coeff > 0:
        final_scores, style_adjustments = apply_style_penalties(
            scores=final_scores,
            response_objs=response_objs,
            group_style_penalty_coeff=group_style_penalty_coeff,
        )
        bonuses = [b + sa for b, sa in zip(bonuses, style_adjustments)]

    # Compute metrics
    metrics: Dict[str, float] = {}

    if all_individual_scores:
        scores_array = np.array(all_individual_scores)
        metrics["mean_individual_score"] = float(np.mean(scores_array))
        metrics["std_individual_score"] = float(np.std(scores_array))

    if comparison_results:
        metrics["tiebreak_usage_rate"] = tiebreak_count / len(comparison_results)

    return final_scores, metrics, base_scores, bonuses
