inference.py

#!/usr/bin/env python3
"""
AdaptiveTutor-Env Baseline Inference Script.

Runs a language-model agent (via OpenAI-compatible API) against all three
AdaptiveTutor-Env tasks and emits structured stdout logs.

Logging format (STRICT — required for evaluation):
  [START] task=<task_name> env=<benchmark> model=<model_name>
  [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>

Environment variables
---------------------
  API_BASE_URL   – Base URL for the LLM API endpoint (default: https://router.huggingface.co/v1).
  MODEL_NAME     – Model identifier for completions (default: Qwen/Qwen2.5-72B-Instruct).
  HF_TOKEN       – Hugging Face API token (REQUIRED).
  API_KEY        – Alternative API key (falls back to HF_TOKEN).
  ENV_BASE_URL   – Base URL for the AdaptiveTutor-Env server (default: http://localhost:7860).

Usage
-----
  python inference.py

The script will run all three tasks sequentially and print a summary table.
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from typing import Any, Dict, List, Optional

import httpx
from openai import OpenAI


# ─── Configuration ────────────────────────────────────────────────────────────

IMAGE_NAME = os.getenv("IMAGE_NAME")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")

API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"

# Validate API_KEY is set (required by hackathon rules)
if API_KEY is None:
    raise ValueError("HF_TOKEN or API_KEY environment variable is required")

ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860")
BENCHMARK = os.getenv("ADAPTIVE_TUTOR_BENCHMARK", "adaptive-tutor-env")

# Valid subjects and topics for validation
VALID_SUBJECTS = {
    "mathematics",
    "science",
    "computer_science",
    "english",
    "social_studies",
}
VALID_TOPICS = {
    "mathematics": {"algebra", "geometry", "calculus", "statistics", "number_theory"},
    "science": {"physics", "chemistry", "biology", "earth_science", "astronomy"},
    "computer_science": {
        "algorithms",
        "data_structures",
        "programming",
        "networking",
        "databases",
    },
    "english": {
        "grammar",
        "vocabulary",
        "reading_comprehension",
        "writing",
        "literature",
    },
    "social_studies": {"history", "geography", "civics", "economics", "culture"},
}
VALID_ACTIVITY_TYPES = {0, 1, 2, 3}
VALID_DIFFICULTIES = {0, 1, 2}
VALID_STRATEGIES = {0, 1, 2}

TASKS: List[str] = [
    "single_subject_mastery",
    "multi_subject_balancing",
    "long_horizon_retention",
]

# Fixed seed for reproducibility
SEED = 42


# ─── Environment Client ────────────────────────────────────────────────────────


class TutorEnvClient:
    """Thin HTTP client wrapping the AdaptiveTutor-Env REST API."""

    def __init__(self, base_url: str = ENV_BASE_URL) -> None:
        self.base_url = base_url.rstrip("/")
        self.http = httpx.AsyncClient(timeout=30.0)

    async def reset(self, task_id: str, seed: int = SEED) -> Dict[str, Any]:
        r = await self.http.post(
            f"{self.base_url}/reset",
            json={"task_id": task_id, "seed": seed},
        )
        r.raise_for_status()
        return r.json()

    async def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
        r = await self.http.post(
            f"{self.base_url}/step",
            json={"action": action},
        )
        r.raise_for_status()
        return r.json()

    async def state(self) -> Dict[str, Any]:
        r = await self.http.get(f"{self.base_url}/state")
        r.raise_for_status()
        return r.json()

    async def grade(self) -> Dict[str, Any]:
        r = await self.http.post(f"{self.base_url}/grade")
        r.raise_for_status()
        return r.json()

    async def close(self) -> None:
        await self.http.aclose()


# ─── LLM Agent ────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are an expert AI tutoring agent. Your goal is to maximize the student's learning score.

## TASK-SPECIFIC GOALS

1. SINGLE_SUBJECT_MASTERY (20 steps max):
   - Goal: Maximize mathematics mastery
   - Only teach mathematics subjects
   - Score = mean(mathematics topic masteries) [range 0-1]
   - Target: Achieve 0.60+ average mastery

2. MULTI_SUBJECT_BALANCING (30 steps max):
   - Goal: Balance mastery across ALL 5 subjects
   - Teach all subjects fairly - avoid neglecting any
   - Score = 0.6×overall_mastery + 0.4×balance_score [range 0-1]
   - Target: 0.50+ with good balance

3. LONG_HORIZON_RETENTION (40 steps max):
   - Goal: Achieve mastery AND retain it (prevent forgetting)
   - Revisit topics regularly - don't let them decay
   - Score = 0.4×mastery + 0.2×min_subject + 0.4×retention [range 0-1]
   - Retention = topics kept at ≥60% of their peak mastery
   - Target: 0.45+ with high retention

## ACTION OUTPUT (JSON only)

{
  "subject": "mathematics|science|computer_science|english|social_studies",
  "topic": "<topic name>",
  "activity_type": 0|1|2|3,
  "difficulty": 0|1|2,
  "strategy": 0|1|2
}

Activity types:
- 0 (video_lesson): Low fatigue, good for introductions, low engagement
- 1 (practice_exercise): High engagement, best for building mastery
- 2 (quiz): Good reinforcement, medium engagement
- 3 (revision): Lowest fatigue, great for recovery

## ZPD (Zone of Proximal Development) - CRITICAL

Set difficulty so student succeeds ~60% of the time:
- mastery < 0.25: use difficulty=0 (easy) → ~80% success
- mastery 0.25-0.45: use difficulty=0 (easy) → ~70% success  
- mastery 0.45-0.60: use difficulty=1 (medium) → ~50% success
- mastery > 0.60: use difficulty=2 (hard) → ~30% success

## STRATEGY SELECTION

- strategy=0 (introduce): For NEW topics or mastery < 0.25
- strategy=1 (reinforce): For topics with mastery 0.25-0.70
- strategy=2 (spaced_repetition): CRITICAL for retention! Use when time_since_visit >= 3

## ENGAGEMENT & FATIGUE MANAGEMENT

Student state affects learning effectiveness:
- engagement < 0.4: Use activity_type=1 (practice) to boost engagement
- fatigue > 0.7: Use activity_type=3 (revision) or 0 (video) to recover
- engagement > 0.6 AND fatigue < 0.5: Good time for challenging activities

## TOPIC SELECTION PRIORITY

For SINGLE_SUBJECT_MASTERY: Always pick lowest mastery mathematics topic
For MULTI_SUBJECT_BALANCING: Pick lowest overall mastery topic (consider all subjects)
For LONG_HORIZON_RETENTION: Pick most overdue topic (highest time_since_visit)

## EXAMPLES

Good action for SINGLE_SUBJECT with mastery=0.3:
{"subject":"mathematics","topic":"algebra","activity_type":1,"difficulty":0,"strategy":1}

Good action for LONG_HORIZON with time_since_visit=10:
{"subject":"science","topic":"physics","activity_type":2,"difficulty":1,"strategy":2}

OUTPUT JSON ONLY. No markdown. No explanation."""


def build_user_prompt(obs: Dict[str, Any], task_id: str, step: int) -> str:
    """Build the per-step prompt for the LLM agent with task-specific guidance."""
    subj_means = obs.get("subject_masteries", {})
    overall = obs.get("overall_mastery", 0.0)
    engagement = obs.get("engagement", 0.0)
    fatigue = obs.get("fatigue", 0.0)
    max_steps = obs.get("max_steps", 20)

    # Task-specific instructions
    task_hints = {
        "single_subject_mastery": "FOCUS ON MATHEMATICS ONLY. Pick lowest mastery math topic.",
        "multi_subject_balancing": "BALANCE ALL SUBJECTS. Avoid letting any subject fall behind.",
        "long_horizon_retention": "SPACED REPETITION IS KEY. Pick topics with highest time_since_visit.",
    }
    task_hint = task_hints.get(task_id, "")

    # Find lowest mastery topic
    lowest: List[tuple] = []
    for subj, topics in obs.get("masteries", {}).items():
        for topic, mastery in topics.items():
            lowest.append((mastery, subj, topic))
    lowest.sort()

    # Find longest unvisited topic
    unvisited: List[tuple] = []
    for subj, topics in obs.get("time_since_last_visit", {}).items():
        for topic, t in topics.items():
            if t > 0:
                unvisited.append((t, subj, topic))
    unvisited.sort(reverse=True)

    # Find lowest subject
    lowest_subject = (
        min(subj_means.items(), key=lambda x: x[1])
        if subj_means
        else ("mathematics", 0.0)
    )
    highest_tsv_subject = unvisited[0] if unvisited else (0, "mathematics", "algebra")

    # Engagement/Fatigue warning
    engagement_warning = ""
    if engagement < 0.4:
        engagement_warning = "⚠️ LOW ENGAGEMENT - Use practice_exercise (1)!"
    elif fatigue > 0.7:
        engagement_warning = "⚠️ HIGH FATIGUE - Use revision (3) or video (0)!"
    elif engagement > 0.6 and fatigue < 0.5:
        engagement_warning = "✅ Good state for challenging activities"

    # Progress indicator
    progress = step / max_steps if max_steps > 0 else 0
    if progress < 0.3:
        time_hint = "Early phase - focus on building foundation"
    elif progress < 0.7:
        time_hint = "Mid phase - balance building and reviewing"
    else:
        time_hint = "Late phase - prioritize retention and weak areas"

    prompt = f"""## {task_id.upper().replace("_", " ")}

**Task Hint:** {task_hint}
**Progress:** Step {step}/{max_steps} ({progress * 100:.0f}%) - {time_hint}

**Student State:**
- Overall mastery: {overall:.3f}
- Engagement: {engagement:.3f} {"⚠️ LOW" if engagement < 0.4 else ""}
- Fatigue: {fatigue:.3f} {"⚠️ HIGH" if fatigue > 0.7 else ""}
{engagement_warning}

**Subject Masteries:**
{json.dumps({k: round(v, 3) for k, v in subj_means.items()}, indent=2)}

**Recommended Focus (lowest mastery):**
{chr(10).join(f"  - {subj}.{topic}: {m:.3f}" for m, subj, topic in lowest[:5])}

**Overdue Topics (prioritize for retention):**
{chr(10).join(f"  - {subj}.{topic}: {t} steps ago" for t, subj, topic in unvisited[:5])}

**Choose the best action based on the task goal.**"""
    return prompt


def validate_action(action: Dict[str, Any], task_id: str) -> tuple[bool, str]:
    """
    Validate that an action has all required fields with valid values.
    Returns (is_valid, error_message).
    """
    required_fields = ["subject", "topic", "activity_type", "difficulty", "strategy"]

    # Check required fields exist
    for field in required_fields:
        if field not in action:
            return False, f"Missing required field: {field}"

    # Validate subject
    subject = action.get("subject", "")
    if subject not in VALID_SUBJECTS:
        return False, f"Invalid subject: '{subject}'. Must be one of: {VALID_SUBJECTS}"

    # Validate topic
    topic = action.get("topic", "")
    valid_topics = VALID_TOPICS.get(subject, set())
    if topic not in valid_topics:
        return (
            False,
            f"Invalid topic: '{topic}' for subject '{subject}'. Must be one of: {valid_topics}",
        )

    # Validate numeric fields
    try:
        activity_type = int(action.get("activity_type", -1))
        difficulty = int(action.get("difficulty", -1))
        strategy = int(action.get("strategy", -1))
    except (ValueError, TypeError):
        return False, "activity_type, difficulty, strategy must be numeric"

    if activity_type not in VALID_ACTIVITY_TYPES:
        return False, f"Invalid activity_type: {activity_type}. Must be 0-3"

    if difficulty not in VALID_DIFFICULTIES:
        return False, f"Invalid difficulty: {difficulty}. Must be 0-2"

    if strategy not in VALID_STRATEGIES:
        return False, f"Invalid strategy: {strategy}. Must be 0-2"

    # Task-specific validation
    if task_id == "single_subject_mastery" and subject != "mathematics":
        return False, f"For single_subject_mastery, must use 'mathematics' subject"

    return True, ""


async def get_llm_action(
    client: OpenAI,
    obs: Dict[str, Any],
    task_id: str,
    step: int,
) -> tuple[Dict[str, Any], bool]:
    """
    Query the LLM for the next action with retry logic.
    Falls back to heuristic after max retries.
    Returns (action, was_llm_used).
    """
    max_retries = 3
    last_error = ""

    for attempt in range(max_retries):
        try:
            # Call LLM
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": build_user_prompt(obs, task_id, step)},
                ],
                max_tokens=300,
                temperature=0.3,
            )
            text = response.choices[0].message.content or ""

            # Strip any accidental markdown fences
            text = text.strip().strip("```json").strip("```").strip()

            # Try to extract JSON if there's extra text
            if "{" in text and "}" in text:
                start = text.find("{")
                end = text.rfind("}") + 1
                text = text[start:end]

            action = json.loads(text)

            # Validate action fields
            is_valid, error_msg = validate_action(action, task_id)
            if not is_valid:
                last_error = f"Action validation failed: {error_msg}"
                if attempt < max_retries - 1:
                    continue  # Retry
                else:
                    raise ValueError(last_error)

            # Convert numeric fields to int
            action["activity_type"] = int(action["activity_type"])
            action["difficulty"] = int(action["difficulty"])
            action["strategy"] = int(action["strategy"])

            # Success!
            if attempt > 0:
                print(
                    f"[INFO] LLM action valid after {attempt + 1} attempts", flush=True
                )
            return action, True

        except json.JSONDecodeError as exc:
            last_error = f"JSON parse error: {exc}"
            if attempt < max_retries - 1:
                print(
                    f"[WARN] JSON parse failed (attempt {attempt + 1}/{max_retries}): {last_error}",
                    flush=True,
                )
                continue
        except Exception as exc:
            last_error = f"LLM error: {exc}"
            if attempt < max_retries - 1:
                print(
                    f"[WARN] LLM call failed (attempt {attempt + 1}/{max_retries}): {last_error}",
                    flush=True,
                )
                await asyncio.sleep(0.5)  # Brief pause before retry
                continue

    # All retries exhausted - use heuristic fallback
    print(
        f"[WARN] LLM failed after {max_retries} attempts: {last_error}. Using heuristic fallback.",
        flush=True,
    )
    return _heuristic_action(obs), False


def _heuristic_action(obs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Enhanced heuristic agent with smart strategies:
    1. Smart Topic Selection: Prioritizes both low mastery AND high decay
    2. Engagement-Aware: Adjusts activity based on fatigue
    3. Spaced Repetition: Prioritizes overdue topics
    4. ZPD-Calibrated: Matches difficulty to mastery level
    """
    masteries = obs.get("masteries", {})
    time_since_visit = obs.get("time_since_last_visit", {})
    engagement = obs.get("engagement", 0.7)
    fatigue = obs.get("fatigue", 0.0)
    task_id = obs.get("task_id", "single_subject_mastery")
    current_step = obs.get("step", 0)
    max_steps = obs.get("max_steps", 20)

    # Calculate decay factor for each topic
    # Topics with high time_since_visit need priority (spaced repetition)
    topic_scores = []
    for subj, topics in masteries.items():
        for topic, mastery in topics.items():
            tsv = time_since_visit.get(subj, {}).get(topic, 0)

            # Different priority calculation based on task
            if task_id == "long_horizon_retention":
                # For retention task: prioritize based on steps remaining and decay
                # Give more weight to topics that haven't been visited recently
                decay_urgency = min(tsv / 8.0, 1.5)  # Stronger spaced repetition
                # Balance mastery building with retention
                priority = decay_urgency + (1.0 - mastery) * 0.2
            elif task_id == "single_subject_mastery":
                # Focus on mastery only
                priority = 1.0 - mastery
            else:
                # Multi-subject balancing: balance mastery and decay
                decay_factor = min(tsv / 10.0, 1.0)
                priority = (1.0 - mastery) * 0.6 + decay_factor * 0.4

            topic_scores.append(
                {
                    "subject": subj,
                    "topic": topic,
                    "mastery": mastery,
                    "time_since_visit": tsv,
                    "priority": priority,
                }
            )

    # Sort by priority (highest first)
    topic_scores.sort(key=lambda x: x["priority"], reverse=True)

    # For single_subject_mastery, focus only on mathematics
    if task_id == "single_subject_mastery":
        topic_scores = [t for t in topic_scores if t["subject"] == "mathematics"]

    # Select best topic
    best = (
        topic_scores[0]
        if topic_scores
        else {
            "subject": "mathematics",
            "topic": "algebra",
            "mastery": 0.2,
            "time_since_visit": 0,
        }
    )
    best_subj = best["subject"]
    best_topic = best["topic"]
    m = best["mastery"]
    tsv = best["time_since_visit"]

    # ZPD-Calibrated Difficulty: match difficulty to mastery for optimal learning
    if m < 0.25:
        difficulty = 0  # Easy - high success rate builds confidence
    elif m < 0.45:
        difficulty = 0  # Still easy for mid-low mastery
    elif m < 0.60:
        difficulty = 1  # Medium - ZPD sweet spot
    else:
        difficulty = 2  # Hard - push toward mastery ceiling

    # Spaced Repetition Strategy - more aggressive for retention task
    if task_id == "long_horizon_retention":
        if tsv >= 3:
            strategy = 2  # Spaced repetition - very important
        elif m < 0.3:
            strategy = 0  # Introduce
        else:
            strategy = 1  # Reinforce
    elif tsv >= 8:
        strategy = 2
    elif tsv >= 5:
        strategy = 2
    elif m < 0.25:
        strategy = 0
    elif m < 0.60:
        strategy = 1
    else:
        strategy = 2

    # Engagement-Aware Activity Selection
    if fatigue > 0.7:
        activity_type = 3  # revision - lowest fatigue
    elif fatigue > 0.5:
        activity_type = 2 if m > 0.3 else 0
    elif engagement < 0.4:
        activity_type = 1
    else:
        activity_type = 1  # practice - best learning gain

    if m < 0.15:
        activity_type = 0

    return {
        "subject": best_subj,
        "topic": best_topic,
        "activity_type": activity_type,
        "difficulty": difficulty,
        "strategy": strategy,
    }


# ─── Episode Runner ───────────────────────────────────────────────────────────


async def run_episode(
    env_client: TutorEnvClient,
    llm_client: OpenAI,
    task_id: str,
) -> Dict[str, Any]:
    """
    Run a single episode for the given task.
    Emits [START], [STEP], and [END] log lines to stdout.
    Tracks LLM vs Heuristic usage for diagnostics.
    """
    # ── [START] ──────────────────────────────────────────────────────────────
    # Format: [START] task=<task_name> env=<benchmark> model=<model_name>
    print(
        f"[START] task={task_id} env={BENCHMARK} model={MODEL_NAME}",
        flush=True,
    )

    obs = await env_client.reset(task_id=task_id, seed=SEED)
    episode_done = False
    steps_taken = 0
    episode_rewards: List[float] = []
    llm_used = 0
    heuristic_used = 0

    while not episode_done:
        step_num = obs.get("step", steps_taken)
        action, was_llm = await get_llm_action(llm_client, obs, task_id, step_num + 1)

        # Track LLM vs Heuristic usage
        if was_llm:
            llm_used += 1
        else:
            heuristic_used += 1

        result = await env_client.step(action)
        obs = result["observation"]
        reward_info = result["reward"]
        episode_done = result["done"]
        info = result.get("info", {})

        step_reward = reward_info.get("total", 0.0)
        episode_rewards.append(step_reward)
        steps_taken += 1

        # ── [STEP] ────────────────────────────────────────────────────────
        # Format: [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
        action_str = f"{action.get('subject')}.{action.get('topic')}:{action.get('activity_type')}:{action.get('difficulty')}:{action.get('strategy')}"
        done_str = str(episode_done).lower()

        print(
            f"[STEP] step={steps_taken} "
            f"action={action_str} "
            f"reward={step_reward:.2f} "
            f"done={done_str} "
            f"error=null",
            flush=True,
        )

    # Final grading
    grade_result = await env_client.grade()
    score = grade_result.get("score", 0.0)
    passed = grade_result.get("passed", False)

    # ── [END] ───────────────────────────────────────────────────────────────
    # Format: [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
    success_str = str(passed).lower()
    rewards_str = ",".join(f"{r:.2f}" for r in episode_rewards)

    print(
        f"[END] success={success_str} steps={steps_taken} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )

    return {
        "task_id": task_id,
        "score": score,
        "passed": passed,
        "steps": steps_taken,
        "episode_return": sum(episode_rewards),
        "llm_used": llm_used,
        "heuristic_used": heuristic_used,
    }


# ─── Main ─────────────────────────────────────────────────────────────────────


async def main() -> None:
    print("=" * 72, flush=True)
    print("AdaptiveTutor-Env — Baseline Inference Script", flush=True)
    print(f"Model:   {MODEL_NAME}", flush=True)
    print(f"API URL: {API_BASE_URL}", flush=True)
    print(f"Env URL: {ENV_BASE_URL}", flush=True)
    print("=" * 72, flush=True)

    # Wait for server to be ready (in Docker / HF Space startup)
    env_client = TutorEnvClient(base_url=ENV_BASE_URL)
    for attempt in range(30):
        try:
            await env_client.http.get(f"{ENV_BASE_URL}/health", timeout=5.0)
            print(
                f"Environment server ready after {attempt + 1} attempt(s).", flush=True
            )
            break
        except Exception:
            print(
                f"Waiting for environment server... (attempt {attempt + 1}/30)",
                flush=True,
            )
            await asyncio.sleep(2)
    else:
        print("ERROR: Environment server did not start. Exiting.", flush=True)
        sys.exit(1)

    # Initialise OpenAI client (matching sample format from requirements)
    llm_client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    results: List[Dict[str, Any]] = []
    for task_id in TASKS:
        print(f"\n{'─' * 72}", flush=True)
        result = await run_episode(env_client, llm_client, task_id)
        results.append(result)
        await asyncio.sleep(1)  # brief pause between episodes

    # ── Summary Table ─────────────────────────────────────────────────────────
    print(f"\n{'=' * 72}", flush=True)
    print("SUMMARY", flush=True)
    print(f"{'=' * 72}", flush=True)
    print(
        f"{'Task':<35} {'Score':>8} {'Passed':>8} {'Steps':>6} {'LLM':>4} {'Heuristic':>10}",
        flush=True,
    )
    print(f"{'-' * 72}", flush=True)
    for r in results:
        print(
            f"{r['task_id']:<35} "
            f"{r['score']:>8.4f} "
            f"{str(r['passed']):>8} "
            f"{r['steps']:>6} "
            f"{r.get('llm_used', 0):>4} "
            f"{r.get('heuristic_used', 0):>10}",
            flush=True,
        )
    print(f"{'=' * 72}", flush=True)

    mean_score = sum(r["score"] for r in results) / len(results)
    print(f"\nMean score across tasks: {mean_score:.4f}", flush=True)

    await env_client.close()


if __name__ == "__main__":
    asyncio.run(main())