import argparse
import os
import json


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('map', metavar='MappingFile', type=str, help="Mapping JSON file.")
    parser.add_argument('src', metavar='SourceFile', type=str, help="Results of span finder.")
    parser.add_argument('dst', metavar='Destination', type=str, help="Output path.")
    args = parser.parse_args()
    assert os.path.exists(args.map), "Mapping file doesn't exist."
    assert os.path.exists(args.src), "Rouce file not found."

    k_raw = json.load(open(args.map))
    k_map = dict()
    for kairos_event, content in k_raw.items():
        for fr in content['framenet']:
            if fr['label'] in k_map:
                print("Duplicate frame: " + fr['label'])
            k_map[fr['label']] = kairos_event
    inputs = list(map(json.loads, open(args.src).readlines()))

    n_total = n_mapped = 0

    for line in inputs:
        new_frames = list()
        n_total += len(line['prediction'])
        for fr in line['prediction']:
            if fr['label'] in k_map:
                fr['label'] = k_map[fr['label']]
                new_frames.append(fr)
                n_mapped += 1
        line['prediction'] = new_frames

    with open(args.dst, 'w') as fp:
        fp.write('\n'.join(map(json.dumps, inputs)))

    print(f'Done. Among {n_total} frames, {n_mapped} are mapped to KAIROS ontology, others are omitted.')


if __name__ == '__main__':
    main()