"""CPU batch inference contract for organizer-run model evaluation."""
import argparse

import numpy as np
import pandas as pd

from model import Model, validate_rows


def predict_file(input_path, output_path):
    frame = pd.read_parquet(input_path)
    if list(frame.columns) != ['id', 'smiles', 'sequence']:
        raise ValueError('Expected input columns: id, smiles, sequence')
    model = Model()
    predictions = []
    for start in range(0, len(frame), 128):
        rows = validate_rows(frame.iloc[start:start + 128].to_dict('records'))
        predictions.extend(model.predict(rows))
    if not np.isfinite(predictions).all():
        raise ValueError('Model returned non-finite predictions')
    pd.DataFrame({'id': frame.id, 'prediction': predictions}).to_csv(output_path, index=False)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    predict_file(args.input, args.output)