import numpy as np
import matplotlib.pyplot as plt
import image_marker as im
import tracking_result as tr
import os
import glob
import argparse
from IPython import embed

#1. Tankkoordinaten
def tankcoordinates(video, dontask=False):
    redo = True
    if os.path.exists('tankcoordinates.py'):
        from tankcoordinates import bottom_left, bottom_right, top_left, top_right
        print("Found tank coordinates top left: %s, top right: %s" % (top_left, top_right))
        if dontask:
            return bottom_left, top_left, top_right, bottom_right
        answer = input('Redo markers? y/n')
        if answer == 'y' or answer == 'Y':
            redo = True
        else:
            redo = False
    if redo:
        tank_task = im.MarkerTask("tank limits", ["bottom left corner", "top left corner", "top right corner", "bottom right corner"], "Mark tank corners")
        image_marker = im.ImageMarker([tank_task])
        marker_positions = image_marker.mark_movie(video, 100)
        bottom_right = marker_positions[0]['bottom right corner']
        bottom_left = marker_positions[0]['bottom left corner']
        top_right = marker_positions[0]['top right corner']
        top_left = marker_positions[0]['top left corner']
        with open('tankcoordinates.py', 'w') as f:
            f.write('bottom_left = %s\n' % str(marker_positions[0]['bottom left corner']))
            f.write('top_left = %s\n' % str(marker_positions[0]['top left corner']))
            f.write('top_right = %s\n' % str(marker_positions[0]['top right corner']))
            f.write('bottom_right = %s\n' % str(marker_positions[0]['bottom right corner']))

    return bottom_left, top_left, top_right, bottom_right
#2. Feederkoordinaten

#3. dark_light Koordinaten
def dark_light_coordinates(video, dontask=False):
    redo = True
    if os.path.exists('dark_light_coordinates.py'):
        from dark_light_coordinates import left, right, dark_center
        print("Found dark_light_coordinates left: %s, right: %s, dark_center: %s" % (left, right, dark_center))
        if dontask:
            return left, right, dark_center
        answer = input('Redo markers? y/n')
        if answer == 'y' or answer == 'Y':
            redo = True
        else:
            redo = False
    if redo:    
        dark_light_task = im.MarkerTask('Dark side', ['left', 'right', 'dark_center'], 'Mark light dark separator line')
        image_marker = im.ImageMarker([dark_light_task])
        marker_positions = image_marker.mark_movie(video, 100)
        
        right = tr.coordinate_transformation(marker_positions[0]['right']) 
        left = tr.coordinate_transformation(marker_positions[0]['left'])
        dark_center = tr.coordinate_transformation(marker_positions[0]['dark_center'])

        with open('dark_light_coordinates.py', 'w') as f:
            f.write('left = %s\n' % str(left))
            f.write('right = %s\n' % str(right))
            f.write('dark_center = %s\n' % str(dark_center))

    return left, right, dark_center

#4. Laden der Trackingresults
def load_tracking_results(dlc_results_file):
    trs = tr.TrackingResult(dlc_results_file)
    t, x, y, l, name = trs.position_values(bodypart="snout", framerate=30)
    return t, x, y, l

#5. Wie lange hält sich der Fisch im Hellen/Dunklen auf? 
    # Anzahl Frames in der Fisch in definiertem, dunklen Bereich ist, bzw. in der der Fisch nicht im Hellen ist
def aufenthaltsort(left, center, fish_y):
    top_is_dark=left[1]>=center[1]
    #bei wie vielen Frames ist der Fisch im Hellen?
    if top_is_dark: 
        hell_count = len(fish_y[fish_y >= left[1]])
    else:
        hell_count = len(fish_y[fish_y < left[1]])
    dark_count=len(fish_y) - hell_count
    total_count = len(fish_y)
    return hell_count, dark_count, total_count


def analysiere_video(v, dlc, left, center):
    t, fish_x, fish_y, likelihood = load_tracking_results(dlc)
    hc, dc, tc = aufenthaltsort(left, center, fish_y)
    print('Der Fisch hat sich %.2f %% im Dunklen aufgehalten.'%(dc/tc*100))
    return (dc/tc*100)


def main():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("day", type=str, help="The day you want to work on")
    parser.add_argument("-f", "--folder", type=str, default="/data/boldness/labeled_videos", help="The base folder in which the labeled videos are stored. Default is /data/boldness/labeled_videos")
    parser.add_argument("-a", "--animal", type=str, default="*", help="The animal id, default is * for all animals")
    parser.add_argument("-e", "--extension", type=str, default=".mp4", help="The video file extension, default is .mp4")
    parser.add_argument("-na", "--noask", action="store_true", help="do not ask for coordinates")
    args = parser.parse_args()
    
    videos = sorted(glob.glob(os.path.join(args.folder, args.day, '*%s*%s' % (args.animal, args.extension))))
    dlc_files = sorted(glob.glob(os.path.join(args.folder, args.day, '*%s*%s' % (args.animal, '.h5'))))
    results = {}
    if len(videos) > 0:
        left, right, center = dark_light_coordinates(videos[0], args.noask)
        # bl, tl, tr, br = tankcoordinates(v, args.noask)
        for video, dlc_file in zip(videos, dlc_files):
            animal = video.split(os.sep)[-1].split('_')[1]
            p_dark = analysiere_video(video, dlc_file, left, center)
            results[animal] = p_dark
        np.save('results_%s.npy' % args.day, results)
    print(results)
    

if __name__ == '__main__':
    main()