"""tasks/easy/grader.py — OpenEnv-compliant grader for the easy task."""
from __future__ import annotations

_CREDIT = {
    ("allow",  "allow"):  1.0,
    ("flag",   "flag"):   1.0,
    ("remove", "remove"): 1.0,
    ("flag",   "remove"): 0.5,
    ("remove", "flag"):   0.5,
    ("flag",   "allow"):  0.5,
    ("allow",  "flag"):   0.5,
    ("allow",  "remove"): 0.0,
    ("remove", "allow"):  0.0,
}


def _compute_score(episode: dict) -> float:
    decision     = (episode.get("final_decision") or
                    episode.get("decision") or "none").lower().strip()
    ground_truth = (episode.get("ground_truth") or
                    episode.get("correct_decision") or "none").lower().strip()

    decision_score = _CREDIT.get((decision, ground_truth), 0.0)

    all_src        = episode.get("all_source_ids") or []
    checked        = episode.get("sources_checked") or []
    evidence_score = min(1.0, len(set(checked)) / len(all_src)) if all_src else 1.0

    all_clm          = episode.get("all_claim_ids") or []
    verified         = episode.get("claims_verified") or []
    reasoning_score  = min(1.0, len(set(verified)) / len(all_clm)) if all_clm else 1.0

    steps            = int(episode.get("step_count") or 0)
    efficiency_score = 1.0 if steps <= 8 else max(0.0, 1.0 - 0.1 * (steps - 8))

    raw = (0.50 * decision_score +
           0.25 * evidence_score +
           0.15 * reasoning_score +
           0.10 * efficiency_score)

    return round(max(0.005, min(0.995, raw)), 4)


# ✅ THIS is what OpenEnv expects
def grade_episode(episode) -> float:
    try:
        return _compute_score(episode if isinstance(episode, dict) else {})
    except Exception:
        return 0.0