import argparse
import json
import os

from utils import load_clusters

def save_clusters(clusters, path):
    data = []
    for i, cluster in enumerate(clusters):
        detections = []
        for feature_id, detection_id in cluster:
            detections.append({'feature_id': feature_id, 'object_id': detection_id})
        data.append({'cluster_id': i, 'objects': detections})
    with open(path, 'w') as f:
        json.dump({'clusters': data}, f, indent=2)


def merge_clusters(clusters_list):
    used_foids = set()
    clusters = []
    for clusters_item in clusters_list:
        for cluster_item in clusters_item:
            for foid in cluster_item:
                if foid in used_foids:
                    raise Exception('{} is already used in other cluster'.format(foid))
                used_foids.add(foid)
        clusters += clusters_item
    return clusters


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--clusters-dir',
        required=True
    )
    parser.add_argument(
        '--clusters',
        required=True
    )
    args = parser.parse_args()

    print('Loading clusters list from {}'.format(args.clusters_dir))
    clusters_list = []
    for name in os.listdir(args.clusters_dir):
        clusters_list.append(load_clusters(os.path.join(args.clusters_dir, name)))

    clusters = merge_clusters(clusters_list)
    print('Generated {} clusters'.format(len(clusters)))
    save_clusters(clusters, args.clusters)

if __name__ == '__main__':
    main()
