metrics_manager.py

# metrics_manager.py
# Phase 4: Session metrics storage, rolling metrics, and baseline computation.

import os
import json
from datetime import datetime
from typing import Optional

from config import (
    METRICS_FILE,
    SESSION_METRICS_DIR,
    EVAL_COMPARE_WINDOW,
    TRAIN_PROGRESS_SIGNAL_WINDOW,
)
from file_lock_utils import get_lock


def ensure_metrics_dirs():
    """Ensure metrics directories exist."""
    os.makedirs(SESSION_METRICS_DIR, exist_ok=True)
    os.makedirs(os.path.dirname(METRICS_FILE), exist_ok=True)


# =============================================================================
# PER-SESSION METRICS
# =============================================================================

def save_session_metrics(session_metrics: dict) -> str:
    """
    Save per-session metrics to a JSON file.
    
    Args:
        session_metrics: Dict with session data (must include session_id)
    
    Returns:
        Path to the saved file
    """
    ensure_metrics_dirs()
    session_id = session_metrics.get("session_id", datetime.now().strftime("%Y%m%d_%H%M%S"))
    filepath = os.path.join(SESSION_METRICS_DIR, f"session_{session_id}.json")
    
    with open(filepath, "w") as f:
        json.dump(session_metrics, f, indent=2)
    
    return filepath


def load_session_metrics(session_id: str) -> Optional[dict]:
    """
    Load a specific session's metrics.
    
    Args:
        session_id: The session identifier
    
    Returns:
        Session metrics dict or None if not found
    """
    filepath = os.path.join(SESSION_METRICS_DIR, f"session_{session_id}.json")
    if os.path.exists(filepath):
        with open(filepath, "r") as f:
            return json.load(f)
    return None


def get_recent_session_metrics(n: int = 10, training_phase: str = None) -> list:
    """
    Get the most recent N session metrics, optionally filtered by training_phase.
    
    Reads from two sources:
      - Legacy per-session JSON files: memory/session_metrics/session_*.json
      - Batch meta JSONL files: logs/batch_*_meta.jsonl (one session per line)
    
    Args:
        n: Number of sessions to return
        training_phase: Filter by training phase ("normal", "posttrain_eval", etc.)
    
    Returns:
        List of session metrics dicts, most recent first
    """
    from config import LOG_DIR
    ensure_metrics_dirs()
    
    all_metrics = []  # List of (timestamp_str, data_dict)
    
    # Source 1: Legacy per-session JSON files
    if os.path.exists(SESSION_METRICS_DIR):
        for fname in os.listdir(SESSION_METRICS_DIR):
            if fname.startswith("session_") and fname.endswith(".json"):
                filepath = os.path.join(SESSION_METRICS_DIR, fname)
                try:
                    with open(filepath, "r") as f:
                        data = json.load(f)
                    ts = data.get("session_id", fname)
                    all_metrics.append((ts, data))
                except Exception as e:
                    print(f"[MetricsManager] Error loading {filepath}: {e}")
    
    # Source 2: Batch meta JSONL files
    if os.path.exists(LOG_DIR):
        for fname in os.listdir(LOG_DIR):
            if fname.startswith("batch_") and fname.endswith("_meta.jsonl"):
                filepath = os.path.join(LOG_DIR, fname)
                try:
                    with open(filepath, "r") as f:
                        for line in f:
                            line = line.strip()
                            if not line:
                                continue
                            data = json.loads(line)
                            ts = data.get("session_id", fname)
                            all_metrics.append((ts, data))
                except Exception as e:
                    print(f"[MetricsManager] Error loading {filepath}: {e}")
    
    # Sort by session_id (timestamp string), newest first
    all_metrics.sort(key=lambda x: x[0], reverse=True)
    
    results = []
    for _, data in all_metrics:
        if len(results) >= n:
            break
        if training_phase is None or data.get("training_phase") == training_phase:
            results.append(data)
    
    return results


# =============================================================================
# ROLLING METRICS (metrics.json)
# =============================================================================

def load_rolling_metrics() -> dict:
    """Load the rolling metrics file."""
    defaults = _default_rolling_metrics()
    
    if os.path.exists(METRICS_FILE):
        try:
            with open(METRICS_FILE, "r") as f:
                loaded = json.load(f)
                # Merge with defaults to ensure all required fields exist
                # This handles migration from old metrics.json format
                merged = defaults.copy()
                merged.update(loaded)
                # Ensure nested dicts are properly initialized
                if "recent_summary" not in merged or not isinstance(merged.get("recent_summary"), dict):
                    merged["recent_summary"] = defaults["recent_summary"]
                return merged
        except Exception as e:
            print(f"[MetricsManager] Error loading metrics: {e}")
            return defaults
    return defaults


def save_rolling_metrics(metrics: dict):
    """Save the rolling metrics file."""
    ensure_metrics_dirs()
    with open(METRICS_FILE, "w") as f:
        json.dump(metrics, f, indent=2)


def _default_rolling_metrics() -> dict:
    """Return default rolling metrics structure."""
    return {
        "total_sessions": 0,
        "total_graded_turns": 0,
        "graded_turns_since_last_train": 0,
        "usable_turns_since_last_train": 0,
        "last_train_timestamp": None,
        "last_train_result": None,  # "kept", "rolled_back", "train_failed"
        "ewma_score": None,
        "ewma_compliance": None,
        "last_n_session_ids": [],  # Most recent session IDs
        "recent_summary": {
            "avg_score": None,
            "avg_compliance": None,
            "progress_signal_rate": None,
            "sessions_count": 0,
        },
        "total_training_runs": 0,
        "total_rollbacks": 0,
    }


def update_rolling_metrics(session_metrics: dict) -> dict:
    """
    Update rolling metrics after a session.
    
    Thread/process-safe: acquires a file lock around the load-modify-save cycle.
    
    Args:
        session_metrics: The just-completed session's metrics
    
    Returns:
        Updated rolling metrics dict
    """
    with get_lock(METRICS_FILE):
        metrics = load_rolling_metrics()
        
        # Update counts
        metrics["total_sessions"] = metrics.get("total_sessions", 0) + 1
        graded_count = session_metrics.get("graded_turns_count", 0)
        metrics["total_graded_turns"] = metrics.get("total_graded_turns", 0) + graded_count
        metrics["graded_turns_since_last_train"] = metrics.get("graded_turns_since_last_train", 0) + graded_count
        usable_count = session_metrics.get("usable_turns_count", 0)
        metrics["usable_turns_since_last_train"] = metrics.get("usable_turns_since_last_train", 0) + usable_count
        
        # Update session ID list (keep last 20)
        session_ids = metrics.get("last_n_session_ids", [])
        session_ids.append(session_metrics.get("session_id"))
        metrics["last_n_session_ids"] = session_ids[-20:]
        
        # Update EWMA (alpha=0.2)
        alpha = 0.2
        session_score = session_metrics.get("avg_score_session")
        session_compliance = session_metrics.get("compliance_rate_session")
        
        if session_score is not None:
            if metrics.get("ewma_score") is None:
                metrics["ewma_score"] = session_score
            else:
                metrics["ewma_score"] = alpha * session_score + (1 - alpha) * metrics["ewma_score"]
        
        if session_compliance is not None:
            if metrics.get("ewma_compliance") is None:
                metrics["ewma_compliance"] = session_compliance
            else:
                metrics["ewma_compliance"] = alpha * session_compliance + (1 - alpha) * metrics["ewma_compliance"]
        
        # Recompute recent summary from last 10 sessions
        metrics["recent_summary"] = _compute_recent_summary(10)
        
        save_rolling_metrics(metrics)
        return metrics


def _compute_recent_summary(n: int) -> dict:
    """Compute summary stats from last N sessions."""
    recent = get_recent_session_metrics(n, training_phase="normal")
    
    if not recent:
        return {
            "avg_score": None,
            "avg_compliance": None,
            "progress_signal_rate": None,
            "sessions_count": 0,
        }
    
    scores = [s.get("avg_score_session") for s in recent if s.get("avg_score_session") is not None]
    compliances = [s.get("compliance_rate_session") for s in recent if s.get("compliance_rate_session") is not None]
    progress_signals = [s.get("progress_signal_session", False) for s in recent]
    
    return {
        "avg_score": sum(scores) / len(scores) if scores else None,
        "avg_compliance": sum(compliances) / len(compliances) if compliances else None,
        "progress_signal_rate": sum(1 for p in progress_signals if p) / len(progress_signals) if progress_signals else None,
        "sessions_count": len(recent),
    }


# =============================================================================
# BASELINE COMPUTATION (for eval/rollback)
# =============================================================================

def compute_baseline(window: int = None) -> dict:
    """
    Compute baseline metrics from the last N "normal" sessions (before training).
    
    Args:
        window: Number of sessions to consider (default: EVAL_COMPARE_WINDOW)
    
    Returns:
        Dict with avg_score, avg_compliance, ewma_score, ewma_compliance
    """
    window = window or EVAL_COMPARE_WINDOW
    
    recent = get_recent_session_metrics(window, training_phase="normal")
    
    if not recent:
        return {
            "avg_score": 0.0,
            "avg_compliance": 0.0,
            "ewma_score": 0.0,
            "ewma_compliance": 0.0,
            "sessions_count": 0,
        }
    
    scores = [s.get("avg_score_session", 0) for s in recent if s.get("avg_score_session") is not None]
    compliances = [s.get("compliance_rate_session", 0) for s in recent if s.get("compliance_rate_session") is not None]
    
    # Also get EWMA from rolling metrics
    metrics = load_rolling_metrics()
    
    return {
        "avg_score": sum(scores) / len(scores) if scores else 0.0,
        "avg_compliance": sum(compliances) / len(compliances) if compliances else 0.0,
        "ewma_score": metrics.get("ewma_score", 0.0) or 0.0,
        "ewma_compliance": metrics.get("ewma_compliance", 0.0) or 0.0,
        "sessions_count": len(recent),
    }


def compute_eval_metrics(eval_session_ids: list) -> dict:
    """
    Compute metrics from specific eval sessions.
    
    Args:
        eval_session_ids: List of session IDs to include
    
    Returns:
        Dict with avg_score, avg_compliance
    """
    scores = []
    compliances = []
    
    for session_id in eval_session_ids:
        session = load_session_metrics(session_id)
        if session:
            if session.get("avg_score_session") is not None:
                scores.append(session["avg_score_session"])
            if session.get("compliance_rate_session") is not None:
                compliances.append(session["compliance_rate_session"])
    
    return {
        "avg_score": sum(scores) / len(scores) if scores else 0.0,
        "avg_compliance": sum(compliances) / len(compliances) if compliances else 0.0,
        "sessions_count": len(scores),
    }


# =============================================================================
# TRAINING TRIGGER HELPERS
# =============================================================================

def get_progress_signal_rate(window: int = None) -> float:
    """
    Get the progress signal rate over the last N sessions.
    
    Args:
        window: Number of sessions (default: TRAIN_PROGRESS_SIGNAL_WINDOW)
    
    Returns:
        Float between 0.0 and 1.0
    """
    window = window or TRAIN_PROGRESS_SIGNAL_WINDOW
    
    recent = get_recent_session_metrics(window, training_phase="normal")
    
    if not recent:
        return 0.0
    
    progress_signals = [s.get("progress_signal_session", False) for s in recent]
    return sum(1 for p in progress_signals if p) / len(progress_signals)


def record_training_attempt(result: str, session_id: str = None):
    """
    Record a training attempt result.
    
    Thread/process-safe: acquires a file lock around the load-modify-save cycle.
    
    Args:
        result: "kept", "rolled_back", or "train_failed"
        session_id: The session ID (timestamp) that triggered training.
                    Used as a training-run boundary for recency weighting:
                    all sessions with keys < this boundary were generated by the
                    previous model version.
    """
    with get_lock(METRICS_FILE):
        metrics = load_rolling_metrics()
        metrics["last_train_timestamp"] = datetime.now().isoformat()
        metrics["last_train_result"] = result
        metrics["total_training_runs"] = metrics.get("total_training_runs", 0) + 1
        
        if result == "rolled_back":
            metrics["total_rollbacks"] = metrics.get("total_rollbacks", 0) + 1
        
        # Record training boundary for recency weighting.
        if session_id:
            boundaries = metrics.get("training_run_boundaries", [])
            boundaries.append(session_id)
            metrics["training_run_boundaries"] = boundaries
        
        # Reset turn counters
        metrics["graded_turns_since_last_train"] = 0
        metrics["usable_turns_since_last_train"] = 0
        
        save_rolling_metrics(metrics)


# =============================================================================
# SESSION METRICS SCHEMA HELPER
# =============================================================================

def create_session_metrics(
    session_id: str,
    subject_of_the_day: str = None,
    lesson_of_the_day: str = None,
    age_band_start: int = 0,
    age_band_end: int = 0,
    graded_turns_count: int = 0,
    usable_turns_count: int = 0,
    avg_score_session: float = 0.0,
    compliance_rate_session: float = 0.0,
    progress_signal_session: bool = False,
    task_category_counts: dict = None,
    avg_basil_tokens: float = None,
    early_stopped: bool = False,
    stop_reason: str = "completed",
    training_phase: str = "normal",
    transcript_path: str = None,
    graded_path: str = None,
    assessed_age_band: int = None,
) -> dict:
    """
    Create a session metrics dict with the standard schema.
    
    Returns:
        Properly structured session metrics dict
    """
    return {
        "session_id": session_id,
        "timestamp": datetime.now().isoformat(),
        "subject_of_the_day": subject_of_the_day,
        "lesson_of_the_day": lesson_of_the_day,
        "age_band_start": age_band_start,
        "age_band_end": age_band_end,
        "graded_turns_count": graded_turns_count,
        "usable_turns_count": usable_turns_count,
        "avg_score_session": avg_score_session,
        "compliance_rate_session": compliance_rate_session,
        "progress_signal_session": progress_signal_session,
        "task_category_counts": task_category_counts or {},
        "avg_basil_tokens": avg_basil_tokens,
        "early_stopped": early_stopped,
        "stop_reason": stop_reason,
        "training_phase": training_phase,
        "transcript_path": transcript_path,
        "graded_path": graded_path,
        "assessed_age_band": assessed_age_band,
    }


if __name__ == "__main__":
    # Test metrics manager
    print("Testing Metrics Manager...")
    
    ensure_metrics_dirs()
    
    # Test creating session metrics
    test_metrics = create_session_metrics(
        session_id="test_001",
        subject_of_the_day="Mathematics",
        lesson_of_the_day="Counting to 5",
        graded_turns_count=10,
        avg_score_session=2.5,
        compliance_rate_session=0.6,
        progress_signal_session=True,
    )
    
    path = save_session_metrics(test_metrics)
    print(f"Saved to: {path}")
    
    # Test loading
    loaded = load_session_metrics("test_001")
    print(f"Loaded: {json.dumps(loaded, indent=2)}")
    
    # Test rolling metrics
    rolling = update_rolling_metrics(test_metrics)
    print(f"Rolling: {json.dumps(rolling, indent=2)}")
    
    # Test baseline
    baseline = compute_baseline(5)
    print(f"Baseline: {json.dumps(baseline, indent=2)}")
    
    # Cleanup test file
    os.remove(path)
    print("Test complete!")