scripts/visualize/data_loader.py

"""Shared loader for reflection artifacts. GLM-aware: rows with no summary.json
are tagged status="in_flight" so viz scripts can plot dashed/hatched segments."""

from __future__ import annotations

import json
import math
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

PROJECT_ROOT = Path(__file__).resolve().parents[2]
ARTIFACTS_ROOT = PROJECT_ROOT / "artifacts"
DATASETS_ROOT = PROJECT_ROOT / "datasets"

# Per-model iter mapping. iter 0 = baseline; iters 1..5 = reflection rollouts.
# Each tuple: (iter_number, run_dir_name, label).
MODEL_RUNS: dict[str, list[tuple[int, str, str]]] = {
    "qwen-qwen3-32b-groq": [
        (0, "20260426T014336Z__train-test__qwen-qwen3-32b-groq", "baseline"),
        (1, "20260426T022442Z__train-test__qwen-qwen3-32b-groq__iter01_reflect", "iter1"),
        (2, "20260426T025319Z__train-test__qwen-qwen3-32b-groq__iter02_reflect", "iter2"),
        (3, "20260426T032611Z__train-test__qwen-qwen3-32b-groq__iter01_reflect", "iter3"),
        (4, "20260426T035335Z__train-test__qwen-qwen3-32b-groq__iter02_reflect", "iter4"),
        (5, "20260426T042137Z__train-test__qwen-qwen3-32b-groq__iter03_reflect", "iter5"),
    ],
    "zai-glm-5.1-together": [
        (0, "20260426T035937Z__train-test__zai-org-glm-5-1-together", "baseline"),
        (1, "20260426T044349Z__train-test__zai-org-glm-5-1-together__iter01_reflect", "iter1"),
        (2, "20260426T052859Z__train-test__zai-org-glm-5-1-together__iter02_reflect", "iter2"),
        (3, "20260426T070351Z__train-test__zai-org-glm-5-1-together__iter03_reflect", "iter3"),
    ],
}

MODEL_LABELS: dict[str, str] = {
    "qwen-qwen3-32b-groq": "Qwen3-32B (Groq)",
    "zai-glm-5.1-together": "GLM-5.1 (Together)",
}


@dataclass
class IterRow:
    iter: int
    label: str
    run_dir: Path
    status: str  # "complete" | "in_flight"
    score_normalized: float | None = None
    final_portfolio_value: float | None = None
    initial_portfolio_value: float | None = None
    roi_pct: float | None = None
    bars_completed: int | None = None
    cumulative_reward: float | None = None
    final_score: float | None = None


def load_model_iters(model_id: str) -> list[IterRow]:
    rows: list[IterRow] = []
    for it, run_name, label in MODEL_RUNS[model_id]:
        run_dir = ARTIFACTS_ROOT / "runs" / run_name
        summary = run_dir / "test" / "summary.json"
        if not summary.is_file():
            rows.append(IterRow(it, label, run_dir, status="in_flight"))
            continue
        s = json.loads(summary.read_text())
        rows.append(
            IterRow(
                iter=it,
                label=label,
                run_dir=run_dir,
                status="complete",
                score_normalized=float(s["score_normalized"]),
                final_portfolio_value=float(s["final_portfolio_value"]),
                initial_portfolio_value=float(s["initial_portfolio_value"]),
                roi_pct=(float(s["final_portfolio_value"]) / float(s["initial_portfolio_value"]) - 1.0) * 100.0,
                bars_completed=int(s["bars_completed"]),
                cumulative_reward=float(s["cumulative_reward"]),
                final_score=float(s["final_score"]),
            )
        )
    return rows


_GATE_REJECT_PREFIX = "Long-horizon planning gate:"


def load_advance_day_rows(run_dir: Path) -> list[dict]:
    """Return all SUCCESSFUL advance_day events from test/trajectory.jsonl.

    Skips per-bar gate rejections (the env returns the row but doesn't tick
    the clock; tool_output_excerpt starts with the gate message).
    """
    traj = run_dir / "test" / "trajectory.jsonl"
    rows: list[dict] = []
    with traj.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if row.get("action_type") != "advance_day":
                continue
            if row.get("error"):
                continue
            excerpt = row.get("tool_output_excerpt") or ""
            if excerpt.startswith(_GATE_REJECT_PREFIX):
                continue
            rows.append(row)
    return rows


def load_action_counts(run_dir: Path) -> Counter:
    """Count test-phase actions from trajectory.jsonl.

    Filters out:
    - rows with explicit error field set
    - advance_day calls that hit the per-bar record_decision gate (these
      did not advance state and shouldn't be counted as "successful actions").
    """
    traj = run_dir / "test" / "trajectory.jsonl"
    counts: Counter = Counter()
    if not traj.is_file():
        return counts
    with traj.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if row.get("error"):
                continue
            atype = row.get("action_type") or "unknown"
            if atype == "advance_day":
                excerpt = row.get("tool_output_excerpt") or ""
                if excerpt.startswith(_GATE_REJECT_PREFIX):
                    continue
            counts[atype] += 1
    return counts


def load_prompt(run_dir: Path) -> str:
    """Read the system prompt that drove a given run."""
    return (run_dir / "system_prompt.txt").read_text(encoding="utf-8")


def episode_manifest() -> dict:
    return json.loads(
        (DATASETS_ROOT / "catalog" / "sample-v2" / "episode_manifests" / "tier_test.json").read_text()
    )


def equal_weight_bnh_series(dates: Iterable[str], initial_value: float = 100_000.0) -> list[float]:
    """Compute equal-weight B&H portfolio value at each agent-date.

    Reads the daily_bars parquet, takes universe close marks at each requested
    date, and returns the B&H portfolio value (set equal weights at the first
    date, no rebalancing thereafter)."""
    import pandas as pd

    manifest = episode_manifest()
    universe = list(manifest["universe_asset_ids"])
    bars_path = DATASETS_ROOT / "catalog" / "sample-v2" / "daily_bars" / "part-000.parquet"
    df = pd.read_parquet(bars_path, columns=["asset_id", "session_date", "close"])
    df = df[df["asset_id"].isin(universe)].copy()
    df["session_date"] = pd.to_datetime(df["session_date"]).dt.strftime("%Y-%m-%d")
    df["close"] = df["close"].astype(float)  # close is stored as Decimal in the parquet

    pivot = df.pivot(index="session_date", columns="asset_id", values="close").sort_index()
    dates = list(dates)
    pivot = pivot.reindex(dates).ffill()  # forward-fill in case a name has no print on a holiday

    # Equal weight at the first agent-date (= bar 0).
    first_close = pivot.iloc[0]
    qty = (initial_value / len(universe)) / first_close  # asset_id -> shares
    bnh_values = (pivot * qty).sum(axis=1).tolist()
    return bnh_values


def replay_reward_components(run_dir: Path) -> list[dict]:
    """Recompute the new convex composite reward + per-component breakdown
    for every advance_day in a run. Uses the live ``compute_composite_reward``
    so drift in the reward module is reflected.
    """
    import sys

    if str(PROJECT_ROOT) not in sys.path:
        sys.path.insert(0, str(PROJECT_ROOT))
    if str(PROJECT_ROOT / "src") not in sys.path:
        sys.path.insert(0, str(PROJECT_ROOT / "src"))
    from tradebench.rewards.composite import compute_composite_reward  # type: ignore

    rows = load_advance_day_rows(run_dir)
    if not rows:
        return []
    manifest = episode_manifest()
    initial_value = float(manifest.get("initial_cash", 100_000))

    dates = [r["current_date"] for r in rows]
    bnh_values = equal_weight_bnh_series(dates, initial_value=initial_value)

    breakdowns: list[dict] = []
    high_water = initial_value
    prev_v = initial_value
    cum_log_bench = 0.0
    recent_alphas: list[float] = []
    prev_bnh = bnh_values[0]

    for i, row in enumerate(rows):
        v_after = float(row["portfolio_value"])
        if v_after <= 0:
            v_after = 1.0
        bnh_t = bnh_values[i]
        bar_log_agent = math.log(v_after / prev_v) if prev_v > 0 else 0.0
        bar_log_bench = math.log(bnh_t / prev_bnh) if prev_bnh > 0 else 0.0
        cum_log_agent = math.log(v_after / initial_value)
        cum_log_bench = math.log(bnh_t / initial_value)

        bar_alpha = bar_log_agent - bar_log_bench
        recent_alphas.append(bar_alpha)
        if len(recent_alphas) > 20:
            recent_alphas.pop(0)

        breakdown = compute_composite_reward(
            value_after=v_after,
            initial_value=initial_value,
            cumulative_log_return=cum_log_agent,
            cumulative_benchmark_log_return=cum_log_bench,
            recent_bar_alphas=tuple(recent_alphas),
            high_watermark=high_water,
            turnover_ratio=0.05,  # approximation; trajectory doesn't capture this directly
            hhi=0.20,
            gross_leverage=0.0,
            violations_rules=False,
            violations_hack=False,
        )
        d = breakdown.to_dict()
        d["bar"] = i
        d["date"] = row["current_date"]
        d["portfolio_value"] = v_after
        d["bnh_value"] = bnh_t
        d["bar_alpha"] = bar_alpha
        d["cum_alpha"] = cum_log_agent - cum_log_bench
        breakdowns.append(d)

        if v_after > high_water:
            high_water = v_after
        prev_v = v_after
        prev_bnh = bnh_t

    return breakdowns


__all__ = [
    "ARTIFACTS_ROOT",
    "DATASETS_ROOT",
    "MODEL_RUNS",
    "MODEL_LABELS",
    "PROJECT_ROOT",
    "IterRow",
    "load_model_iters",
    "load_advance_day_rows",
    "load_action_counts",
    "load_prompt",
    "episode_manifest",
    "equal_weight_bnh_series",
    "replay_reward_components",
]