"""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()
    model.prepare(frame.sequence.to_numpy())
    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')
    predictions = np.asarray(predictions)
    if predictions.shape != (len(frame), 2):
        raise ValueError('Model must return two logged predictions per input')
    output = pd.DataFrame(predictions, columns=['log10_kcat', 'log10_km'])
    output.insert(0, 'id', frame.id.to_numpy())
    output.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)