curriculum_manager.py

# curriculum_manager.py
# Manages subject/lesson rotation to prevent topic stagnation.
#
# Key functions:
#   normalize_subject(name) -> str   Normalize subject name for comparison
#   record_subject(subject)          Record chosen subject into rotation state
#   record_lesson(subject, lesson)   Record chosen lesson into rotation state
#   get_blacklists() -> dict         Get recent subjects/lessons for prompt injection
#
# Note: Subject generation is now handled by subject_generator.py
# This module only handles rotation state persistence and blacklists.

import os
import json
import re
from datetime import datetime
from typing import Optional
from openai import OpenAI

from llm_client import create_smart_client
client = create_smart_client()

from config import (
    MEMORY_DIR,
    ROTATION_STATE_FILE,
    SUBJECT_COOLDOWN,
    LESSON_COOLDOWN,
    USED_LESSONS_FILE,
    TASK_AGENT_MODEL,
)
from dedup_utils import fuzzy_match as dedup_fuzzy_match, SEMANTIC_CHECK_MAX_ITEMS
from file_lock_utils import get_lock

# Defaults (from config, with fallbacks)
DEFAULT_SUBJECT_COOLDOWN = SUBJECT_COOLDOWN
DEFAULT_LESSON_COOLDOWN = LESSON_COOLDOWN


def normalize_subject(name: str) -> str:
    """
    Normalize a subject name for consistent comparison.
    
    - lowercase
    - strip punctuation (except hyphens and ampersands)
    - collapse multiple spaces
    - strip whitespace
    
    This helps prevent "Math" vs "Mathematics" or "Math & Science" vs "Math and Science"
    from being treated as different when they're essentially the same.
    """
    if not name:
        return ""
    
    # Lowercase
    normalized = name.lower()
    
    # Replace common variations
    normalized = normalized.replace(" and ", " & ")
    
    # Remove punctuation except hyphens, ampersands, and spaces
    normalized = re.sub(r'[^\w\s&-]', '', normalized)
    
    # Collapse multiple spaces
    normalized = re.sub(r'\s+', ' ', normalized)
    
    # Strip whitespace
    normalized = normalized.strip()
    
    return normalized


def normalize_lesson_key(subject: str, lesson: str) -> str:
    """
    Normalize a lesson key for consistent storage and comparison.
    
    Format: "Subject :: Lesson Title"
    Normalization: lowercase, trim, collapse spaces, keep :: and —
    """
    key = f"{subject} :: {lesson}"
    # Lowercase
    key = key.lower()
    # Trim whitespace
    key = key.strip()
    # Collapse multiple spaces
    key = re.sub(r'\s+', ' ', key)
    return key


def load_rotation_state() -> dict:
    """Load the rotation state from file."""
    if not os.path.exists(ROTATION_STATE_FILE):
        return _default_rotation_state()
    
    try:
        with open(ROTATION_STATE_FILE, "r") as f:
            state = json.load(f)
        # Ensure all expected keys exist
        default = _default_rotation_state()
        for key in default:
            if key not in state:
                state[key] = default[key]
        return state
    except Exception as e:
        print(f"[Curriculum] Error loading rotation state: {e}")
        return _default_rotation_state()


def save_rotation_state(state: dict):
    """Save the rotation state to file."""
    os.makedirs(os.path.dirname(ROTATION_STATE_FILE), exist_ok=True)
    with open(ROTATION_STATE_FILE, "w") as f:
        json.dump(state, f, indent=2)


def _default_rotation_state() -> dict:
    """Return default rotation state."""
    return {
        "recent_subjects": [],
        "recent_subjects_normalized": [],  # Normalized versions for comparison
        "recent_lessons": [],
        "subject_cooldown": DEFAULT_SUBJECT_COOLDOWN,
        "lesson_cooldown": DEFAULT_LESSON_COOLDOWN,
        "last_subject": None,
        "last_lesson": None,
        "subject_counts": {},
        "lesson_counts": {},
        "subject_scores": {},
    }


def get_recent_subjects_normalized() -> set:
    """
    Get the set of normalized recent subjects for blacklist comparison.
    
    Returns:
        Set of normalized subject name strings
    """
    state = load_rotation_state()
    recent = state.get("recent_subjects", [])
    return set(normalize_subject(s) for s in recent)


def record_subject(subject: str):
    """
    Record a chosen subject into the rotation state. Process-safe via file lock.
    
    This should be called after subject_of_the_day is selected.
    
    Args:
        subject: The subject name (display version)
    """
    with get_lock(ROTATION_STATE_FILE):
        state = load_rotation_state()
        cooldown = state.get("subject_cooldown", DEFAULT_SUBJECT_COOLDOWN)
        
        # Add to recent subjects (display version)
        state["recent_subjects"].append(subject)
        state["recent_subjects"] = state["recent_subjects"][-cooldown:]
        
        # Also track normalized version for efficient comparison
        normalized = normalize_subject(subject)
        if "recent_subjects_normalized" not in state:
            state["recent_subjects_normalized"] = []
        state["recent_subjects_normalized"].append(normalized)
        state["recent_subjects_normalized"] = state["recent_subjects_normalized"][-cooldown:]
        
        state["last_subject"] = subject
        
        # Update counts
        counts = state.get("subject_counts", {})
        counts[subject] = counts.get(subject, 0) + 1
        state["subject_counts"] = counts
        
        save_rotation_state(state)


def record_lesson(subject: str, lesson: str, score: float = None):
    """
    Record a chosen lesson into the rotation state. Process-safe via file lock.
    
    Args:
        subject: The subject name
        lesson: The lesson title (as chosen by Sophie)
        score: Optional session average score for this lesson
    """
    with get_lock(ROTATION_STATE_FILE):
        state = load_rotation_state()
        cooldown = state.get("lesson_cooldown", DEFAULT_LESSON_COOLDOWN)
        
        # Normalize and create lesson key
        lesson_key = normalize_lesson_key(subject, lesson)
        
        # Update recent lessons
        state["recent_lessons"].append(lesson_key)
        state["recent_lessons"] = state["recent_lessons"][-cooldown:]
        state["last_lesson"] = lesson_key
        
        # Update counts
        counts = state.get("lesson_counts", {})
        counts[lesson_key] = counts.get(lesson_key, 0) + 1
        state["lesson_counts"] = counts
        
        # Update scores if provided
        if score is not None:
            scores = state.get("subject_scores", {})
            if subject not in scores:
                scores[subject] = {"total": 0, "count": 0}
            scores[subject]["total"] += score
            scores[subject]["count"] += 1
            state["subject_scores"] = scores
        
        save_rotation_state(state)


def get_blacklists(n_subjects: int = 25, n_lessons: int = 50) -> dict:
    """
    Get recent subjects and lessons for prompt injection.
    
    Args:
        n_subjects: Number of recent subjects to return (default 25 for subject generator)
        n_lessons: Number of recent lessons to return (default 50 for lesson picker)
    
    Returns:
        dict with keys: recent_subjects, recent_lessons, last_subject, last_lesson
    """
    state = load_rotation_state()
    
    recent_subjects = state.get("recent_subjects", [])[-n_subjects:]
    recent_lessons = state.get("recent_lessons", [])[-n_lessons:]
    
    return {
        "recent_subjects": recent_subjects,
        "recent_lessons": recent_lessons,
        "last_subject": state.get("last_subject"),
        "last_lesson": state.get("last_lesson"),
    }


def get_subject_stats() -> dict:
    """Get statistics about subject/lesson usage."""
    state = load_rotation_state()
    
    return {
        "subject_counts": state.get("subject_counts", {}),
        "lesson_counts": state.get("lesson_counts", {}),
        "subject_scores": state.get("subject_scores", {}),
        "total_subjects_used": len(state.get("subject_counts", {})),
        "total_lessons_used": len(state.get("lesson_counts", {})),
    }


def format_blacklist_for_prompt(blacklist: list, max_items: int = 25) -> str:
    """
    Format a blacklist as a string for prompt injection.
    
    Args:
        blacklist: List of items to avoid
        max_items: Maximum items to include
    
    Returns:
        Formatted string like "Animals, Geography, Weather"
    """
    items = blacklist[-max_items:] if blacklist else []
    if not items:
        return "(none)"
    return ", ".join(items)


# =============================================================================
# USED LESSONS TRACKING (scoped to current training run)
# =============================================================================

def load_used_lessons() -> list[str]:
    """
    Load lessons used in the current training run.
    
    Returns:
        List of lesson title strings
    """
    if not os.path.exists(USED_LESSONS_FILE):
        return []
    
    try:
        with open(USED_LESSONS_FILE, "r") as f:
            data = json.load(f)
        return data.get("lessons", [])
    except Exception as e:
        print(f"[Curriculum] Error loading used lessons: {e}")
        return []


def add_used_lesson(lesson: str):
    """
    Add a lesson to the used lessons list for the current training run.
    Process-safe via file lock.
    
    Args:
        lesson: The lesson title string
    """
    with get_lock(USED_LESSONS_FILE):
        data = {
            "lessons": load_used_lessons(),
            "since_training": None,
        }
        
        # Add lesson if not already present (case-insensitive)
        lesson_lower = lesson.lower().strip()
        existing_lower = [l.lower().strip() for l in data["lessons"]]
        
        if lesson_lower not in existing_lower:
            data["lessons"].append(lesson)
            
            # Set timestamp if this is the first lesson
            if data["since_training"] is None:
                data["since_training"] = datetime.now().isoformat()
            
            # Save
            os.makedirs(os.path.dirname(USED_LESSONS_FILE), exist_ok=True)
            with open(USED_LESSONS_FILE, "w") as f:
                json.dump(data, f, indent=2)


def clear_used_lessons():
    """
    Clear the used lessons list (called after training completes).
    """
    data = {
        "lessons": [],
        "since_training": None,
    }
    os.makedirs(os.path.dirname(USED_LESSONS_FILE), exist_ok=True)
    with open(USED_LESSONS_FILE, "w") as f:
        json.dump(data, f, indent=2)


def check_lesson_overlap(candidate_lesson: str, used_lessons: list) -> bool:
    """
    Check if a candidate lesson substantially overlaps with any lessons used in the current run.
    
    Two-layer check:
      1. Fast fuzzy word-overlap (no API call) — catches obvious near-dupes
      2. LLM-based semantic comparison — catches subtle overlaps
    
    Args:
        candidate_lesson: The candidate lesson title to check
        used_lessons: List of lesson titles used in the current training run
    
    Returns:
        True if candidate overlaps with any used lesson, False otherwise.
        Returns False (allow) on API failure to avoid blocking valid lessons.
    """
    if not used_lessons:
        return False  # No used lessons to compare against
    
    # Layer 1: Fast fuzzy word-overlap pre-filter (no API call)
    is_fuzzy, matched = dedup_fuzzy_match(candidate_lesson, used_lessons)
    if is_fuzzy:
        return True
    
    # Layer 2: LLM semantic check (only if fuzzy didn't flag)
    # Cap to most recent items to avoid token explosion on large lists.
    # Fuzzy check (Layer 1) already scanned the full list.
    capped_lessons = used_lessons[-SEMANTIC_CHECK_MAX_ITEMS:] if len(used_lessons) > SEMANTIC_CHECK_MAX_ITEMS else used_lessons
    used_list = ", ".join(capped_lessons)
    
    prompt = f"""You are checking if a new lesson topic overlaps with lessons already used in this training run.

Lessons already used: {used_list}

New candidate lesson: "{candidate_lesson}"

Does the candidate lesson substantially overlap with any of the already used lessons?
- Consider semantic similarity (e.g., "Ancient Egypt" vs "Egyptian Pyramids", "The Water Cycle" vs "How Water Moves")
- Consider if they cover the same core topic area
- Ignore minor variations in wording if they're essentially the same topic

Answer with ONLY: YES or NO"""
    
    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,  # gpt-4o-mini
            messages=[
                {"role": "system", "content": "You check if lessons overlap. Answer YES or NO only."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,  # Low temperature for consistent judgment
            max_tokens=10,
        )
        
        result = response.choices[0].message.content.strip().upper()
        return result.startswith("YES")
        
    except Exception as e:
        print(f"[Curriculum] Semantic overlap check error: {e}, allowing candidate")
        return False  # On error, allow the candidate


if __name__ == "__main__":
    # Test the curriculum manager
    print("Testing Curriculum Manager...")
    print()
    
    # Test normalization
    print("Normalization tests:")
    test_names = ["Math", "Mathematics", "Math & Science", "math and science", "  Language Arts  "]
    for name in test_names:
        print(f"  '{name}' -> '{normalize_subject(name)}'")
    print()
    
    # Record some subjects
    print("Recording subjects:")
    for subj in ["Language Arts", "Mathematics", "Science"]:
        record_subject(subj)
        print(f"  Recorded: {subj}")
    print()
    
    # Record some lessons
    print("Recording lessons:")
    record_lesson("Language Arts", "Vowel Sounds")
    record_lesson("Mathematics", "Counting to 10")
    print("  Recorded: Language Arts :: Vowel Sounds")
    print("  Recorded: Mathematics :: Counting to 10")
    print()
    
    # Get blacklists
    blacklists = get_blacklists()
    print("Current blacklists:")
    print(f"  Recent subjects: {blacklists['recent_subjects']}")
    print(f"  Recent lessons: {blacklists['recent_lessons']}")
    print()
    
    # Get stats
    stats = get_subject_stats()
    print(f"Stats: {json.dumps(stats, indent=2)}")