"""
Source and credits: https://github.com/ZubinGou/math-evaluation-harness/blob/main/python_executor.py
"""
import argparse
import json
from concurrent.futures import TimeoutError

import numpy as np
from pebble import ProcessPool
from tqdm import tqdm

from grader import math_equal_process


def evaluate(samples: list=None, file_path: str=None):
    assert samples or file_path, "samples or file_path must be provided"
    if not samples:
        with open(file_path, 'r') as f:
            samples = [json.loads(line) for line in f]

    # dedup by idx
    if 'idx' in samples[0]:
        samples = {sample['idx']: sample for sample in samples}.values()
        samples = sorted(samples, key=lambda x: x['idx']) 
    else:
        samples = [dict(idx=idx, **sample) for idx, sample in enumerate(samples)]

    params = [(idx, sample['pred'], sample['gt']) for idx, sample in enumerate(samples)]

    scores = []
    timeout_cnt = 0 

    with ProcessPool() as pool:
        future = pool.map(math_equal_process, params, timeout=3)
        iterator = future.result()
        with tqdm(total=len(samples), desc="Evaluate") as progress_bar:
            while True:
                try:
                    result = next(iterator)
                    scores.append(result)
                except StopIteration:
                    break
                except TimeoutError as error:
                    print(error)
                    scores.append(False)
                    timeout_cnt += 1
                except Exception as error:
                    print(error.traceback)
                    exit()
                progress_bar.update(1) 

    assert len(samples) == len(scores)

    for i in range(len(samples)):
        samples[i]['score'] = scores[i]

    mean_score = np.round(np.mean([score for score in scores if score is not False]), decimals=2)

    result_json = {
        "num_samples": len(samples),
        "num_scores": len(scores),
        "timeout_samples": timeout_cnt,
        "acc": mean_score
    }

    return samples, result_json


if __name__ == "__main__":
    samples, results_json = evaluate(file_path="output/MATH.jsonl")
    print('test')