"""
tasks/graders.py
================
Deterministic graders for easy / medium / hard task difficulties.

Each grader scores a completed episode on a 0.0–1.0 scale using four
independent sub-scores:

  1. decision_score     (0–0.40) — correctness of the final moderation call
  2. evidence_score     (0–0.30) — how well the agent used available sources
  3. reasoning_score    (0–0.20) — depth of claim verification
  4. efficiency_score   (0–0.10) — did the agent avoid unnecessary steps?

Weights differ per difficulty to reflect what matters most:
  Easy   → decision is paramount; evidence and reasoning are straightforward
  Medium → evidence quality and reasoning depth matter more
  Hard   → all four dimensions matter roughly equally; efficiency penalised less

Usage
-----
    from tasks.graders import GraderFactory

    grader = GraderFactory.get("medium")
    score  = grader.score(
        decision="flag",
        ground_truth="flag",
        sources_checked=["src_001", "src_002", "src_003"],
        claims_verified=["claim_001", "claim_002"],
        step_count=8,
        all_source_ids=["src_001", ..., "src_006"],
        all_claim_ids=["claim_001", "claim_002"],
    )
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import List


# ---------------------------------------------------------------------------
# Partial-credit adjacency for decisions
# ---------------------------------------------------------------------------
# (predicted, ground_truth) → credit (0–1)
_DECISION_CREDIT = {
    # Exact match
    ("allow",  "allow"):  1.0,
    ("flag",   "flag"):   1.0,
    ("remove", "remove"): 1.0,
    # One step away
    ("flag",   "remove"): 0.5,
    ("remove", "flag"):   0.5,
    ("flag",   "allow"):  0.5,
    ("allow",  "flag"):   0.5,
    # Far apart
    ("allow",  "remove"): 0.0,
    ("remove", "allow"):  0.0,
    ("none",   "allow"):  0.0,
    ("none",   "flag"):   0.0,
    ("none",   "remove"): 0.0,
}


# ---------------------------------------------------------------------------
# Base grader
# ---------------------------------------------------------------------------

class BaseGrader(ABC):

    # Subclasses set these weights (must sum to 1.0)
    W_DECISION  : float = 0.40
    W_EVIDENCE  : float = 0.30
    W_REASONING : float = 0.20
    W_EFFICIENCY: float = 0.10

    # Subclasses set these thresholds
    IDEAL_MAX_STEPS : int   = 10
    MIN_SOURCE_FRAC : float = 0.40   # fraction of all sources that should be checked
    MIN_CLAIM_FRAC  : float = 0.50   # fraction of all claims that should be verified

    def grade(self, episode: dict) -> float:
        """
        OpenEnv-standard grader interface.
        Accepts a completed episode dict and returns a score in (0, 1).
        """
        return self.score(
            decision=episode.get("final_decision") or episode.get("decision") or "none",
            ground_truth=episode.get("ground_truth") or episode.get("correct_decision") or "none",
            sources_checked=episode.get("sources_checked") or [],
            claims_verified=episode.get("claims_verified") or [],
            step_count=int(episode.get("step_count") or 0),
            all_source_ids=episode.get("all_source_ids") or [],
            all_claim_ids=episode.get("all_claim_ids") or [],
        )

    def score(
        self,
        decision: str,
        ground_truth: str,
        sources_checked: List[str],
        claims_verified: List[str],
        step_count: int,
        all_source_ids: List[str],
        all_claim_ids: List[str],
    ) -> float:
        """Return a normalised score strictly in (0, 1)."""
        d = self._decision_score(decision, ground_truth)
        e = self._evidence_score(sources_checked, all_source_ids)
        r = self._reasoning_score(claims_verified, all_claim_ids)
        f = self._efficiency_score(step_count)

        raw = (
            self.W_DECISION   * d
            + self.W_EVIDENCE   * e
            + self.W_REASONING  * r
            + self.W_EFFICIENCY * f
        )
        # Clamp strictly inside (0, 1) — validator requires scores never equal 0.0 or 1.0
        return round(max(0.005, min(0.995, raw)), 4)

    def score_breakdown(
        self,
        decision: str,
        ground_truth: str,
        sources_checked: List[str],
        claims_verified: List[str],
        step_count: int,
        all_source_ids: List[str],
        all_claim_ids: List[str],
    ) -> dict:
        """Return full breakdown for audit / dashboards."""
        d = self._decision_score(decision, ground_truth)
        e = self._evidence_score(sources_checked, all_source_ids)
        r = self._reasoning_score(claims_verified, all_claim_ids)
        f = self._efficiency_score(step_count)
        total = self.W_DECISION*d + self.W_EVIDENCE*e + self.W_REASONING*r + self.W_EFFICIENCY*f
        return {
            "decision_sub_score":    round(d, 4),
            "evidence_sub_score":    round(e, 4),
            "reasoning_sub_score":   round(r, 4),
            "efficiency_sub_score":  round(f, 4),
            "weights": {
                "decision":    self.W_DECISION,
                "evidence":    self.W_EVIDENCE,
                "reasoning":   self.W_REASONING,
                "efficiency":  self.W_EFFICIENCY,
            },
            "total": round(max(0.005, min(0.995, total)), 4),
        }

    # ------------------------------------------------------------------
    # Sub-score methods (shared; can be overridden per difficulty)
    # ------------------------------------------------------------------

    def _decision_score(self, decision: str, ground_truth: str) -> float:
        key = (decision.lower(), ground_truth.lower())
        return _DECISION_CREDIT.get(key, 0.0)

    def _evidence_score(
        self, sources_checked: List[str], all_source_ids: List[str]
    ) -> float:
        if not all_source_ids:
            return 1.0
        frac = len(set(sources_checked)) / len(all_source_ids)
        # Reward checking at least MIN_SOURCE_FRAC of available sources
        if frac >= 1.0:
            return 1.0
        if frac >= self.MIN_SOURCE_FRAC:
            # Scale linearly from MIN_SOURCE_FRAC→1 mapping to 0.5→1.0
            slope = 0.5 / (1.0 - self.MIN_SOURCE_FRAC)
            return 0.5 + slope * (frac - self.MIN_SOURCE_FRAC)
        # Below minimum — partial credit only
        return frac / self.MIN_SOURCE_FRAC * 0.5

    def _reasoning_score(
        self, claims_verified: List[str], all_claim_ids: List[str]
    ) -> float:
        if not all_claim_ids:
            return 1.0
        frac = len(set(claims_verified)) / len(all_claim_ids)
        if frac >= self.MIN_CLAIM_FRAC:
            return min(1.0, 0.6 + 0.4 * (frac - self.MIN_CLAIM_FRAC) / (1.0 - self.MIN_CLAIM_FRAC + 1e-9))
        return frac / self.MIN_CLAIM_FRAC * 0.6

    def _efficiency_score(self, step_count: int) -> float:
        if step_count <= 0:
            return 0.0
        if step_count <= self.IDEAL_MAX_STEPS:
            return 1.0
        # Linear decay past ideal; floor at 0
        overage = step_count - self.IDEAL_MAX_STEPS
        return max(0.0, 1.0 - 0.1 * overage)

    @property
    @abstractmethod
    def difficulty(self) -> str: ...


# ---------------------------------------------------------------------------
# Difficulty-specific graders
# ---------------------------------------------------------------------------

class EasyGrader(BaseGrader):
    """
    Easy scenarios have clear-cut answers.
    Decision correctness dominates; efficiency is important.
    """
    difficulty = "easy"
    W_DECISION   = 0.50
    W_EVIDENCE   = 0.25
    W_REASONING  = 0.15
    W_EFFICIENCY = 0.10

    IDEAL_MAX_STEPS = 8
    MIN_SOURCE_FRAC = 0.33   # at least 1/3 of sources
    MIN_CLAIM_FRAC  = 0.50


class MediumGrader(BaseGrader):
    """
    Medium scenarios are partially true / misleading.
    Evidence quality and reasoning depth matter more.
    """
    difficulty = "medium"
    W_DECISION   = 0.35
    W_EVIDENCE   = 0.35
    W_REASONING  = 0.20
    W_EFFICIENCY = 0.10

    IDEAL_MAX_STEPS = 12
    MIN_SOURCE_FRAC = 0.50
    MIN_CLAIM_FRAC  = 0.75

    def _evidence_score(
        self, sources_checked: List[str], all_source_ids: List[str]
    ) -> float:
        """Medium: penalise lightly for checking < 50%, harshly for < 33%."""
        if not all_source_ids:
            return 1.0
        frac = len(set(sources_checked)) / len(all_source_ids)
        if frac >= 0.75:
            return 1.0
        if frac >= 0.50:
            return 0.6 + 0.4 * (frac - 0.50) / 0.25
        if frac >= 0.33:
            return 0.3 + 0.3 * (frac - 0.33) / 0.17
        return frac / 0.33 * 0.3


class HardGrader(BaseGrader):
    """
    Hard scenarios have conflicting evidence and require nuanced reasoning.
    All four dimensions count equally; efficiency is barely penalised.
    """
    difficulty = "hard"
    W_DECISION   = 0.30
    W_EVIDENCE   = 0.30
    W_REASONING  = 0.30
    W_EFFICIENCY = 0.10

    IDEAL_MAX_STEPS = 18
    MIN_SOURCE_FRAC = 0.60
    MIN_CLAIM_FRAC  = 1.00   # all claims must be verified for full credit

    def _decision_score(self, decision: str, ground_truth: str) -> float:
        """Hard: never give 0 for a flag/remove swap — these are genuinely ambiguous."""
        base = super()._decision_score(decision, ground_truth)
        # Give 0.25 minimum for any non-trivially wrong answer
        if base == 0.0 and decision in ("flag", "remove"):
            return 0.25
        return base

    def _efficiency_score(self, step_count: int) -> float:
        """Hard: efficiency barely matters — thorough exploration is rewarded."""
        if step_count <= self.IDEAL_MAX_STEPS:
            return 1.0
        overage = step_count - self.IDEAL_MAX_STEPS
        return max(0.0, 1.0 - 0.05 * overage)


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

class GraderFactory:
    _REGISTRY = {
        "easy":   EasyGrader,
        "medium": MediumGrader,
        "hard":   HardGrader,
    }

    @classmethod
    def get(cls, difficulty: str) -> BaseGrader:
        key = difficulty.lower().strip()
        if key not in cls._REGISTRY:
            raise ValueError(
                f"Unknown difficulty '{difficulty}'. "
                f"Valid options: {list(cls._REGISTRY.keys())}"
            )
        return cls._REGISTRY[key]()

    @classmethod
    def register(cls, difficulty: str, grader_cls: type) -> None:
        """Allow enterprise teams to plug in custom graders."""
        cls._REGISTRY[difficulty.lower()] = grader_cls


# ---------------------------------------------------------------------------
# Top-level grader functions — referenced as tasks.graders:easy_grader etc.
# in openenv.yaml.  Each delegates to the per-task grader module.
# ---------------------------------------------------------------------------

def easy_grader(episode) -> float:
    """OpenEnv entry-point grader for the easy task."""
    try:
        from tasks.easy.grader import grade_episode
        return grade_episode(episode if isinstance(episode, dict) else {})
    except Exception:
        return 0.005


def medium_grader(episode) -> float:
    """OpenEnv entry-point grader for the medium task."""
    try:
        from tasks.medium.grader import grade_episode
        return grade_episode(episode if isinstance(episode, dict) else {})
    except Exception:
        return 0.005


def hard_grader(episode) -> float:
    """OpenEnv entry-point grader for the hard task."""
    try:
        from tasks.hard.grader import grade_episode
        return grade_episode(episode if isinstance(episode, dict) else {})
    except Exception:
        return 0.005