source/project/data/close_distribution_v2.py

"""Prepare multi-asset, horizon-masked close data for forecasting-v4."""

from __future__ import annotations

import argparse
import hashlib
import json
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any

import numpy as np


ADAPTER_VERSION = "close-distribution-v2.0"
REPO_ID = "tmmycruise/autoresearch-market-data"
REVISION = "93ec44a9918e42fb5900af1690753d3bf175709e"
SELECTION_ASSETS = ("AAPL", "ABBV", "MCD")
CONFIRMATION_ASSETS = ("AMD", "MA")
ALL_ASSETS = SELECTION_ASSETS + CONFIRMATION_ASSETS
DEVELOPMENT_START = date(2016, 8, 8)
RESEARCH_END = date(2022, 8, 8)
SEALED_START = date(2023, 8, 8)
SEALED_END = date(2026, 8, 8)
EVALUATION_HORIZONS = np.arange(2, 33, dtype=np.int16)
CANONICAL_HORIZONS = np.asarray((2, 4, 8, 16, 32), dtype=np.int16)
CLASS_COUNT = 21
MINUTE_NS = 60_000_000_000
BASE_TOKEN_FEATURE_NAMES = (
    "close_log_return_since_last_observed_close",
    "close_return_observed",
    "log1p_elapsed_wall_clock_minutes_since_last_observed_close",
)
CANDIDATE_TOKEN_FEATURE_NAMES = (
    "regular_session_progress",
    "regular_session_progress_sin",
    "regular_session_progress_cos",
    "log1p_minutes_since_previous_session_last_observed_close",
    "log_raw_close",
    "close_log_return_since_session_first_observed",
    "close_log_return_since_previous_session_last_observed",
)
TOKEN_FEATURE_NAMES = BASE_TOKEN_FEATURE_NAMES + CANDIDATE_TOKEN_FEATURE_NAMES
DISCOVERY_FOLDS = {
    "discovery_1": {
        "train": ("2016-08-08", "2020-08-08"),
        "validation": ("2020-08-08", "2021-08-08"),
    },
    "discovery_2": {
        "train": ("2016-08-08", "2021-08-08"),
        "validation": ("2021-08-08", "2022-08-08"),
    },
}


@dataclass(frozen=True)
class Session:
    day: date
    open_ns: int
    close_ns: int


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _stable_hash(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
    return hashlib.sha256(payload).hexdigest()


def _save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, **arrays)


def _timestamp_ns(value: Any) -> int:
    if value.tzinfo is None:
        raise ValueError("session timestamp must be timezone-aware")
    return int(value.timestamp() * 1_000_000_000)


def _read_sessions(path: Path) -> list[Session]:
    import pyarrow.dataset as ds

    table = ds.dataset(str(path), format="parquet").to_table(
        columns=[
            "session_date",
            "open_utc",
            "close_utc",
            "regular_minutes",
        ],
        filter=(
            (ds.field("session_date") >= DEVELOPMENT_START)
            & (ds.field("session_date") < RESEARCH_END)
            & (ds.field("calendar") == "XNYS")
        ),
    )
    sessions: list[Session] = []
    for row in table.to_pylist():
        open_ns = _timestamp_ns(row["open_utc"])
        close_ns = _timestamp_ns(row["close_utc"])
        expected = (close_ns - open_ns) // MINUTE_NS
        if expected != int(row["regular_minutes"]):
            raise ValueError(
                f"calendar duration mismatch on {row['session_date']}"
            )
        if expected < 180 or expected > 390:
            raise ValueError(
                f"invalid regular-session length on {row['session_date']}"
            )
        sessions.append(
            Session(
                day=row["session_date"],
                open_ns=open_ns,
                close_ns=close_ns,
            )
        )
    sessions.sort(key=lambda item: item.day)
    if not sessions or sessions[-1].day >= RESEARCH_END:
        raise ValueError("session input crossed the research boundary")
    return sessions


def _read_bars(path: Path, ticker: str) -> dict[str, np.ndarray]:
    import pyarrow.compute as pc
    import pyarrow.dataset as ds

    table = ds.dataset(
        str(path),
        format="parquet",
        partitioning=None,
    ).to_table(
        columns=[
            "ticker",
            "window_start_ns",
            "close",
            "source_date",
            "adjusted",
        ],
        filter=(
            (ds.field("ticker") == ticker)
            & (ds.field("source_date") >= DEVELOPMENT_START)
            & (ds.field("source_date") < RESEARCH_END)
        ),
    )
    if table.num_rows == 0:
        raise ValueError(f"no research bars found for {ticker}")
    if pc.any(table["adjusted"]).as_py():
        raise ValueError(f"{ticker} source unexpectedly contains adjusted rows")
    if pc.max(table["source_date"]).as_py() >= RESEARCH_END:
        raise ValueError(f"{ticker} bars crossed the research boundary")
    arrays = {
        name: table[name].combine_chunks().to_numpy(zero_copy_only=False)
        for name in ("window_start_ns", "close", "source_date")
    }
    order = np.argsort(arrays["window_start_ns"], kind="stable")
    arrays = {name: values[order] for name, values in arrays.items()}
    timestamp = arrays["window_start_ns"].astype(np.int64, copy=False)
    if np.any(timestamp[1:] == timestamp[:-1]):
        raise ValueError(f"{ticker} contains duplicate minute timestamps")
    close = arrays["close"].astype(np.float64, copy=False)
    if np.any(~np.isfinite(close)) or np.any(close <= 0.0):
        raise ValueError(f"{ticker} contains invalid closes")
    return arrays


def _read_splits(path: Path, ticker: str) -> list[dict[str, Any]]:
    import pyarrow.dataset as ds

    table = ds.dataset(str(path), format="parquet").to_table(
        columns=[
            "id",
            "execution_date",
            "split_from",
            "split_to",
        ],
        filter=(
            (ds.field("ticker") == ticker)
            & (ds.field("execution_date") >= DEVELOPMENT_START)
            & (ds.field("execution_date") < RESEARCH_END)
        ),
    )
    result = sorted(
        table.to_pylist(),
        key=lambda row: (row["execution_date"], str(row["id"])),
    )
    for row in result:
        if float(row["split_from"]) <= 0 or float(row["split_to"]) <= 0:
            raise ValueError(f"invalid {ticker} split event: {row}")
    return result


def _split_adjusted_close(
    close: np.ndarray,
    source_date: np.ndarray,
    splits: list[dict[str, Any]],
) -> np.ndarray:
    adjusted = np.asarray(close, dtype=np.float64).copy()
    for event in splits:
        factor = float(event["split_from"]) / float(event["split_to"])
        adjusted[source_date < event["execution_date"]] *= factor
    if np.any(~np.isfinite(adjusted)) or np.any(adjusted <= 0.0):
        raise ValueError("split adjustment produced invalid closes")
    return adjusted


def _align_sessions(
    sessions: list[Session],
    bars: dict[str, np.ndarray],
    splits: list[dict[str, Any]],
) -> dict[str, np.ndarray]:
    timestamp = bars["window_start_ns"].astype(np.int64, copy=False)
    adjusted_close = _split_adjusted_close(
        bars["close"],
        bars["source_date"],
        splits,
    )
    parts: dict[str, list[np.ndarray]] = {
        "timestamp_ns": [],
        "close": [],
        "raw_close": [],
        "observed": [],
        "session_date": [],
        "minute_of_session": [],
        "session_length": [],
    }
    for session in sessions:
        expected = np.arange(
            session.open_ns,
            session.close_ns,
            MINUTE_NS,
            dtype=np.int64,
        )
        positions = np.searchsorted(timestamp, expected)
        matched = positions < len(timestamp)
        matched[matched] &= (
            timestamp[positions[matched]] == expected[matched]
        )
        session_close = np.full(len(expected), np.nan, dtype=np.float64)
        raw_session_close = np.full(len(expected), np.nan, dtype=np.float64)
        session_close[matched] = adjusted_close[positions[matched]]
        raw_session_close[matched] = bars["close"][positions[matched]]
        valid = matched & np.isfinite(session_close) & (session_close > 0.0)
        day = np.datetime64(session.day.isoformat(), "D").astype(np.int32)
        parts["timestamp_ns"].append(expected)
        parts["close"].append(session_close)
        parts["raw_close"].append(raw_session_close)
        parts["observed"].append(valid)
        parts["session_date"].append(
            np.full(len(expected), day, dtype=np.int32)
        )
        parts["minute_of_session"].append(
            np.arange(len(expected), dtype=np.int16)
        )
        parts["session_length"].append(
            np.full(len(expected), len(expected), dtype=np.int16)
        )
    return {
        name: np.concatenate(values)
        for name, values in parts.items()
    }


def _token_features(
    grid: dict[str, np.ndarray],
) -> tuple[np.ndarray, np.ndarray]:
    timestamp = grid["timestamp_ns"]
    close = grid["close"]
    raw_close = grid["raw_close"]
    observed = grid["observed"]
    features = np.zeros(
        (len(timestamp), len(TOKEN_FEATURE_NAMES)),
        dtype=np.float32,
    )
    volatility = np.full(len(timestamp), np.nan, dtype=np.float32)
    previous_observed = -1
    recent_returns: list[float] = []
    current_session = None
    session_first_close = np.nan
    session_last_close = np.nan
    previous_session_last_close = np.nan
    previous_session_last_timestamp = -1
    session_gap = 0.0
    for index in range(len(timestamp)):
        session = int(grid["session_date"][index])
        if current_session != session:
            if current_session is not None and np.isfinite(session_last_close):
                previous_session_last_close = session_last_close
                previous_session_last_timestamp = int(
                    timestamp[previous_observed]
                )
            current_session = session
            session_first_close = np.nan
            session_last_close = np.nan
            if previous_session_last_timestamp >= 0:
                gap_minutes = max(
                    1.0,
                    (
                        timestamp[index]
                        - previous_session_last_timestamp
                    )
                    / MINUTE_NS,
                )
                session_gap = np.log1p(gap_minutes)
            else:
                session_gap = 0.0
        denominator = max(int(grid["session_length"][index]) - 1, 1)
        progress = float(grid["minute_of_session"][index]) / denominator
        angle = 2.0 * np.pi * progress
        features[index, 3] = progress
        features[index, 4] = np.sin(angle)
        features[index, 5] = np.cos(angle)
        features[index, 6] = session_gap
        if previous_observed >= 0:
            elapsed = max(
                1.0,
                (timestamp[index] - timestamp[previous_observed]) / MINUTE_NS,
            )
            features[index, 2] = np.log1p(elapsed)
        if observed[index] and previous_observed >= 0:
            value = np.log(close[index] / close[previous_observed])
            features[index, 0] = value
            features[index, 1] = 1.0
            features[index, 7] = np.log(raw_close[index])
            if not np.isfinite(session_first_close):
                session_first_close = close[index]
            features[index, 8] = np.log(
                close[index] / session_first_close
            )
            if np.isfinite(previous_session_last_close):
                features[index, 9] = np.log(
                    close[index] / previous_session_last_close
                )
            session_last_close = close[index]
            recent_returns.append(float(value))
            if len(recent_returns) > 32:
                recent_returns.pop(0)
            if len(recent_returns) >= 8:
                volatility[index] = np.sqrt(
                    np.mean(np.square(recent_returns))
                )
            previous_observed = index
        elif observed[index]:
            features[index, 1] = 1.0
            features[index, 7] = np.log(raw_close[index])
            session_first_close = close[index]
            session_last_close = close[index]
            if np.isfinite(previous_session_last_close):
                features[index, 9] = np.log(
                    close[index] / previous_session_last_close
                )
            previous_observed = index
        elif previous_observed >= 0 and len(recent_returns) >= 8:
            volatility[index] = np.sqrt(
                np.mean(np.square(recent_returns))
            )
    if not np.all(np.isfinite(features)):
        raise ValueError("token features contain nonfinite values")
    return features, volatility


def _targets(grid: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
    close = grid["close"]
    observed = grid["observed"]
    minute = grid["minute_of_session"].astype(np.int64)
    length = grid["session_length"].astype(np.int64)
    y = np.full(
        (len(close), len(EVALUATION_HORIZONS)),
        np.nan,
        dtype=np.float64,
    )
    mask = np.zeros_like(y, dtype=np.bool_)
    for horizon_index, horizon in enumerate(
        EVALUATION_HORIZONS.astype(np.int64)
    ):
        candidate = np.arange(len(close), dtype=np.int64) + horizon
        valid = (
            observed
            & (minute + horizon < length)
            & (candidate < len(close))
        )
        rows = np.flatnonzero(valid)
        same_session = (
            grid["session_date"][candidate[rows]]
            == grid["session_date"][rows]
        )
        target_observed = observed[candidate[rows]]
        rows = rows[same_session & target_observed]
        y[rows, horizon_index] = np.log(
            close[candidate[rows]] / close[rows]
        )
        mask[rows, horizon_index] = True
    return y, mask


def _label_bundle(
    grid: dict[str, np.ndarray],
    y: np.ndarray,
    mask: np.ndarray,
) -> dict[str, np.ndarray]:
    selected = np.any(mask, axis=1)
    row_index = np.flatnonzero(selected).astype(np.int64)
    return {
        "row_index": row_index,
        "y": y[row_index],
        "target_mask": mask[row_index],
        "session_date": grid["session_date"][row_index],
        "minute_of_session": grid["minute_of_session"][row_index],
        "session_third": np.minimum(
            (
                3.0
                * grid["minute_of_session"][row_index]
                / grid["session_length"][row_index]
            ).astype(np.int8),
            2,
        ),
    }


def _input_paths(
    study_dir: Path,
    ticker: str,
    boundary: dict[str, Any],
) -> tuple[Path, Path, Path]:
    asset = boundary["assets"][ticker]
    paths = {}
    for name in ("bars", "splits"):
        record = asset["outputs"][name]
        path = study_dir / record["path"]
        if _sha256(path) != record["sha256"]:
            raise ValueError(f"{ticker} {name} source hash mismatch")
        paths[name] = path
    session_record = boundary["sessions"]
    sessions = study_dir / session_record["path"]
    if _sha256(sessions) != session_record["sha256"]:
        raise ValueError("session source hash mismatch")
    return paths["bars"], sessions, paths["splits"]


def _asset_output_dir(study_dir: Path, ticker: str) -> Path:
    if ticker in CONFIRMATION_ASSETS:
        return study_dir / "data" / "protected" / "confirmation"
    return study_dir / "data" / "runner"


def _scope_configuration(
    study_dir: Path,
    scope: str,
) -> tuple[tuple[str, ...], Path, Path, dict[str, Any] | None]:
    if scope == "selection":
        return (
            SELECTION_ASSETS,
            study_dir / "data" / "source" / "source-manifest.json",
            study_dir / "data" / "prepared-manifest.json",
            None,
        )
    if scope != "confirmation":
        raise ValueError("scope must be selection or confirmation")
    marker_path = study_dir / "confirmation" / "CONFIRMATION_OPENED.json"
    if not marker_path.exists():
        raise PermissionError(
            "confirmation preparation requires the frozen open marker"
        )
    marker = json.loads(marker_path.read_text())
    if marker.get("state") != "opened" or not marker.get("freeze_sha256"):
        raise ValueError("confirmation open marker is invalid")
    return (
        CONFIRMATION_ASSETS,
        (
            study_dir
            / "data"
            / "protected_source"
            / "confirmation-source-manifest.json"
        ),
        (
            study_dir
            / "data"
            / "protected"
            / "confirmation"
            / "prepared-manifest.json"
        ),
        marker,
    )


def prepare(
    study_dir: Path,
    *,
    scope: str = "selection",
) -> dict[str, Any]:
    assets, boundary_path, manifest_path, marker = _scope_configuration(
        study_dir,
        scope,
    )
    if not boundary_path.exists():
        raise FileNotFoundError(
            f"forecasting-v4 {scope} isolated source is missing"
        )
    boundary = json.loads(boundary_path.read_text())
    if boundary.get("research_end_exclusive") != RESEARCH_END.isoformat():
        raise ValueError("source manifest has the wrong research boundary")
    if boundary.get("revision") != REVISION:
        raise ValueError("source manifest revision differs from contract")
    if set(boundary.get("assets", {})) != set(assets):
        raise ValueError(f"source manifest has the wrong {scope} assets")

    derived_hashes: dict[str, str] = {}
    asset_metadata: dict[str, Any] = {}
    for ticker in assets:
        bars_path, sessions_path, splits_path = _input_paths(
            study_dir,
            ticker,
            boundary,
        )
        sessions = _read_sessions(sessions_path)
        bars = _read_bars(bars_path, ticker)
        splits = _read_splits(splits_path, ticker)
        grid = _align_sessions(sessions, bars, splits)
        features, volatility = _token_features(grid)
        y, target_mask = _targets(grid)
        labels = _label_bundle(grid, y, target_mask)

        output_dir = _asset_output_dir(study_dir, ticker)
        slug = ticker.lower()
        feature_path = output_dir / f"{slug}_features.npz"
        label_path = output_dir / f"{slug}_labels.npz"
        _save_npz(
            feature_path,
            {
                "X": features,
                "timestamp_ns": grid["timestamp_ns"],
                "available_at_ns": grid["timestamp_ns"] + MINUTE_NS,
                "session_date": grid["session_date"],
                "minute_of_session": grid["minute_of_session"],
                "session_length": grid["session_length"],
                "causal_volatility_32": volatility,
            },
        )
        _save_npz(label_path, labels)
        for path in (feature_path, label_path):
            relative = str(path.relative_to(study_dir))
            derived_hashes[relative] = _sha256(path)

        valid_target_rows = target_mask.sum(axis=0).astype(np.int64)
        target_times = []
        for horizon_index, horizon in enumerate(EVALUATION_HORIZONS):
            rows = np.flatnonzero(target_mask[:, horizon_index])
            target_times.append(
                int(
                    (
                        grid["timestamp_ns"][rows]
                        + (int(horizon) + 1) * MINUTE_NS
                    ).max()
                )
            )
        asset_metadata[ticker] = {
            "confirmation": ticker in CONFIRMATION_ASSETS,
            "feature_rows": int(len(features)),
            "label_rows": int(len(labels["row_index"])),
            "maximum_anchor_available_at_ns": int(
                (
                    grid["timestamp_ns"][labels["row_index"]]
                    + MINUTE_NS
                ).max()
            ),
            "maximum_target_available_at_ns": int(max(target_times)),
            "observed_token_fraction": float(np.mean(grid["observed"])),
            "split_events": [
                {
                    "execution_date": row["execution_date"].isoformat(),
                    "id": str(row["id"]),
                    "split_from": float(row["split_from"]),
                    "split_to": float(row["split_to"]),
                }
                for row in splits
            ],
            "valid_target_rows_by_horizon": {
                str(int(horizon)): int(count)
                for horizon, count in zip(
                    EVALUATION_HORIZONS,
                    valid_target_rows,
                    strict=True,
                )
            },
        }

    sealed_ns = int(
        np.datetime64(SEALED_START.isoformat(), "ns").astype(np.int64)
    )
    research_end_ns = int(
        np.datetime64(RESEARCH_END.isoformat(), "ns").astype(np.int64)
    )
    for ticker, values in asset_metadata.items():
        for key in (
            "maximum_anchor_available_at_ns",
            "maximum_target_available_at_ns",
        ):
            if values[key] >= research_end_ns or values[key] >= sealed_ns:
                raise ValueError(
                    f"{ticker} {key} crossed the research boundary"
                )

    source_snapshot = {
        "assets": list(assets),
        "boundary": boundary,
        "development_start": DEVELOPMENT_START.isoformat(),
        "repo_id": REPO_ID,
        "research_end_exclusive": RESEARCH_END.isoformat(),
        "revision": REVISION,
        "sealed_end_exclusive": SEALED_END.isoformat(),
        "sealed_start": SEALED_START.isoformat(),
    }
    metadata = {
        "adapter_version": ADAPTER_VERSION,
        "assets": asset_metadata,
        "canonical_horizons": CANONICAL_HORIZONS.tolist(),
        "class_count": CLASS_COUNT,
        "confirmation_assets": (
            list(CONFIRMATION_ASSETS) if scope == "confirmation" else []
        ),
        "derived_sha256": derived_hashes,
        "evaluation_horizons": EVALUATION_HORIZONS.tolist(),
        "folds": DISCOVERY_FOLDS,
        "scope": scope,
        "selection_assets": (
            list(SELECTION_ASSETS) if scope == "selection" else []
        ),
        "snapshot_sha256": _stable_hash(
            {
                "source": source_snapshot,
                "derived": derived_hashes,
                "evaluation_horizons": EVALUATION_HORIZONS.tolist(),
                "canonical_horizons": CANONICAL_HORIZONS.tolist(),
                "class_count": CLASS_COUNT,
                "token_features": list(TOKEN_FEATURE_NAMES),
            }
        ),
        "source_snapshot": source_snapshot,
        "token_feature_names": list(TOKEN_FEATURE_NAMES),
    }
    if marker is not None:
        selection_manifest = json.loads(
            (study_dir / "data" / "prepared-manifest.json").read_text()
        )
        metadata["freeze_sha256"] = marker["freeze_sha256"]
        metadata["selection_snapshot_sha256"] = selection_manifest[
            "snapshot_sha256"
        ]
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(
        json.dumps(metadata, indent=2, sort_keys=True) + "\n"
    )
    return metadata


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--study-dir", type=Path, required=True)
    parser.add_argument(
        "--scope",
        choices=("selection", "confirmation"),
        default="selection",
    )
    args = parser.parse_args()
    print(json.dumps(prepare(**vars(args)), sort_keys=True))


if __name__ == "__main__":
    main()