brain_client.py

"""
WorldQuant Brain API Client Library

A Python client for interacting with the WorldQuant Brain platform API.
Supports authentication, alpha simulation, and result retrieval.
"""

import json
import getpass
import requests
from time import sleep
from os.path import expanduser
from typing import Optional


BASE_URL = "https://api.worldquantbrain.com"


class BrainClient:
    """Client for the WorldQuant Brain API."""

    def __init__(
        self,
        email: Optional[str] = None,
        password: Optional[str] = None,
        credentials_file: Optional[str] = None,
    ):
        """
        Initialize the client. Credentials priority:
          1. email + password passed directly
          2. credentials_file path (JSON: ["email", "password"])
          3. ~/.brain_credentials (if the file exists)
          4. Interactive prompt (asks at runtime)

        Args:
            email: Brain platform email.
            password: Brain platform password.
            credentials_file: Path to a JSON credentials file.
        """
        self._session = requests.Session()
        self._authenticated = False

        if email and password:
            self._session.auth = (email, password)
            return

        if credentials_file:
            with open(expanduser(credentials_file), "r") as f:
                self._session.auth = tuple(json.load(f))
            return

        default_file = expanduser("~/.brain_credentials")
        try:
            with open(default_file, "r") as f:
                self._session.auth = tuple(json.load(f))
        except FileNotFoundError:
            # Fall back to interactive prompt
            email = input("WorldQuant Brain email: ").strip()
            password = getpass.getpass("WorldQuant Brain password: ")
            self._session.auth = (email, password)

    @classmethod
    def login(cls, email: Optional[str] = None, password: Optional[str] = None) -> "BrainClient":
        """
        Create a BrainClient and authenticate in one step.

        If email/password are not provided, prompts interactively.

        Args:
            email: Brain platform email (optional).
            password: Brain platform password (optional).

        Returns:
            An authenticated BrainClient instance.

        Example:
            client = BrainClient.login()                          # interactive prompt
            client = BrainClient.login("me@email.com", "pass")   # direct
        """
        client = cls(email=email, password=password)
        client.authenticate()
        return client

    # -------------------------------------------------------------------------
    # Authentication
    # -------------------------------------------------------------------------

    def authenticate(self) -> dict:
        """
        Sign in and obtain a JWT token.

        Returns:
            Response JSON from the /authentication endpoint.

        Raises:
            requests.HTTPError: If authentication fails.
        """
        response = self._session.post(f"{BASE_URL}/authentication")
        response.raise_for_status()
        self._authenticated = True
        return response.json()

    # -------------------------------------------------------------------------
    # Simulations
    # -------------------------------------------------------------------------

    def simulate(self, expression: str, settings: dict = None, alpha_type: str = "REGULAR", regular: str = "close") -> "SimulationResult":
        """
        Submit an alpha expression for simulation.

        Args:
            expression: The alpha expression string.
            settings: Simulation settings dict. Defaults to standard equity settings.
            alpha_type: Type of simulation, e.g. "REGULAR".
            regular: Price field, e.g. "close".

        Returns:
            A SimulationResult object to poll for completion.

        Raises:
            requests.HTTPError: If the submission fails.
        """
        default_settings = {
            "instrumentType": "EQUITY",
            "region": "USA",
            "universe": "TOP3000",
            "delay": 1,
            "decay": 15,
            "neutralization": "SUBINDUSTRY",
            "truncation": 0.08,
            "maxTrade": "ON",
            "pasteurization": "ON",
            "testPeriod": "P1Y6M",
            "unitHandling": "VERIFY",
            "nanHandling": "OFF",
            "language": "FASTEXPR",
            "visualization": False,
        }
        if settings:
            default_settings.update(settings)

        payload = {
            "type": alpha_type,
            "settings": default_settings,
            "regular": expression if regular == "close" else regular,
        }

        # If expression is provided separately from regular field
        if regular == "close":
            payload["regular"] = expression

        response = self._session.post(f"{BASE_URL}/simulations", json=payload)
        response.raise_for_status()

        progress_url = response.headers.get("Location")
        return SimulationResult(self._session, progress_url)

    # -------------------------------------------------------------------------
    # Alphas
    # -------------------------------------------------------------------------

    def get_alpha(self, alpha_id: str) -> dict:
        """
        Retrieve details of a completed alpha by ID.

        Args:
            alpha_id: The alpha ID string.

        Returns:
            Alpha details as a dict.
        """
        response = self._session.get(f"{BASE_URL}/alphas/{alpha_id}")
        response.raise_for_status()
        return response.json()

    def get_pnl(self, alpha_id: str, poll_interval: float = 5.0) -> dict:
        """
        Retrieve PnL record set for an alpha, polling until ready.

        Args:
            alpha_id: The alpha ID string.
            poll_interval: Fallback seconds to wait between polls if no Retry-After header.

        Returns:
            PnL record set as a dict.
        """
        return self._poll_recordset(alpha_id, "pnl", poll_interval)

    def get_recordset(self, alpha_id: str, record_set_name: str, poll_interval: float = 5.0) -> dict:
        """
        Retrieve any named record set for an alpha, polling until ready.

        Args:
            alpha_id: The alpha ID string.
            record_set_name: Name of the record set (e.g. "pnl", "sharpe").
            poll_interval: Fallback seconds between polls if no Retry-After header.

        Returns:
            Record set data as a dict.
        """
        return self._poll_recordset(alpha_id, record_set_name, poll_interval)

    def _poll_recordset(self, alpha_id: str, record_set_name: str, poll_interval: float) -> dict:
        url = f"{BASE_URL}/alphas/{alpha_id}/recordsets/{record_set_name}"
        while True:
            response = self._session.get(url)
            retry_after = float(response.headers.get("Retry-After", 0))
            if retry_after == 0:
                response.raise_for_status()
                return response.json()
            print(f"Sleeping for {retry_after} seconds...")
            sleep(retry_after)


class SimulationResult:
    """Represents a pending or completed alpha simulation."""

    def __init__(self, session: requests.Session, progress_url: str):
        self._session = session
        self.progress_url = progress_url
        self.alpha_id: str = None
        self._result: dict = None

    def wait(self, verbose: bool = True) -> dict:
        """
        Poll until the simulation completes.

        Args:
            verbose: If True, print polling status messages.

        Returns:
            The completed simulation result JSON containing the alpha id.
        """
        while True:
            response = self._session.get(self.progress_url)
            retry_after = float(response.headers.get("Retry-After", 0))
            if retry_after == 0:
                response.raise_for_status()
                self._result = response.json()
                self.alpha_id = self._result.get("alpha")
                if verbose:
                    print(f"Alpha simulation complete. Alpha ID: {self.alpha_id}")
                return self._result
            if verbose:
                print(f"Simulating... sleeping for {retry_after} seconds")
            sleep(retry_after)

    def get_alpha(self) -> dict:
        """
        Fetch full alpha details after simulation completes.

        Returns:
            Alpha details dict.

        Raises:
            RuntimeError: If wait() has not been called yet.
        """
        if not self.alpha_id:
            raise RuntimeError("Simulation not complete. Call wait() first.")
        response = self._session.get(f"{BASE_URL}/alphas/{self.alpha_id}")
        response.raise_for_status()
        return response.json()

    def get_pnl(self, poll_interval: float = 5.0) -> dict:
        """
        Retrieve PnL record set after simulation completes.

        Args:
            poll_interval: Fallback seconds between polls if no Retry-After header.

        Returns:
            PnL record set dict.

        Raises:
            RuntimeError: If wait() has not been called yet.
        """
        if not self.alpha_id:
            raise RuntimeError("Simulation not complete. Call wait() first.")
        client = BrainClient.__new__(BrainClient)
        client._session = self._session
        client._authenticated = True
        return client.get_pnl(self.alpha_id, poll_interval)