import argparse
import json
import os

from assignments import get_corrections

from utils import load_assignments, load_clusters

def apply_assignments(clusters, assignments_list):
    result = []
    new_clusters = clusters
    for assignments in assignments_list:
        _, rejected = get_corrections(assignments)

        cluster_elems_by_master_id = {}

        for master_id, slave_id in rejected:
            if master_id not in cluster_elems_by_master_id:
                cluster_elems_by_master_id[master_id] = []
            cluster_elems_by_master_id[master_id].append(slave_id)
            for i in range(len(new_clusters)):
                new_clusters[i] = [elem for elem in new_clusters[i] if elem != slave_id]
        result += new_clusters
        new_clusters = [elems for master_id, elems in cluster_elems_by_master_id.items()]
        print(len(new_clusters))
    if len(new_clusters) != 0:
        result += new_clusters
    return result

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 main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--clusters',
        required=True
    )
    parser.add_argument(
        '--assignments',
        required=True
    )
    parser.add_argument(
        '--new_clusters',
        required=True
    )
    args = parser.parse_args()

    clusters = load_clusters(args.clusters)
    print('Loaded {} clusters from {}'.format(len(clusters), args.clusters))
    assignments_list = []
    for name in sorted(os.listdir(args.assignments)):
        assignments_list.append(load_assignments(os.path.join(args.assignments, name)))
    print('Loaded {} assignments from folder {}'.format(len(assignments_list), args.assignments))

    new_clusters = apply_assignments(clusters, assignments_list)
    print('Generated {} new clusters'.format(len(new_clusters)))

    save_clusters(new_clusters, args.new_clusters)

if __name__ == '__main__':
    main()
