src/submission/process_zip.py

import json
import os
import shutil
import tempfile
import zipfile
from dataclasses import dataclass

from src.about import MMRED_GROUP_KEYS, MMRED_SUBTASK_KEYS, Tasks
from src.leaderboard.metrics import extract_exact_match_max_flexible_strict, extract_metric_value


LOGS_PUBLIC_ZIP = "logs_public.zip"
LOGS_PUBLIC_DIR = "logs_public"
RESULTS_PREFIX = "results_"


@dataclass
class ParsedResultFile:
    path: str
    benchmark: str
    data: dict


TASK_KEY_ALIASES: dict[str, set[str]] = {
    task.value.benchmark: {task.value.benchmark, task.value.col_name.lower()}
    for task in Tasks
}
TASK_KEY_ALIASES["t_math"].update({"tmath", "t-math"})
for task in Tasks:
    benchmark = task.value.benchmark
    TASK_KEY_ALIASES[benchmark].add(benchmark.replace("_", ""))


def _normalize_task_key(task_key: str) -> str:
    key = task_key.lower().removesuffix("_gen")
    return key.replace("-", "_")


def _benchmark_for_task_key(task_key: str) -> str | None:
    normalized = _normalize_task_key(task_key)
    if normalized == "tmath":
        return "t_math"
    if normalized == "mmred" or normalized.startswith("mmred_"):
        return "mmred"
    for benchmark, aliases in TASK_KEY_ALIASES.items():
        alias_norm = {_normalize_task_key(alias) for alias in aliases}
        if normalized in alias_norm or normalized.startswith(benchmark):
            return benchmark
    return None


def _submission_name_candidates(task) -> set[str]:
    benchmark = task.value.benchmark
    col_name = task.value.col_name
    return {
        benchmark.lower(),
        col_name.lower(),
        benchmark.replace("_", "").lower(),
        col_name.replace("-", "").lower(),
        col_name.replace("-", "_").lower(),
        "tmath",
    }


BENCHMARKS_MAX_EXACT_MATCH = {"ruaime", "luzitania"}


def _extract_exact_match_score(benchmark: str, task_result: dict) -> float | None:
    if benchmark in BENCHMARKS_MAX_EXACT_MATCH:
        return extract_exact_match_max_flexible_strict(task_result)
    return extract_metric_value(task_result, "exact_match")


def _extract_mmred_scores(task_key: str, task_result: dict, scores: dict[str, float]) -> None:
    normalized = _normalize_task_key(task_key)
    if normalized == "mmred":
        value = extract_metric_value(task_result, "em.dc_aggregate")
        if value is not None:
            scores["em.dc_aggregate"] = value
        return
    if normalized in MMRED_GROUP_KEYS:
        value = extract_metric_value(task_result, "em.dc_aggregate")
        if value is not None:
            scores[f"{normalized}::em.dc_aggregate"] = value
        return
    if normalized in MMRED_SUBTASK_KEYS:
        value = extract_metric_value(task_result, "exact_match")
        if value is not None:
            scores[f"{normalized}::exact_match"] = value


def _extract_benchmark_scores(task_key: str, task_result: dict) -> dict[str, float]:
    benchmark = _benchmark_for_task_key(task_key)
    if benchmark is None:
        return {}

    if benchmark == "mmred":
        scores: dict[str, float] = {}
        _extract_mmred_scores(task_key, task_result, scores)
        return scores

    value = _extract_exact_match_score(benchmark, task_result)
    if value is None:
        return {}
    return {"exact_match": value}


def parse_results_data(data: dict) -> dict[str, dict[str, float]]:
    merged: dict[str, dict[str, float]] = {}
    for task_key, task_result in data.get("results", {}).items():
        if not isinstance(task_result, dict):
            continue
        benchmark = _benchmark_for_task_key(task_key)
        if benchmark is None:
            continue
        scores = _extract_benchmark_scores(task_key, task_result)
        if scores:
            merged.setdefault(benchmark, {}).update(scores)
    return merged


def get_scores_from_result(parsed: ParsedResultFile) -> dict[str, float]:
    return parse_results_data(parsed.data).get(parsed.benchmark, {})


def extract_submission_archive(zip_path: str) -> str:
    if not zip_path or not os.path.isfile(zip_path):
        raise ValueError("Please upload a valid zip submission file.")
    if not zipfile.is_zipfile(zip_path):
        raise ValueError("Uploaded file is not a zip archive.")

    extract_dir = tempfile.mkdtemp(prefix="mera_reason_submission_")
    with zipfile.ZipFile(zip_path, "r") as archive:
        archive.extractall(extract_dir)
    return extract_dir


def _find_submission_root(extract_dir: str) -> str:
    entries = [
        name
        for name in os.listdir(extract_dir)
        if not name.startswith(".") and name not in ("__MACOSX",)
    ]
    if len(entries) == 1:
        only_entry = os.path.join(extract_dir, entries[0])
        if os.path.isdir(only_entry):
            return only_entry
    return extract_dir


def _resolve_logs_public_dir(submission_root: str) -> str:
    nested_zip = os.path.join(submission_root, LOGS_PUBLIC_ZIP)
    nested_dir = os.path.join(submission_root, LOGS_PUBLIC_DIR)

    if os.path.isfile(nested_zip):
        extract_dir = tempfile.mkdtemp(prefix="mera_reason_logs_public_")
        with zipfile.ZipFile(nested_zip, "r") as archive:
            archive.extractall(extract_dir)
        return extract_dir

    if os.path.isdir(nested_dir):
        return nested_dir

    raise ValueError(
        f"Submission must contain `{LOGS_PUBLIC_ZIP}` or `{LOGS_PUBLIC_DIR}/` with evaluation logs."
    )


def _collect_submission_json_files(submission_root: str) -> dict[str, str]:
    files = {}
    for name in os.listdir(submission_root):
        if not name.endswith(".json") or name == LOGS_PUBLIC_ZIP:
            continue
        path = os.path.join(submission_root, name)
        if os.path.isfile(path):
            files[os.path.splitext(name)[0].lower()] = path
    return files


def validate_submission_files(submission_root: str) -> list[str]:
    submission_files = _collect_submission_json_files(submission_root)
    missing = []
    for task in Tasks:
        candidates = _submission_name_candidates(task)
        if not any(candidate in submission_files for candidate in candidates):
            missing.append(task.value.col_name)
    return missing


def _iter_results_files(*roots: str):
    seen = set()
    for root in roots:
        if not os.path.isdir(root):
            continue
        for name in sorted(os.listdir(root)):
            if not name.startswith(RESULTS_PREFIX) or not name.endswith(".json"):
                continue
            path = os.path.join(root, name)
            if path in seen:
                continue
            seen.add(path)
            yield path


def _count_results_files(*roots: str) -> int:
    count = 0
    for root in roots:
        if not os.path.isdir(root):
            continue
        for name in os.listdir(root):
            if name.startswith(RESULTS_PREFIX) and name.endswith(".json"):
                count += 1
    return count


def parse_logs_public_results(logs_public_dir: str, submission_root: str | None = None) -> dict[str, dict[str, float]]:
    merged: dict[str, dict[str, float]] = {}
    roots = [logs_public_dir]
    if submission_root:
        roots.append(submission_root)

    for path in _iter_results_files(*roots):
        with open(path, encoding="utf-8") as fp:
            data = json.load(fp)
        for benchmark, scores in parse_results_data(data).items():
            merged.setdefault(benchmark, {}).update(scores)
    return merged


def _fill_missing_benchmarks(
    parsed_results: dict[str, dict[str, float]],
) -> tuple[dict[str, dict[str, float]], list[str]]:
    """Ensure all benchmarks are present; missing tasks score 0."""
    complete: dict[str, dict[str, float]] = {}
    missing_names: list[str] = []

    for task in Tasks:
        benchmark = task.value.benchmark
        primary = task.value.primary_metric
        raw_scores = parsed_results.get(benchmark, {})

        if raw_scores and primary in raw_scores:
            complete[benchmark] = raw_scores
            continue

        if benchmark not in parsed_results or not raw_scores:
            missing_names.append(task.value.col_name)

        if raw_scores:
            complete[benchmark] = {primary: 0.0, **raw_scores}
        else:
            complete[benchmark] = {primary: 0.0}

    return complete, missing_names


def build_normalized_result_payload(
    results: dict[str, dict[str, float]],
    model_name: str,
    team: str,
) -> dict:
    return {
        "model": model_name,
        "team": team,
        "results": results,
    }


def process_submission_zip(
    zip_path: str,
    model_name: str,
    team: str,
) -> tuple[dict, dict[str, dict[str, float]], list[str], list[str]]:
    extract_dir = extract_submission_archive(zip_path)
    logs_public_dir = None
    try:
        submission_root = _find_submission_root(extract_dir)
        missing_files = validate_submission_files(submission_root)
        logs_public_dir = _resolve_logs_public_dir(submission_root)
        parsed_results = parse_logs_public_results(logs_public_dir, submission_root)

        if not parsed_results:
            roots = [logs_public_dir]
            if submission_root:
                roots.append(submission_root)
            has_results = _count_results_files(*roots) > 0
            if has_results:
                detail = "Found results_*.json files but could not extract leaderboard metrics from them."
            else:
                detail = (
                    "logs_public contains no results_*.json files (only samples are not enough). "
                    "Re-pack the submission with log_to_reasoning_submission.py after evaluation."
                )
            raise ValueError(
                f"No leaderboard task scores found in logs_public. {detail} "
                f"Expected results for: {', '.join(t.value.col_name for t in Tasks)}."
            )

        parsed_results, missing_benchmarks = _fill_missing_benchmarks(parsed_results)

        resolved_model = model_name.strip()
        if not resolved_model:
            raise ValueError("Model name is required.")

        payload = build_normalized_result_payload(parsed_results, resolved_model, team)
        return payload, parsed_results, missing_files, missing_benchmarks
    finally:
        shutil.rmtree(extract_dir, ignore_errors=True)
        if logs_public_dir and logs_public_dir.startswith(tempfile.gettempdir()):
            shutil.rmtree(logs_public_dir, ignore_errors=True)