data/episodes/src/graph.py

"""Graph state manager for the social network simulation."""

from __future__ import annotations

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

from schemas import Post, round2


class GraphStateManager:
    """Mutable graph state for a single episode."""

    def __init__(self, episode_id: int, neutral_ids: List[int] = None):
        self.episode_id = episode_id
        self.step: int = 0
        self.posts: List[Post] = []
        self.post_counter: int = 0

        self.account_posts: Dict[int, List[int]] = {}
        self.account_citations_given: Dict[int, List[int]] = {}
        self.account_citations_received: Dict[int, List[int]] = {}

        self.flags: Dict[int, int] = {}
        self.burned_accounts: set = set()

        self.neutral_ids: set = set(neutral_ids) if neutral_ids else set()
        self.neutral_post_window: List[Post] = []  # last 10 neutral posts

        self.prop_pivot_log: Dict[int, List[int]] = {}  # account_id -> [step, step, ...]

    def add_post(
        self,
        account_id: int,
        text: str,
        topic: str,
        framing: str,
        veracity: bool,
        cites_post_id: Optional[int] = None,
        cites_account_id: Optional[int] = None,
    ) -> Post:
        self.post_counter += 1
        post = Post(
            post_id=self.post_counter,
            step=self.step,
            account_id=account_id,
            text=text,
            topic=topic,
            framing=framing,
            veracity=veracity,
            cites_post_id=cites_post_id,
            cites_account_id=cites_account_id,
        )
        self.posts.append(post)
        self.account_posts.setdefault(account_id, []).append(post.post_id)

        # Track neutral posts in rolling window
        if account_id in self.neutral_ids and cites_post_id is None:
            self.neutral_post_window.append(post)
            if len(self.neutral_post_window) > 10:
                self.neutral_post_window.pop(0)

        if cites_post_id is not None:
            self.account_citations_given.setdefault(account_id, []).append(post.post_id)
            if cites_account_id is not None:
                self.account_citations_received.setdefault(cites_account_id, []).append(post.post_id)

        return post

    def flag_account(self, account_id: int, step: int) -> None:
        self.flags[account_id] = step

    def burn_account(self, account_id: int) -> None:
        self.burned_accounts.add(account_id)

    def is_burned(self, account_id: int) -> bool:
        return account_id in self.burned_accounts

    def is_flagged(self, account_id: int) -> bool:
        return account_id in self.flags

    def log_pivot(self, account_id: int, step: int) -> None:
        """Record when a propagandist takes MIMIC or PIVOT action."""
        self.prop_pivot_log.setdefault(account_id, []).append(step)

    def get_account_velocity(self, account_id: int, window: int = 10) -> float:
        recent = [p for p in self.posts if p.account_id == account_id and self.step - p.step <= window]
        if window == 0:
            return 0.0
        return round2(len(recent) / window)

    def get_recent_posts(self, n_steps: int = 8) -> List[Post]:
        return [p for p in self.posts if self.step - p.step < n_steps]

    def get_last_n_posts(self, n: int = 5) -> List[Post]:
        return self.posts[-n:] if len(self.posts) >= n else self.posts[:]

    def to_dict(self, full: bool = True) -> dict:
        base = {
            "episode_id": self.episode_id,
            "step": self.step,
            "num_posts": len(self.posts),
            "num_citations": sum(1 for p in self.posts if p.cites_post_id is not None),
            "flags": {str(k): v for k, v in self.flags.items()},
            "burned": sorted(list(self.burned_accounts)),
            "prop_pivot_log": {str(k): v for k, v in self.prop_pivot_log.items()},
        }
        if full:
            base["posts"] = [p.model_dump() for p in self.posts]
        else:
            base["recent_posts"] = [p.model_dump() for p in self.get_recent_posts(n_steps=10)]
        return base

    def save(self, path: Path, full: bool = True) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(self.to_dict(full=full), f, indent=2)