#!/usr/bin/env python3

import argparse
import openem
import os
import cv2
import numpy as np
from openem.tracking import *
import json
import sys
import datetime
import tator
from pprint import pprint
from collections import defaultdict

import yaml
import math

import subprocess
import sys

def crop_localization(frame_bgr, localization):
    img_width = frame_bgr.shape[1]
    img_height = frame_bgr.shape[0]
    box_x = round(localization['x'] * img_width)
    box_y = round(localization['y'] * img_height)
    box_width = round(localization['width'] * img_width)
    box_height = round(localization['height'] * img_height)
    img_crop = frame_bgr[box_y:box_y+box_height,box_x:box_x+box_width,:]
    return img_crop

def join_up_iteration(detections, track_ids):
            tracklets = defaultdict(list)
            num_tracklets = np.max(track_ids) + 1
            assert(len(detections) == len(track_ids))
            for d,tid in zip(detections, track_ids):
                tracklets[tid].append(d)
            return tracklets

def extend_tracklets(tracklets, length):
    for track_id,track in tracklets.items():
        if len(track) <= 16:
            continue

        ext_length = min(length,len(track))
        sum_h=0.0
        sum_w=0.0

        track.sort(key=lambda x:x['frame'])

        def restore_det(det):
            det['x'] = det.get('orig_x',det['x'])
            det['y'] = det.get('orig_y',det['y'])
            det['width'] = det.get('orig_w',det['width'])
            det['height'] = det.get('orig_h',det['height'])
            det['orig_x'] = det['x']
            det['orig_y'] = det['y']
            det['orig_w'] = det['width']
            det['orig_h'] = det['height']
        restore_det(track[0])
        restore_det(track[-1])

        for d in track:
            sum_h += d['height']
            sum_w += d['width']
        angle,vel,comps = track_vel(track)
        vel_x = comps[0]
        vel_y = comps[1]
        avg_h = sum_h / len(track)
        avg_w = sum_w / len(track)
        new_x = min(1,max(0,track[-1]['x']+(vel_x*ext_length)))
        new_y = min(1,max(0,track[-1]['y']+(vel_y*ext_length)))
        old_x = min(1,max(0,track[0]['x']-(vel_x*ext_length)))
        old_y = min(1,max(0,track[0]['y']-(vel_y*ext_length)))


        min_x = min(track[-1]['x'],new_x)
        min_y = min(track[-1]['y'],new_y)
        if min_x > 0 and min_y > 0:
            track[-1]['x'] = min_x
            track[-1]['y'] = min_y
            track[-1]['width'] = min(max(0,abs(new_x-track[-1]['x'])+avg_w),1)
            track[-1]['height'] = min(max(0,abs(new_x-track[-1]['y'])+avg_h),1)
        else:
            track[-1]['width'] = 0
            track[-1]['height'] = 0


        min_x = min(track[0]['x'],old_x)
        min_y = min(track[0]['y'],old_y)
        if min_x > 0 and min_y > 0:
            track[0]['x'] = min(max(0,min_x),1)
            track[0]['y'] = min(max(0,min_y),1)
            track[0]['width'] = min(max(abs(old_x-track[0]['x'])+avg_w,0),1)
            track[0]['height'] = min(max(abs(old_x-track[0]['y'])+avg_h,0),1)
        else:
            track[0]['width'] = 0
            track[0]['height'] = 0
    return tracklets


def split_tracklets(tracklets):
    track_ids=[]
    detections=[]
    for track_id,track in tracklets.items():
        for d in track:
            track_ids.append(track_id)
            detections.append(d)
    return detections,track_ids

def trim_tracklets(detections, track_ids, max_length):
    tracklets = join_up_iteration(detections, track_ids)
    next_track_id = 1
    new_tracklets = {}
    for track_id,detections in tracklets.items():
        new_track_count=math.ceil(len(detections)/max_length)
        for i in range(new_track_count):
            start=max_length*i
            end=max_length+(max_length*i)
            new_tracklets[next_track_id] = detections[start:end]
            next_track_id += 1
    detections, track_ids = split_tracklets(new_tracklets)
    track_ids = renumber_track_ids(track_ids)
    return detections, track_ids


if __name__=="__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    tator.get_parser(parser)
    parser.add_argument("--detection-type-id", type=int, required=True)
    parser.add_argument("--tracklet-type-id", type=int, required=True)
    parser.add_argument("--version-id", type=int)
    parser.add_argument("--input-version-id", type=int)
    parser.add_argument("--strategy-config", type=str)
    parser.add_argument("--dry-run", action='store_true')
    parser.add_argument('media_files', type=str, nargs='*')
    args = parser.parse_args()

    # Weight methods
    methods = ['hybrid', 'iou', 'iou-motion', 'iou-global-motion']

    # Weight methods that require the video
    visual_methods = ['hybrid', 'iou-global-motion']

    api = tator.get_api(args.host, args.token)
    detection_type = api.get_localization_type(args.detection_type_id)
    project = detection_type.project
    version_id = args.version_id


    default_strategy = {"method": "hybrid",
                        "frame-diffs": [1,2,4,8,16,32,64,128,256],
                        "args": {},
                        "extension": {'method' : None},
                        "max-length": {},
                        "min-length": 0}

    if args.strategy_config:
        strategy = {**default_strategy}
        with open(args.strategy_config, "r") as strategy_file:
            strategy.update(yaml.load(strategy_file))
    else:
        strategy = default_strategy

    if strategy['method'] == 'hybrid':
        model_file = strategy['args']['model_file']
        batch_size = strategy['args'].get('batch_size', 4)
        comparator=FeaturesComparator(model_file)
        #extractor=FeaturesExtractor(args.model_file)


    class_method = strategy.get('class-method',None)
    classify_function = None
    classify_args = {}
    if class_method:
        pip_package=class_method.get('pip',None)
        if pip_package:
            p = subprocess.run([sys.executable,
                                "-m",
                                "pip",
                                "install",
                                pip_package])
            print("Finished process.", flush=True)
        function_name = class_method.get('function',None)
        classify_args = class_method.get('args',None)
        names = function_name.split('.')
        module = __import__(names[0])
        for name in names[1:-1]:
            module = getattr(module,name)
        classify_function = getattr(module,names[-1])


    print("Strategy: ", flush=True)
    pprint(strategy)
    print(args.media_files, flush=True)

    optional_fetch_args = {}
    if args.input_version_id:
        optional_fetch_args['version'] = [args.input_version_id]
    for media_file in args.media_files:
        comps=os.path.splitext(os.path.basename(media_file))[0]
        media_id=comps.split('_')[0]
        media = api.get_media(media_id)
        if media.attributes.get("Tracklet Generator Processed") != "No":
            print(f"Skipping media ID {media.id}, name {media.name} due to "
                  f"'Tracklet Generator Processed' attribute being set to "
                  f"something other than 'No'!")
            continue
        media_shape = (media.height, media.width)
        fps = media.fps

        localizations_by_frame = {}
        localizations = api.get_localization_list(project,
                                                  type=args.detection_type_id,
                                                  media_id=[media_id],
                                                  **optional_fetch_args)

        localizations = [l.to_dict() for l in localizations]
        if len(localizations) == 0:
            print(f"No localizations present in media {media_file}", flush=True)
            continue
        print(f"Processing {len(localizations)} detections", flush=True)
        # Group by localizations by frame
        for lid, local in enumerate(localizations):
            frame = local['frame']
            if frame in localizations_by_frame:
                localizations_by_frame[frame].append(local)
            else:
                localizations_by_frame[frame] = [local]

        detections=[]
        track_ids=[]
        track_id=1

        # If media does not exist, download it.
        if strategy['method'] == 'iou-global-motion':
            if not os.path.exists(media_file):
                temp_path = f'/tmp/{os.path.basename(media_file)}'
                for progress in tator.util.download_media(api, media, temp_path):
                    print(f"Downloading {media_file}, {progress}%...")
                print("Download finished!")
                # Unfrag the file
                subprocess.run(["ffmpeg", '-i', temp_path, '-c:v', 'copy', media_file])
                os.remove(temp_path)

        if strategy['method'] == 'hybrid': # Not all visual methods need detection images
            vid=cv2.VideoCapture(media_file)
            ok=True
            frame = 0
            while ok:
                ok,frame_bgr = vid.read()
                if frame in localizations_by_frame:
                    for l in localizations_by_frame[frame]:
                        l['bgr'] = crop_localization(frame_bgr, l)
                        if l['attributes']['Confidence'] < 0.50:
                            continue
                        detections.append(l)
                        track_ids.append(track_id)
                        track_id += 1
                frame+=1
        else:
            # The method is analytical on the detections coordinates
            # and does not require processing the video
            for frame,frame_detections in localizations_by_frame.items():
                for det in frame_detections:
                    detections.append(det)
                    track_ids.append(track_id)
                    track_id += 1

        print("Loaded all detections", flush=True)

        track_ids = renumber_track_ids(track_ids)

        if strategy['method'] == 'hybrid':
            weights_strategy = HybridWeights(comparator,
                                             None,
                                             None,
                                             media_shape,
                                             fps,
                                             0.0,
                                             batch_size)
        elif strategy['method'] == 'iou':
            weights_strategy = IoUWeights(media_shape, **strategy['args'])
        elif strategy['method'] == 'iou-motion':
            weights_strategy = IoUMotionWeights(media_shape, **strategy['args'])
        elif strategy['method'] == 'iou-global-motion':
            weights_strategy = IoUGlobalMotionWeights(media_shape, media_file, **strategy['args'])
        # Generate localization bgr based on grouped localizations
        for x in strategy['frame-diffs']:
            print(f"Started {x}", flush=True)
            detections, track_ids, pairs, weights, is_cut, constraints = join_tracklets(
                detections,
                track_ids,
                x,
                weights_strategy)

            if x in strategy['max-length']:
                trim_to = strategy['max-length'][x]
                print(f"Trimming track to max length of {trim_to}")
                detections, track_ids = trim_tracklets(detections, track_ids, trim_to)
            _,det_counts_per_track=np.unique(track_ids,return_counts=True)
            print(f"frame-diff {x}: {len(detections)} to {len(det_counts_per_track)}", flush=True)

            if x > 1 and strategy['extension']['method'] == 'linear-motion':
                ext_frames=x
                print(f"Extending by linear motion, {ext_frames}")
                tracklets = join_up_iteration(detections,track_ids)
                tracklets = extend_tracklets(tracklets, ext_frames)
                detections, track_ids = split_tracklets(tracklets)

        # Now we make new track objects based on the result
        # from the graph solver
        # [ detection, detection, detection, ...]
        # [ track#, track#, track#,...]
        # [ 133, 33, 13, 133,]
        # [ 0,0,1,1]
        # TODO: Handle is_cut?
        def join_up_final(detections, track_ids):
            tracklets = defaultdict(list)
            num_tracklets = np.max(track_ids) + 1
            assert(len(detections) == len(track_ids))
            for d,tid in zip(detections, track_ids):
                tracklets[tid].append(d)
            return tracklets

        def make_object(track):
            track.sort(key=lambda x:x['frame'])
            if classify_function:
                valid,attrs = classify_function(media.to_dict(),
                                                track,
                                                **classify_args)
            elif len(track) >= strategy['min-length']:
                valid = True
                attrs = {}
            else:
                valid = False
                attrs = {}
            if valid:
                obj={"type": args.tracklet_type_id,
                     "media_ids": [int(media_id)],
                     "localization_ids": [x['id'] for x in track],
                     **attrs,
                     "version": version_id}
                return obj
            else:
                return None

        tracklets = join_up_final(detections, track_ids)
        new_objs=[make_object(tracklet) for tracklet in tracklets.values()]
        new_objs=[x for x in new_objs if x is not None]
        print(f"New objects = {len(new_objs)}")
        with open(f"/work/{media_id}.json", "w") as f:
            json.dump(new_objs,f)
        if not args.dry_run:
            for response in tator.util.chunked_create(api.create_state_list,project,
                                                      state_spec=new_objs):
                pass
            try:
                api.update_media(int(media_id), {"attributes":{"Tracklet Generator Processed": str(datetime.datetime.now())}})
            except:
                print("WARNING: Unable to set 'Tracklet Generator Processed' attribute")

[((33, 24, 33, 41), 'collections.defaultdict', 'defaultdict', ({(33, 36, 33, 40): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((128, 13, 128, 57), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((129, 4, 129, 28), 'tator.get_parser', 'tator.get_parser', ({(129, 21, 129, 27): 'parser'}, {}), '(parser)', False, 'import tator\n'), ((145, 10, 145, 46), 'tator.get_api', 'tator.get_api', ({(145, 24, 145, 33): 'args.host', (145, 35, 145, 45): 'args.token'}, {}), '(args.host, args.token)', False, 'import tator\n'), ((194, 4, 194, 20), 'pprint.pprint', 'pprint', ({(194, 11, 194, 19): 'strategy'}, {}), '(strategy)', False, 'from pprint import pprint\n'), ((34, 28, 34, 45), 'numpy.max', 'np.max', ({(34, 35, 34, 44): 'track_ids'}, {}), '(track_ids)', True, 'import numpy as np\n'), ((178, 16, 182, 45), 'subprocess.run', 'subprocess.run', ({(178, 31, 182, 44): "[sys.executable, '-m', 'pip', 'install', pip_package]"}, {}), "([sys.executable, '-m', 'pip', 'install', pip_package])", False, 'import subprocess\n'), ((247, 16, 247, 44), 'cv2.VideoCapture', 'cv2.VideoCapture', ({(247, 33, 247, 43): 'media_file'}, {}), '(media_file)', False, 'import cv2\n'), ((301, 35, 301, 74), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((319, 24, 319, 41), 'collections.defaultdict', 'defaultdict', ({(319, 36, 319, 40): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((353, 12, 353, 33), 'json.dump', 'json.dump', ({(353, 22, 353, 30): 'new_objs', (353, 31, 353, 32): 'f'}, {}), '(new_objs, f)', False, 'import json\n'), ((355, 28, 356, 74), 'tator.util.chunked_create', 'tator.util.chunked_create', (), '', False, 'import tator\n'), ((161, 28, 161, 52), 'yaml.load', 'yaml.load', ({(161, 38, 161, 51): 'strategy_file'}, {}), '(strategy_file)', False, 'import yaml\n'), ((201, 31, 201, 59), 'os.path.basename', 'os.path.basename', ({(201, 48, 201, 58): 'media_file'}, {}), '(media_file)', False, 'import os\n'), ((237, 19, 237, 45), 'os.path.exists', 'os.path.exists', ({(237, 34, 237, 44): 'media_file'}, {}), '(media_file)', False, 'import os\n'), ((239, 32, 239, 80), 'tator.util.download_media', 'tator.util.download_media', ({(239, 58, 239, 61): 'api', (239, 63, 239, 68): 'media', (239, 70, 239, 79): 'temp_path'}, {}), '(api, media, temp_path)', False, 'import tator\n'), ((243, 16, 243, 87), 'subprocess.run', 'subprocess.run', ({(243, 31, 243, 86): "['ffmpeg', '-i', temp_path, '-c:v', 'copy', media_file]"}, {}), "(['ffmpeg', '-i', temp_path, '-c:v', 'copy', media_file])", False, 'import subprocess\n'), ((244, 16, 244, 36), 'os.remove', 'os.remove', ({(244, 26, 244, 35): 'temp_path'}, {}), '(temp_path)', False, 'import os\n'), ((320, 28, 320, 45), 'numpy.max', 'np.max', ({(320, 35, 320, 44): 'track_ids'}, {}), '(track_ids)', True, 'import numpy as np\n'), ((238, 36, 238, 64), 'os.path.basename', 'os.path.basename', ({(238, 53, 238, 63): 'media_file'}, {}), '(media_file)', False, 'import os\n'), ((359, 99, 359, 122), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n')]