import argparse
import json
import itertools

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
from coordinates import distance, heading_abs_diff

def generate_tasks(feature_by_id, detection_by_feature_detection_id, clusters, max_distance, max_heading_diff):
    base_feature_detection_id_by_cluster_indx = {}
    for i, cluster in enumerate(clusters):
        base_feature_detection_id_by_cluster_indx[i] = select_base_feature_detection_id(
            cluster,
            feature_by_id, detection_by_feature_detection_id
        )

    clusters_indx_by_type = {}
    for i, cluster in enumerate(clusters):
        detection = detection_by_feature_detection_id[cluster[0]]
        sign_type = detection.sign_type
        if sign_type not in clusters_indx_by_type:
            clusters_indx_by_type[sign_type] = []
        clusters_indx_by_type[sign_type].append(i)

    tasks = []
    for sign_type, clusters_indx in clusters_indx_by_type.items():
        for indx1, indx2 in itertools.combinations(clusters_indx, 2):
            maybe_one_cluster = False
            for cluster1_feature_detection_id in clusters[indx1]:
                if maybe_one_cluster:
                    break
                cluster1_feature_id = cluster1_feature_detection_id[0]
                cluster1_feature = feature_by_id[cluster1_feature_id]
                for cluster2_feature_detection_id in clusters[indx2]:
                    cluster2_feature_id = cluster2_feature_detection_id[0]
                    cluster2_feature = feature_by_id[cluster2_feature_id]
                    if distance(cluster1_feature.pos, cluster2_feature.pos) > max_distance:
                        continue
                    if heading_abs_diff(cluster1_feature.heading, cluster2_feature.heading) > max_heading_diff:
                        continue
                    maybe_one_cluster = True
                    break
            if not maybe_one_cluster:
                continue

            feature_detection_id1 = base_feature_detection_id_by_cluster_indx[indx1]
            feature_id1, detection_id1 = feature_detection_id1
            feature1 = feature_by_id[feature_id1]
            detection1 = detection_by_feature_detection_id[feature_detection_id1]

            feature_detection_id2 = base_feature_detection_id_by_cluster_indx[indx2]
            feature_id2, detection_id2 = feature_detection_id2
            feature2 = feature_by_id[feature_id2]
            detection2 = detection_by_feature_detection_id[feature_detection_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(
        '--distance',
        type = float,
        required = True
    )
    parser.add_argument(
        '--heading-diff',
        type = float,
        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, args.distance, args.heading_diff)
    print('Generated {} tasks'.format(len(tasks)))

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

if __name__ == '__main__':
    main()
