data/episodes/src/utils.py

"""Utility functions for loading data and common operations."""

import json
import random
from pathlib import Path
from typing import Dict, List

from config import NEUTRAL_PERSONAS, EPISODES_DIR


def load_data_pool(pool_dir: str = "data/neutral_data_pool/generated") -> Dict[str, Dict[str, List[str]]]:
    """Load the neutral data pool into memory.

    Returns:
        {persona: {topic: [text1, text2, ...]}}
    """
    pool_dir = Path(pool_dir)
    pool: Dict[str, Dict[str, List[str]]] = {}

    for persona in NEUTRAL_PERSONAS:
        persona_dir = pool_dir / persona
        if not persona_dir.exists():
            print(f"Warning: Missing data directory for persona '{persona}' at {persona_dir}")
            pool[persona] = {}
            continue

        pool[persona] = {}
        for jsonl_file in persona_dir.glob("*.jsonl"):
            topic = jsonl_file.stem.replace(f"{persona}_", "")
            texts = []
            with open(jsonl_file, "r") as f:
                for line in f:
                    try:
                        record = json.loads(line)
                        texts.append(record.get("text", ""))
                    except json.JSONDecodeError:
                        continue
            if texts:
                pool[persona][topic] = texts

    return pool


def create_minimal_data_pool() -> Dict[str, Dict[str, List[str]]]:
    """Create a minimal in-memory data pool for testing without files."""
    topics = ["politics", "sports", "technology", "health", "entertainment"]
    pool = {}
    for persona in NEUTRAL_PERSONAS:
        pool[persona] = {}
        for topic in topics:
            pool[persona][topic] = [
                f"[{persona}] Sample post about {topic} #{i}"
                for i in range(50)
            ]
    return pool


def load_agenda_config(path: str = "data/episodes/agenda_config.json") -> Dict[str, List[str]]:
    """Load agenda configuration."""
    path = Path(path)
    if not path.exists():
        return {}
    with open(path) as f:
        return json.load(f)


def load_propaganda_bank(bank_dir: str = "data/episodes/propaganda_bank") -> Dict[str, Dict[int, Dict[str, List[str]]]]:
    """Load propaganda bank into memory.

    Returns:
        {topic_slug: {angle_idx: {persona: [text1, text2, ...]}}}
    """
    bank_dir = Path(bank_dir)
    if not bank_dir.exists():
        return {}

    bank: Dict[str, Dict[int, Dict[str, List[str]]]] = {}
    personas = ["hothead", "clickbait", "careful_journalist", "partisan", "rando"]

    for topic_dir in bank_dir.iterdir():
        if not topic_dir.is_dir():
            continue
        topic_slug = topic_dir.name
        bank[topic_slug] = {}

        for angle_dir in topic_dir.iterdir():
            if not angle_dir.is_dir() or not angle_dir.name.startswith("angle_"):
                continue
            angle_idx = int(angle_dir.name.replace("angle_", ""))
            bank[topic_slug][angle_idx] = {}

            for persona_file in angle_dir.glob("*.jsonl"):
                persona = persona_file.stem
                if persona not in personas:
                    continue
                texts = []
                with open(persona_file, "r") as f:
                    for line in f:
                        try:
                            record = json.loads(line)
                            texts.append(record.get("text", ""))
                        except json.JSONDecodeError:
                            continue
                if texts:
                    bank[topic_slug][angle_idx][persona] = texts

    return bank


def get_episode_pool_path() -> Path:
    """Get the directory path for episode pool (relative to project root)."""
    # Find project root by looking for PLANS.md or README.md
    current = Path.cwd()
    while current != current.parent:
        if (current / "PLANS.md").exists() or (current / "README.md").exists():
            path = current / EPISODES_DIR
            path.mkdir(parents=True, exist_ok=True)
            return path
        current = current.parent
    # Fallback to relative path
    path = Path(EPISODES_DIR)
    path.mkdir(parents=True, exist_ok=True)
    return path


def get_all_episode_dirs() -> List[Path]:
    """Get all episode directories."""
    pool_path = get_episode_pool_path()
    return sorted(pool_path.glob("ep_*"))