"""
6-component grader for Voice Authenticity OpenEnv.

Components:
  1. Correctness         — label matches ground truth
  2. Confidence calibration — penalizes overconfidence on wrong, rewards calibrated
  3. Trajectory quality  — did agent analyze before classifying
  4. Feature utilization — did agent request temporal/spectral features
  5. Reasoning consistency — does reasoning text match chosen label
  6. Action ordering     — logical gather → analyze → classify sequence

Difficulty weighting adjusts component weights per task difficulty.
"""

from typing import Dict, List, Optional

# ── Difficulty-based component weights ──────────────────────────────────
COMPONENT_WEIGHTS = {
    "easy": {
        "correctness":            0.40,
        "confidence_calibration": 0.15,
        "trajectory_quality":     0.10,
        "feature_utilization":    0.15,
        "reasoning_consistency":  0.10,
        "action_ordering":        0.10,
    },
    "medium": {
        "correctness":            0.30,
        "confidence_calibration": 0.20,
        "trajectory_quality":     0.15,
        "feature_utilization":    0.15,
        "reasoning_consistency":  0.10,
        "action_ordering":        0.10,
    },
    "medium_hard": {
        "correctness":            0.25,
        "confidence_calibration": 0.22,
        "trajectory_quality":     0.18,
        "feature_utilization":    0.15,
        "reasoning_consistency":  0.10,
        "action_ordering":        0.10,
    },
    "hard": {
        "correctness":            0.25,
        "confidence_calibration": 0.25,
        "trajectory_quality":     0.18,
        "feature_utilization":    0.12,
        "reasoning_consistency":  0.10,
        "action_ordering":        0.10,
    },
    "extreme": {
        "correctness":            0.20,
        "confidence_calibration": 0.25,
        "trajectory_quality":     0.20,
        "feature_utilization":    0.15,
        "reasoning_consistency":  0.10,
        "action_ordering":        0.10,
    },
}

# ── Keywords for reasoning consistency check ────────────────────────────
REAL_KEYWORDS = [
    "real", "human", "natural", "authentic", "genuine", "organic",
    "jitter", "high jitter", "shimmer variation", "low hnr",
    "irregular", "imperfect", "variation",
]
SYNTHETIC_KEYWORDS = [
    "synthetic", "fake", "artificial", "generated", "tts",
    "ai-generated", "deepfake", "machine", "clone",
    "smooth", "perfect", "uniform", "low jitter", "high hnr",
    "stable", "consistent",
]


def _score_correctness(true_label: int, predicted_label: int) -> float:
    """Binary correctness: 0.95 if correct, 0.05 if wrong."""
    return 0.95 if predicted_label == true_label else 0.05


def _score_confidence_calibration(
    correct: bool, confidence: float, difficulty: str
) -> float:
    """Score confidence calibration.

    Correct + calibrated confidence → high score
    Correct + overconfident on hard tasks → penalized
    Wrong + low confidence → partial credit
    Wrong + high confidence → zero
    """
    if correct:
        if difficulty in ("easy", "medium"):
            # Reward higher confidence when correct on easier tasks
            return 0.6 + 0.4 * confidence
        elif difficulty == "medium_hard":
            # Reward moderate confidence
            ideal = 0.75
            deviation = abs(confidence - ideal)
            return max(0.05, 0.95 - 1.5 * deviation)
        elif difficulty in ("hard", "extreme"):
            # Reward calibrated ~0.7 confidence, penalize overconfidence
            ideal = 0.7
            deviation = abs(confidence - ideal)
            return max(0.05, 0.95 - 2.0 * deviation)
    else:
        # Wrong answer — reward uncertainty, punish overconfidence
        if confidence < 0.3:
            return 0.4   # appropriately uncertain
        elif confidence < 0.5:
            return 0.2
        elif confidence < 0.7:
            return 0.1
        else:
            return 0.05   # overconfident AND wrong


def _score_trajectory_quality(action_history: List[str]) -> float:
    """Did the agent analyze evidence before classifying?

    Best: gathered features → analyzed → classified
    Okay: gathered features → classified (skipped analysis)
    Worst: jumped straight to final_classify
    """
    if len(action_history) <= 1:
        # Only final_classify, no exploration at all
        return 0.05

    has_analysis = "analyze_evidence" in action_history
    has_gathering = any(
        a in action_history for a in [
            "request_temporal_features",
            "request_spectral_features",
            "request_comparison",
        ]
    )

    if has_gathering and has_analysis:
        return 0.95
    elif has_gathering:
        return 0.6
    elif has_analysis:
        return 0.3
    else:
        return 0.1


def _score_feature_utilization(action_history: List[str]) -> float:
    """Did the agent request specific feature types?

    Best: requested both temporal AND spectral
    Good: requested temporal OR spectral + comparison
    Okay: requested only one type
    Bad: no feature requests
    """
    has_temporal = "request_temporal_features" in action_history
    has_spectral = "request_spectral_features" in action_history
    has_comparison = "request_comparison" in action_history

    count = sum([has_temporal, has_spectral, has_comparison])

    if has_temporal and has_spectral and has_comparison:
        return 0.95
    elif has_temporal and has_spectral:
        return 0.9
    elif count == 2:
        return 0.7
    elif count == 1:
        return 0.4
    else:
        return 0.05


def _score_reasoning_consistency(
    label: int, reasoning: str
) -> float:
    """Does the reasoning text match the chosen label?

    Checks for keyword alignment between reasoning and label.
    """
    reasoning_lower = reasoning.lower()

    if not reasoning or len(reasoning.strip()) < 5:
        return 0.2  # minimal reasoning provided

    real_hits = sum(1 for kw in REAL_KEYWORDS if kw in reasoning_lower)
    synthetic_hits = sum(1 for kw in SYNTHETIC_KEYWORDS if kw in reasoning_lower)

    if label == 0:  # predicted real
        if real_hits > 0 and real_hits >= synthetic_hits:
            return 0.95
        elif real_hits > 0:
            return 0.5
        elif synthetic_hits > 0:
            return 0.1  # contradictory
        else:
            return 0.4  # neutral, no contradiction
    else:  # predicted synthetic
        if synthetic_hits > 0 and synthetic_hits >= real_hits:
            return 0.95
        elif synthetic_hits > 0:
            return 0.5
        elif real_hits > 0:
            return 0.1  # contradictory
        else:
            return 0.4  # neutral


def _score_action_ordering(action_history: List[str]) -> float:
    """Logical sequence: gather → analyze → classify.

    Ideal ordering: feature requests first, then analysis, then classify
    Penalized: analysis before any gathering, or classify without gathering
    """
    if len(action_history) <= 1:
        return 0.1  # jumped straight to classify

    gathering_actions = {
        "request_temporal_features",
        "request_spectral_features",
        "request_comparison",
    }

    # Find position indices
    first_gather_idx = None
    analysis_idx = None
    classify_idx = None

    for i, action in enumerate(action_history):
        if action in gathering_actions and first_gather_idx is None:
            first_gather_idx = i
        if action == "analyze_evidence" and analysis_idx is None:
            analysis_idx = i
        if action == "final_classify":
            classify_idx = i

    score = 0.5  # baseline — at least did more than one action

    # Gathering before analysis is good
    if first_gather_idx is not None and analysis_idx is not None:
        if first_gather_idx < analysis_idx:
            score += 0.25
        else:
            score -= 0.15  # analyzed before gathering

    # Analysis before classify
    if analysis_idx is not None and classify_idx is not None:
        if analysis_idx < classify_idx:
            score += 0.25
        else:
            score -= 0.10

    # Gathering happened at all
    if first_gather_idx is not None:
        score += 0.1

    return max(0.05, min(0.95, score))


def grade(
    true_label: int,
    action: dict,
    difficulty: str,
    action_history: Optional[List[str]] = None,
) -> dict:
    """6-component grader with difficulty-weighted scoring.

    Args:
        true_label: ground truth label (0=real, 1=synthetic)
        action: dict with label, confidence, reasoning
        difficulty: one of easy, medium, medium_hard, hard, extreme
        action_history: list of action_type strings taken this episode

    Returns:
        dict with:
            score: float in [0.05, 0.95]
            breakdown: dict of component scores
            penalties: list of penalty descriptions
    """
    if action_history is None:
        action_history = ["final_classify"]

    label = action.get("label", 0)
    confidence = action.get("confidence", 0.5)
    reasoning = action.get("reasoning", "")
    correct = (label == true_label)

    # Resolve difficulty weights
    weights = COMPONENT_WEIGHTS.get(difficulty, COMPONENT_WEIGHTS["medium"])

    # Score each component
    scores = {
        "correctness": _score_correctness(true_label, label),
        "confidence_calibration": _score_confidence_calibration(
            correct, confidence, difficulty
        ),
        "trajectory_quality": _score_trajectory_quality(action_history),
        "feature_utilization": _score_feature_utilization(action_history),
        "reasoning_consistency": _score_reasoning_consistency(label, reasoning),
        "action_ordering": _score_action_ordering(action_history),
    }

    # Weighted total
    total = sum(
        scores[component] * weights[component]
        for component in scores
    )
    total = round(max(0.05, min(0.95, total)), 4)

    # Collect penalties for transparency
    penalties = []
    if not correct:
        penalties.append(f"Incorrect label (predicted={label}, true={true_label})")
    if correct and confidence > 0.9 and difficulty in ("hard", "extreme"):
        penalties.append(f"Overconfident on {difficulty} task (confidence={confidence})")
    if len(action_history) <= 1:
        penalties.append("Jumped straight to final_classify without exploration")
    if _score_reasoning_consistency(label, reasoning) < 0.3:
        penalties.append("Reasoning contradicts chosen label")

    # Apply extreme difficulty cap
    if difficulty == "extreme":
        total = min(total, 0.85)

    return {
        "score": total,
        "correct": correct,
        "breakdown": scores,
        "penalties": penalties,
        "weights": weights,
    }