from collections import defaultdict


def convert_idsw_to_heatmap_format(filepath, dest_file):
    """Convert idsw format to heatmap formats
    idsw: <frame> <id1_gt> <id1> <bb1_left> <bb1_top> <bb1_width> <bb1_height> <id2_gt> <id2> <bb2_left> <bb2_top> <bb2_width> ...
    heatmap:  <frame> <bb1_left> <bb1_top> <bb1_width> <bb1_height> <bb2_left> <bb2_top> <bb2_width> <bb2_height> ...

    Args:
        filepath ([str]): idsw path
        dest_file : save file path
    """

    ids_group = defaultdict(list)
    obj_infos = []

    with open(filepath, "r") as f:
        for line in f:
            # <frame> <id1_gt> <id1> <bb1_left> <bb1_top> <bb1_width> <bb1_height> <id2_gt> <id2> <bb2_left> <bb2_top> <bb2_width>
            p = line.rstrip().split(" ")
            p = list(map(int, p))
            # get number of objects in current frame 
            num_obj = int((len(p) - 1) / 6)
            for idx in range(num_obj):
                # <frame> <id1_gt> <id1> <bb1_left> <bb1_top> <bb1_width> <bb1_height>\n <frame> <id2_gt> <id2> <bb2_left> <bb2_top> <bb2_width>
                obj_infos.append([p[0]] + p[1 + 6 * idx: 7 + 6 * idx])

    for obj in obj_infos:
        ids_group[obj[1]].append(obj)

    # print(ids_group)
    with open(dest_file, 'w') as f:
        for _, objs in ids_group.items():
            objs = sorted(objs, key=lambda x: x[0])
            if len(objs) % 2 == 0:
                for i in range(1, len(objs), 2):
                    tmp = list(map(str, objs[i]))
                    # print(tmp)
                    line = tmp[0] + " " + tmp[3] + " " + tmp[4] + " " + tmp[5] + " " + tmp[6] + "\n"
                    f.write(line)

    return 1