From protein sequence,
to real molecules,
to discovery, in weeks.

from pathlib import Path
from uuid import uuid4

import polars as pl
from omtx import OmClient

JAK2_V617F_SEQUENCE = "YOUR_JAK2_V617F_PROTEIN_SEQUENCE"
client = OmClient(api_key="YOUR_API_KEY")

job = client.lula2.score(
    protein_sequence=JAK2_V617F_SEQUENCE,
    source="om",
    tier=50,
    n=50_000,
    top_k=10_000,
    idempotency_key="jak2-v617f-lula2-r1",
)

artifact_paths = []
result_dir = Path("outputs/jak2-v617f-lula2-r1")

for job_id in job["job_ids"]:
    client.jobs.wait(job_id, poll_interval=5, timeout=3600)
    artifact_paths.extend(
        client.jobs.download_all_artifacts(
            job_id,
            output_dir=result_dir / job_id,
            overwrite=True,
        )
    )

score_tables = [
    pl.read_parquet(path)
    for path in artifact_paths
    if path.suffix == ".parquet"
]

score_rows = pl.concat(score_tables).sort("score", descending=True)
selected_hits = score_rows.head(100).to_dicts()

addresses = client.molecules.shipping_addresses()

order = client.molecules.order(
    items=selected_hits,
    shipping_address_id=addresses["default_shipping_address_id"],
    idempotency_key=f"jak2-v617f-round-1-{uuid4()}",
)

print(order["order_number"])