import argparse
import json

import yt.wrapper as yt
from feature import Feature, Size, Orientation
from detection import Box, Detection
from utils import load_features, load_detections, load_clusters, select_base_feature_detection_id
from task import Task

def generate_tasks(feature_by_id, detection_by_feature_detection_id, clusters):
    tasks = []
    for cluster in clusters:
        base_feature_detection_id = select_base_feature_detection_id(
            cluster,
            feature_by_id, detection_by_feature_detection_id
        )

        cluster_tasks = []
        for feature_detection_id in cluster:
            if feature_detection_id == base_feature_detection_id:
                continue
            cluster_tasks.append((base_feature_detection_id, feature_detection_id))

        for feature_detection_id1, feature_detection_id2 in cluster_tasks:
            feature_id1, detection_id1 = feature_detection_id1
            detection1 = detection_by_feature_detection_id[feature_detection_id1]
            feature1 = feature_by_id[feature_id1]

            feature_id2, detection_id2 = feature_detection_id2
            detection2 = detection_by_feature_detection_id[feature_detection_id2]
            feature2 = feature_by_id[feature_id2]

            tasks.append(Task(detection1, feature1, detection2, feature2).to_json())

    return tasks

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--yt-proxy',
        dest = 'yt_proxy',
        default = 'hahn',
        help = 'name of yt proxy'
    )
    parser.add_argument(
        '--feature-table',
        dest='feature_table',
        required=True
    )
    parser.add_argument(
        '--detections',
        required=True
    )
    parser.add_argument(
        '--clusters',
        required=True
    )
    parser.add_argument(
        '--tasks',
        required=True
    )
    args = parser.parse_args()

    yt_client = yt.YtClient(proxy=args.yt_proxy)

    feature_by_id = load_features(yt_client, args.feature_table)
    print('Loaded {} features from {}'.format(len(feature_by_id), args.feature_table))
    detection_by_feature_detection_id = load_detections(args.detections)
    print('Loaded {} detections from {}'.format(len(detection_by_feature_detection_id), args.detections))
    clusters = load_clusters(args.clusters)
    print('Loaded {} clusters from {}'.format(len(clusters), args.clusters))

    tasks = generate_tasks(feature_by_id, detection_by_feature_detection_id, clusters)
    print('Generated {} tasks'.format(len(tasks)))

    with open(args.tasks, 'w') as f:
        json.dump(tasks, f, indent=2)

if __name__ == '__main__':
    main()
