viral_script_engine/tests/test_training_pipeline.py

"""
Tests for Phase 3 — Training Pipeline.

Run: pytest viral_script_engine/tests/test_training_pipeline.py -v
"""
import json
import tempfile
from pathlib import Path
from unittest.mock import patch, MagicMock

import pytest

BASE_DIR = Path(__file__).parent.parent
CURRICULUM_DIR = BASE_DIR / "data" / "curriculum"


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def dummy_episode_config():
    return {
        "episode_config_id": "easy_001",
        "difficulty": "easy",
        "script_id": "S01",
        "script_text": (
            "Okay so real talk — I've been broke my whole life. "
            "One trick changed everything. Mutual funds. Just SIPs."
        ),
        "region": "Mumbai Gen Z",
        "platform": "Reels",
        "niche": "personal finance",
        "dominant_flaw": "buried_hook",
        "expected_critique_class": "hook_weakness",
        "expected_action": "hook_rewrite",
        "curriculum_notes": "One obvious flaw.",
    }


@pytest.fixture
def mock_env(dummy_episode_config):
    env = MagicMock()
    env.max_steps = 5
    obs = {
        "current_script": dummy_episode_config["script_text"],
        "original_script": dummy_episode_config["script_text"],
        "region": dummy_episode_config["region"],
        "platform": dummy_episode_config["platform"],
        "niche": dummy_episode_config["niche"],
        "step_num": 0,
        "max_steps": 5,
        "debate_history": [],
        "reward_components": {
            "r1_hook_strength": 0.4,
            "r2_coherence": 0.6,
            "r3_cultural_alignment": 0.5,
            "r4_debate_resolution": None,
            "r5_defender_preservation": None,
            "total": 0.5,
        },
        "difficulty_level": "easy",
        "episode_id": "test-episode-001",
    }
    env.reset.return_value = (obs, {})
    env.reset_from_config.return_value = (obs, {})
    env.step.return_value = (
        obs,
        0.65,
        True,
        False,
        {
            "reward_components": obs["reward_components"],
            "anti_gaming_triggered": False,
            "anti_gaming_log": {"triggered": False, "penalty_applied": 0.0},
        },
    )
    return env


@pytest.fixture
def mock_model():
    def _model(prompt: str) -> str:
        return json.dumps({
            "action_type": "hook_rewrite",
            "target_section": "hook",
            "instruction": "Open with the most surprising claim immediately.",
            "critique_claim_id": "C1",
            "reasoning": "The hook is buried — move the key reveal to line 1.",
        })
    return _model


# ---------------------------------------------------------------------------
# Test 1: build_training_prompts returns non-empty dataset with correct format
# ---------------------------------------------------------------------------

def test_build_training_prompts_easy():
    """build_training_prompts('easy') returns non-empty list with correct prompt format."""
    if not (CURRICULUM_DIR / "easy_tier.jsonl").exists():
        pytest.skip("easy_tier.jsonl not found — run build_curriculum.py first")

    from viral_script_engine.training.rollout_function import build_training_prompts
    prompts = build_training_prompts("easy")

    assert len(prompts) > 0, "Should return at least one prompt"
    first = prompts[0]
    assert "##EPISODE_CONFIG##" in first, "Prompt must contain embedded episode config header"
    assert "##END_CONFIG##" in first, "Prompt must contain end-config marker"
    assert "<|system|>" in first, "Prompt must include system role tag"
    assert "CURRENT SCRIPT:" in first, "Prompt must include script section"
    assert "AVAILABLE ACTIONS:" in first, "Prompt must list available actions"


def test_build_training_prompts_config_parseable():
    """Episode config embedded in prompt must be valid JSON."""
    if not (CURRICULUM_DIR / "easy_tier.jsonl").exists():
        pytest.skip("easy_tier.jsonl not found")

    import re
    from viral_script_engine.training.rollout_function import build_training_prompts
    prompts = build_training_prompts("easy")

    for prompt in prompts[:3]:
        match = re.search(r"##EPISODE_CONFIG##\s*(\{.*?\})\s*##END_CONFIG##", prompt, re.DOTALL)
        assert match, "Config header must be parseable"
        config = json.loads(match.group(1))
        assert "script_text" in config
        assert "region" in config
        assert "difficulty" in config


# ---------------------------------------------------------------------------
# Test 2: rollout_fn completes one episode given a mock model returning valid JSON
# ---------------------------------------------------------------------------

def test_rollout_fn_single_episode(mock_env, mock_model, dummy_episode_config):
    """rollout_fn completes one episode and returns (completions, rewards)."""
    from viral_script_engine.training.rollout_function import build_rollout_fn, _config_to_prompt

    rollout_fn = build_rollout_fn(mock_env, max_steps=5)
    prompt = _config_to_prompt(dummy_episode_config)

    completions, rewards = rollout_fn([prompt], model=mock_model, tokenizer=None)

    assert len(completions) == 1
    assert len(rewards) == 1
    assert isinstance(rewards[0], float)
    assert 0.0 <= rewards[0] <= 1.0, "Reward should be in [0, 1]"


def test_rollout_fn_batch(mock_env, mock_model, dummy_episode_config):
    """rollout_fn handles a batch of prompts."""
    from viral_script_engine.training.rollout_function import build_rollout_fn, _config_to_prompt

    rollout_fn = build_rollout_fn(mock_env, max_steps=5)
    prompt = _config_to_prompt(dummy_episode_config)

    completions, rewards = rollout_fn([prompt] * 3, model=mock_model, tokenizer=None)

    assert len(completions) == 3
    assert len(rewards) == 3


# ---------------------------------------------------------------------------
# Test 3: GRPOConfig builds without error
# ---------------------------------------------------------------------------

def test_grpo_config_builds():
    """GRPOConfig builds without error when trl is installed and pyarrow DLL is available."""
    try:
        from trl import GRPOConfig
    except Exception:
        pytest.skip("trl/GRPOConfig not available on this machine (pyarrow DLL or import issue)")

    from viral_script_engine.training.train_grpo import build_grpo_config
    with tempfile.TemporaryDirectory() as tmpdir:
        config = build_grpo_config(output_dir=tmpdir, num_steps=200, dry_run=True)
        assert config.max_steps == 5
        assert config.per_device_train_batch_size == 1


# ---------------------------------------------------------------------------
# Test 4: Model saving uses save_pretrained_merged
# ---------------------------------------------------------------------------

def test_model_save_uses_merged():
    """Training script uses save_pretrained_merged, not save_pretrained."""
    train_script = Path(__file__).parent.parent / "training" / "train_grpo.py"
    content = train_script.read_text(encoding="utf-8")

    assert "save_pretrained_merged" in content, (
        "train_grpo.py must use model.save_pretrained_merged() — "
        "naive upcast from 4-bit is not supported"
    )
    # Ensure the naive form is only in comments or strings, not as a bare call
    import re
    bare_calls = re.findall(r"model\.save_pretrained\(", content)
    assert len(bare_calls) == 0, (
        "train_grpo.py must NOT use model.save_pretrained() — use save_pretrained_merged"
    )


# ---------------------------------------------------------------------------
# Test 5: plot_training_curves generates PNG given valid JSON inputs
# ---------------------------------------------------------------------------

def test_plot_training_curves_generates_png():
    """plot_training_curves() generates a PNG file given valid JSON inputs."""
    from viral_script_engine.training.reward_curves import plot_training_curves

    episode_template = {
        "episode_num": 1,
        "difficulty": "easy",
        "total_reward": 0.55,
        "steps": [
            {"r1": 0.6, "r2": 0.5, "r3": 0.4, "r4": 0.5, "r5": 0.6, "total": 0.55}
        ],
    }

    baseline = [dict(episode_template, episode_num=i, total_reward=0.4 + i * 0.01)
                for i in range(1, 21)]
    trained = [dict(episode_template, episode_num=i, total_reward=0.55 + i * 0.01)
               for i in range(1, 21)]

    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)
        base_path = tmpdir / "baseline_results.json"
        train_path = tmpdir / "training_results.json"
        out_path = tmpdir / "training_vs_baseline.png"

        base_path.write_text(json.dumps(baseline))
        train_path.write_text(json.dumps(trained))

        plot_training_curves(
            baseline_log_path=str(base_path),
            training_log_path=str(train_path),
            output_path=str(out_path),
        )

        assert out_path.exists(), "PNG file must be created"
        assert out_path.stat().st_size > 1000, "PNG file must be non-trivial"
        pdf_path = out_path.with_suffix(".pdf")
        assert pdf_path.exists(), "PDF file must also be created"


# ---------------------------------------------------------------------------
# Test 6: Env reset_from_config works correctly
# ---------------------------------------------------------------------------

def test_env_reset_from_config(dummy_episode_config, monkeypatch):
    """ViralScriptEnv.reset_from_config() resets state from a given config."""
    from viral_script_engine.environment.env import ViralScriptEnv
    from viral_script_engine.rewards import r2_coherence, r5_defender_preservation
    from viral_script_engine.rewards.r2_coherence import CoherenceRewardResult
    from viral_script_engine.rewards.r5_defender_preservation import DefenderPreservationResult

    monkeypatch.setattr(
        r2_coherence.CoherenceReward, "score",
        lambda self, a, b: CoherenceRewardResult(score=0.75, raw_similarity=0.85, interpretation="good_coherence"),
    )
    monkeypatch.setattr(
        r5_defender_preservation.DefenderPreservationReward, "score",
        lambda self, d, s: DefenderPreservationResult(score=0.70, max_similarity=0.80, best_matching_sentence="[test mock]"),
    )

    env = ViralScriptEnv(
        scripts_path=str(BASE_DIR / "data" / "test_scripts" / "scripts.json"),
        cultural_kb_path=str(BASE_DIR / "data" / "cultural_kb.json"),
        max_steps=5,
        difficulty="easy",
    )

    obs, info = env.reset_from_config(dummy_episode_config)

    assert obs["current_script"] == dummy_episode_config["script_text"]
    assert obs["region"] == dummy_episode_config["region"]
    assert obs["platform"] == dummy_episode_config["platform"]
    assert obs["niche"] == dummy_episode_config["niche"]
    assert obs["step_num"] == 0