#!/usr/bin/env python3
"""Reproduce the quality-review ID selection and curve values, not assay judgments.

Usage: python selection_script.txt REFERENCE_DIR OUTPUT_DIR
Dependencies: pandas, pyarrow. Source data: OpenADMET, CC-BY-4.0.
No network access is used. This script is package-authored, not upstream code.
"""
import json
import sys
from pathlib import Path
import pandas as pd

source, output = map(Path, sys.argv[1:])
output.mkdir(parents=True, exist_ok=True)
train = pd.read_csv(source / 'training_metadata.csv')
curves = pd.read_parquet(source / 'original_curves.parquet')
eligible = train[train['Molecule Name'].isin(set(curves['Molecule Name']))].copy()
eligible = eligible.sort_values(['pEC50_std.error (-log10(molarity))', 'Molecule Name'], kind='stable')
ranks = [round(i * (len(eligible) - 1) / 15) for i in range(16)]
selected = eligible.iloc[ranks].copy()
selected.insert(0, 'selection_rank_zero_based', ranks)
selected.to_csv(output / 'selection.csv', index=False)
selected_ids = set(selected['Molecule Name'])
records = []
for index, row in curves.iterrows():
    if row['Molecule Name'] not in selected_ids:
        continue
    record = {'source_row_zero_based': int(index)}
    for key, value in row.items():
        record[key] = value.tolist() if hasattr(value, 'tolist') else int(value) if key == 'npoints' else value
    records.append(record)
(output / 'curve_values.json').write_text(json.dumps(records, indent=2, allow_nan=False) + '\n')
print(json.dumps({'selected_ids': sorted(selected_ids), 'curve_rows': len(records)}))