import argparse
import json
from collections import defaultdict
import pandas


def update_summary(summary, res_doc):
    for key in res_doc:
        summary[key].append(res_doc[key])
    return summary

def get_params():
    parser = argparse.ArgumentParser(description='tomita_lstm')
    parser.add_argument('--result_file', type=str, required = True)
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    session_args = get_params()
    with open(session_args.result_file, "r+") as ff:
        data = json.load(ff)

    summary = defaultdict(list)
    for sample in data:
        summary["seed"].append(sample["seed"])
        if "n_parameters" in sample:
            summary["n_params"].append(sample["n_parameters"])
        elif "n_states" in sample:
            summary["n_states"].append(sample["n_states"])
        summary = update_summary(summary, sample["result"])
    
    pd = pandas.DataFrame.from_dict(summary)

    print("=====================================================================")
    print("File : ", session_args.result_file)
    print("=====================================================================")

    print("-------------------------- Results ----------------------------------")
    print(pd.to_markdown())    

    new = pd[filter( lambda x : False if x in ["seed", "n_params", "n_states"] else True, list(summary.keys()))].copy()
    
    print("--------------------------- Mean -----------------------------------")
    print(new.mean())

    print("--------------------------- Std ------------------------------------")
    print(new.std()) 

