src/diagnostics/vectorized_eval.py

"""
Day 29 — Vectorized Simulator Architecture.

Processes states in parallel across N independent simulator instances.
Stack N embeddings → shape (N, 384), query MaskablePPO → predict N actions
simultaneously, apply action vector → advance all N simulators.

RUN:
  python -m src.diagnostics.vectorized_eval
"""
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))

import numpy as np
import json
import time
from collections import defaultdict

from sb3_contrib import MaskablePPO
from src.env.discord_env import DiscordEnv, LANG_TO_IDX, NUM_LANGUAGES


ACTION_NAMES = ["ALLOW", "WARN", "DELETE", "TIMEOUT", "BAN"]


class VectorizedSimulator:
    """
    Runs N episodes in parallel by batching observations and
    querying the policy once per step across all active episodes.
    """

    def __init__(self, model_path: str = "data/models/best/best_model.zip",
                 data_dir: str = "data/processed", n_parallel: int = 32):
        self.model = MaskablePPO.load(model_path)
        self.data_dir = data_dir
        self.n_parallel = n_parallel

        # Load shared data
        with open(f"{data_dir}/episodes.json", "r") as f:
            self.all_episodes = json.load(f)
        self.embeddings = np.load(f"{data_dir}/context_embeddings.npy")
        self.toxicity_scores = np.load(f"{data_dir}/toxicity_scores.npy")

        print(f"Vectorized simulator: {len(self.all_episodes)} episodes, "
              f"N={n_parallel} parallel")

    def _make_env(self, ep_data):
        """Create a lightweight episode state tracker."""
        return {
            "thread_id": ep_data["thread_id"],
            "step_indices": ep_data["step_indices"],
            "user_ids": ep_data["user_ids"],
            "languages": ep_data.get("languages", ["en"] * len(ep_data["step_indices"])),
            "current_step": 0,
            "ledger": {},
            "recent_tox": [],
            "recent_actions": [],
            "done": False,
            "results": [],
        }

    def _get_obs_and_mask(self, env):
        """Get observation and action mask for a single env."""
        step = env["current_step"]
        idx = env["step_indices"][step]
        tox = float(self.toxicity_scores[idx])
        uid = env["user_ids"][step]
        lang = env["languages"][step] if step < len(env["languages"]) else "en"

        # Ensure ledger
        if uid not in env["ledger"]:
            env["ledger"][uid] = {"warns": 0.0, "timeouts": 0.0, "total_infractions": 0.0}

        led = env["ledger"][uid]

        # Embedding
        emb = np.nan_to_num(self.embeddings[idx].astype(np.float32), nan=0.0, posinf=1.0, neginf=-1.0)
        emb = np.clip(emb, -1.0, 1.0)

        # Language one-hot
        lang_vec = np.zeros(NUM_LANGUAGES, dtype=np.float32)
        li = LANG_TO_IDX.get(lang, -1)
        if li >= 0:
            lang_vec[li] = 1.0

        obs = {
            "message_embedding": emb,
            "toxicity_score": np.array([np.clip(tox, 0.0, 1.0)], dtype=np.float32),
            "user_history": np.array([
                min(led["warns"] / 5.0, 1.0),
                min(led["timeouts"] / 3.0, 1.0),
                min(led["total_infractions"] / 10.0, 1.0),
            ], dtype=np.float32),
            "server_heat": np.array([0.0, 0.0], dtype=np.float32),
            "language_id": lang_vec,
        }

        # Action mask
        mask = np.ones(5, dtype=bool)
        if led["timeouts"] < 1:
            mask[4] = False
        if led["total_infractions"] < 2:
            mask[3] = False
        if tox < 0.15:
            mask[2:] = False
        if tox >= 0.85:
            mask[0] = False
        if not mask.any():
            mask[0] = True

        return obs, mask, tox, uid, lang

    def _apply_action(self, env, action, tox, uid):
        """Update ledger and record result."""
        led = env["ledger"][uid]

        if action == 1:
            led["warns"] += 1.0
            led["total_infractions"] += 1.0
        elif action == 2:
            led["total_infractions"] += 1.0
        elif action == 3:
            led["timeouts"] += 1.0
            led["total_infractions"] += 1.0
        elif action == 4:
            led["total_infractions"] += 1.0

        env["results"].append({
            "action": int(action),
            "toxicity": tox,
            "user_id": uid,
        })

        env["current_step"] += 1
        if env["current_step"] >= len(env["step_indices"]):
            env["done"] = True

    def run(self, n_episodes: int = None) -> dict:
        """Run vectorized evaluation across all or n_episodes episodes."""
        if n_episodes is None:
            n_episodes = len(self.all_episodes)
        n_episodes = min(n_episodes, len(self.all_episodes))

        rng = np.random.default_rng(42)
        episode_indices = rng.choice(len(self.all_episodes), size=n_episodes, replace=False)

        # Track metrics
        total_steps = 0
        action_counts = defaultdict(int)
        lang_action_counts = defaultdict(lambda: defaultdict(int))
        false_positives = 0
        false_negatives = 0

        start_time = time.time()

        # Process in batches of n_parallel
        completed = 0
        remaining = list(episode_indices)

        while remaining:
            # Fill batch
            batch_indices = remaining[:self.n_parallel]
            remaining = remaining[self.n_parallel:]
            envs = [self._make_env(self.all_episodes[i]) for i in batch_indices]

            # Step all envs until all done
            while any(not e["done"] for e in envs):
                active = [e for e in envs if not e["done"]]

                for env in active:
                    obs, mask, tox, uid, lang = self._get_obs_and_mask(env)
                    action, _ = self.model.predict(obs, deterministic=True, action_masks=mask)
                    action = int(action)

                    self._apply_action(env, action, tox, uid)

                    # Track metrics
                    total_steps += 1
                    action_counts[action] += 1
                    lang_action_counts[lang][action] += 1

                    if tox < 0.30 and action > 0:
                        false_positives += 1
                    if tox >= 0.70 and action == 0:
                        false_negatives += 1

            completed += len(batch_indices)

        elapsed = time.time() - start_time
        steps_per_sec = total_steps / max(elapsed, 0.001)

        # Compile results
        results = {
            "episodes": n_episodes,
            "total_steps": total_steps,
            "elapsed_seconds": round(elapsed, 2),
            "steps_per_second": round(steps_per_sec, 1),
            "action_distribution": {ACTION_NAMES[k]: v for k, v in sorted(action_counts.items())},
            "false_positive_rate": round(false_positives / max(total_steps, 1), 4),
            "false_negative_rate": round(false_negatives / max(total_steps, 1), 4),
            "language_action_counts": {
                lang: {ACTION_NAMES[a]: c for a, c in sorted(actions.items())}
                for lang, actions in sorted(lang_action_counts.items())
            },
        }

        # Print report
        print(f"\n{'=' * 70}")
        print(f"VECTORIZED EVALUATION — {n_episodes} episodes, {total_steps} steps")
        print(f"{'=' * 70}")
        print(f"   Throughput: {steps_per_sec:.0f} steps/sec ({elapsed:.1f}s total)")
        print(f"\n   Action Distribution:")
        for a in range(5):
            count = action_counts.get(a, 0)
            pct = count / max(total_steps, 1) * 100
            bar = "█" * int(pct / 2)
            print(f"     {ACTION_NAMES[a]:8s}: {count:5d} ({pct:5.1f}%) {bar}")
        print(f"\n   FP Rate: {results['false_positive_rate']:.4f}")
        print(f"   FN Rate: {results['false_negative_rate']:.4f}")

        return results


if __name__ == "__main__":
    sim = VectorizedSimulator(n_parallel=32)
    results = sim.run()