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.25,
    ("remove", "allow"):  0.25,
}


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 []
    frac_src       = (len(set(checked)) / len(all_src)) if all_src else 1.0

    if frac_src >= 1.0:
        evidence_score = 1.0
    elif frac_src >= 0.60:
        evidence_score = 0.5 + 0.5 * (frac_src - 0.60) / 0.40
    else:
        evidence_score = frac_src / 0.60 * 0.5

    all_clm         = episode.get("all_claim_ids") or []
    verified        = episode.get("claims_verified") or []
    frac_clm        = (len(set(verified)) / len(all_clm)) if all_clm else 1.0
    reasoning_score = min(1.0, 0.6 + 0.4 * frac_clm) if frac_clm >= 0.5 else frac_clm * 0.6 / 0.5

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

    raw = (0.30 * decision_score +
           0.30 * evidence_score +
           0.30 * reasoning_score +
           0.10 * efficiency_score)

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


def grade_episode(episode) -> float:
    try:
        return _compute_score(episode if isinstance(episode, dict) else {})
    except Exception:
        return 0.0