polyguard-rl/app/simulator/scenario_generator.py

"""Scenario generation entrypoints."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

from app.common.enums import Difficulty
from app.common.types import PatientProfile
from app.simulator.patient_generator import generate_patient_profile


def generate_patient_scenario(difficulty: Difficulty, patient_id: Optional[str], seed: int) -> PatientProfile:
    return generate_patient_profile(seed=seed, difficulty=difficulty, patient_id=patient_id)


def build_scenario_library(root: Path, easy: int, medium: int, hard: int, seed: int = 42) -> None:
    counts = {
        Difficulty.EASY: easy,
        Difficulty.MEDIUM: medium,
        Difficulty.HARD: hard,
    }
    for diff, count in counts.items():
        out_dir = root / "data" / "scenarios" / diff.value
        out_dir.mkdir(parents=True, exist_ok=True)
        for i in range(count):
            profile = generate_patient_profile(seed=seed + i, difficulty=diff, patient_id=f"{diff.value}_{i:04d}")
            target = out_dir / f"{profile.patient_id}.json"
            target.write_text(json.dumps(profile.model_dump(mode="json"), ensure_ascii=True, indent=2), encoding="utf-8")