"""Episode grading: turn validation attempts + job outcome into a scalar reward.

Current regime (mock compute): the reward is dominated by the validation
component — passing on the first submission scores 1.0, decaying per extra
attempt. The job component (post-training eval reward of the inner agent) is
computed and recorded but weighted 0 until real GPU results exist; flip the
weights when the compute provider is real.
"""

from __future__ import annotations

from dataclasses import dataclass

from autonomous_trainer.schemas import JobRecord, JobStatus


@dataclass(frozen=True)
class GradeWeights:
    validation: float = 1.0
    job: float = 0.0
    # Token efficiency: fewer total tokens to a successful submission scores
    # higher. Weighted 0 for now; intended as a tie-breaker during eval/RL so the
    # trainer learns to be economical, not just correct.
    efficiency: float = 0.0
    # Training-time pressure: faster training scores higher, gated on job
    # success. Keep this weight SMALL relative to `job` — the final eval reward
    # must dominate; this only nudges the trainer toward leaner configs over time.
    train_speed: float = 0.0


# Weights for live evaluation/RL with real compute: eval-transfer dominates,
# validation efficiency matters, training time exerts slight pressure.
# train_speed is capped LOW deliberately: at 0.10 a fast-but-flat small-model
# job outscored a genuine improvement on the larger base model (scenario
# analysis, pilot-3) — speed must never offset real transfer, only break ties.
LIVE_WEIGHTS = GradeWeights(validation=0.35, job=0.60, efficiency=0.0, train_speed=0.05)


def validation_component(validated: bool, attempts_used: int, max_attempts: int) -> float:
    """1.0 for a first-try pass, decaying linearly per failed attempt; 0.0 if never validated."""
    if not validated or attempts_used < 1:
        return 0.0
    return max(0.0, 1.0 - (attempts_used - 1) / max_attempts)


def efficiency_component(validated: bool, total_tokens: int, token_budget: int) -> float:
    """1.0 at zero tokens, linearly to 0.0 at the episode token budget.

    Only a *successful* episode earns efficiency credit — otherwise the cheapest
    strategy would be doing nothing.
    """
    if not validated or token_budget <= 0:
        return 0.0
    return max(0.0, min(1.0, 1.0 - total_tokens / token_budget))


def job_component(job: JobRecord | None, best_pre: float | None = None) -> float:
    """Value the trained model against the best UNTRAINED option for the task.

    job = 0.25*post + 0.75*clip01(0.5 + (post - best_pre) / (2*max(1-best_pre, 0.3)))

    The dominant term is signed uplift over `best_pre` (the best pre-training
    score across the task's allowed base models — harness-measured once, then
    frozen): coasting on the strongest model's floor lands at ~0.5 of that
    term, beating the floor climbs toward 1.0, degrading sinks below 0.5 — so
    failure < harmful < coast < genuine gain is monotone with real contrast at
    every step. Using best_pre (not the chosen model's own pre) kills weak-base
    headroom farming; the denominator floor of 0.3 caps eval-noise
    amplification; the 0.25 absolute term keeps reward anchored to shipped
    quality and keeps hard families faintly live. Failed/absent jobs score 0.

    best_pre=None (unmeasured task) deliberately falls back to plain
    clip01(post) — run scripts/seed_baselines.py before RL so this path never
    fires during training.
    """
    if job is None or job.result is None or job.result.status is not JobStatus.SUCCEEDED:
        return 0.0
    post = max(0.0, min(1.0, job.result.metrics.get("eval_reward_post", 0.0)))
    if best_pre is None:
        return post
    uplift = 0.5 + (post - best_pre) / (2 * max(1.0 - best_pre, 0.3))
    return 0.25 * post + 0.75 * max(0.0, min(1.0, uplift))


def train_speed_component(job: JobRecord | None, train_time_budget_s: float) -> float:
    """1.0 for instant training, linearly to 0.0 at the budget. Success-gated:
    a failed job earns no speed credit (fast failures must not pay)."""
    if (
        job is None
        or job.result is None
        or job.result.status is not JobStatus.SUCCEEDED
        or train_time_budget_s <= 0
    ):
        return 0.0
    t_train = job.result.metrics.get("t_train_s")
    if t_train is None:
        return 0.0
    return max(0.0, min(1.0, 1.0 - t_train / train_time_budget_s))


def grade_episode(
    validated: bool,
    attempts_used: int,
    max_attempts: int,
    job: JobRecord | None,
    weights: GradeWeights = GradeWeights(),
    total_tokens: int = 0,
    token_budget: int = 0,
    train_time_budget_s: float = 0.0,
    best_pre: float | None = None,
) -> tuple[float, dict[str, float]]:
    """Return (reward, components). Components are always recorded, whatever the weights."""
    vc = validation_component(validated, attempts_used, max_attempts)
    jc = job_component(job, best_pre)
    ec = efficiency_component(validated, total_tokens, token_budget)
    tc = train_speed_component(job, train_time_budget_s)
    total_weight = weights.validation + weights.job + weights.efficiency + weights.train_speed
    reward = (
        (
            weights.validation * vc
            + weights.job * jc
            + weights.efficiency * ec
            + weights.train_speed * tc
        )
        / total_weight
        if total_weight
        else 0.0
    )
    t_train = (
        job.result.metrics.get("t_train_s") if job and job.result else None
    )
    post = (
        job.result.metrics.get("eval_reward_post") if job and job.result else None
    )
    return reward, {
        "validation": vc,
        "job": jc,
        "efficiency": ec,
        "train_speed": tc,
        "total_tokens": float(total_tokens),
        "t_train_s": float(t_train) if t_train is not None else -1.0,
        # Debuggability: the raw inputs behind the job component (-1 = absent).
        "eval_reward_post": float(post) if post is not None else -1.0,
        "best_pre": float(best_pre) if best_pre is not None else -1.0,
        "uplift": float(post - best_pre) if post is not None and best_pre is not None else 0.0,
        "weight_validation": weights.validation,
        "weight_job": weights.job,
        "weight_efficiency": weights.efficiency,
        "weight_train_speed": weights.train_speed,
    }