from src.surprise import calculate_surprise
import json
import pandas as pd

if __name__ == "__main__":
    input_file = 'data/test_data_output2.json'
    # read the data from the output file
    with open(input_file, 'r') as json_file:
        processed_data = json.load(json_file)
    new_processed_data = []
    for item in processed_data:
        r_obs = item['r_val']
        r_pred = item['predicted_coef']
        sigma = item['predicted_std']
        surprise = calculate_surprise(r_obs, r_pred, sigma)
        item['surprise'] = round(float(surprise), 2)
        new_processed_data.append(item)
    # write the data to the output file
    with open('data/test_data_output_w_surprise2.json', 'w') as json_file:
        json.dump(new_processed_data, json_file, indent=4)
    # save the data to a csv file
    df = pd.DataFrame(new_processed_data)
    df.to_csv('data/test_data_output2.csv', index=False)

