import cv2
import pathlib
import time
import argparse
import os
import glob

# Construct the argument parser.
parser = argparse.ArgumentParser()
parser.add_argument('-i',
                    '--input',
                    help='path to input video',
                    default='input/video_1.mp4')
parser.add_argument('-t',
                    '--template',
                    help='path to the template',
                    default='input/video_1_template.jpg')
args = vars(parser.parse_args())

print(args['template'].split(os.path.sep)[-1])

res = []

# Get all folders in the input folder.
folder = args['template']

sub_folders = [
    name for name in os.listdir(folder)
    if os.path.isdir(os.path.join(folder, name))
]

print(sub_folders)

for f in sub_folders:

    # Read the video input.
    cap = cv2.VideoCapture(args['input'])
    if (cap.isOpened() == False):
        print('Error while trying to read video. Please check path again')
    # Get the frame width and height
    frame_width = int(cap.get(3))
    frame_height = int(cap.get(4))
    # String name with which to save the resulting video.
    save_name = str(pathlib.Path(args['input'])).split(
        os.path.sep)[-1].split('.')[0]
    # define codec and create VideoWriter object
    out = cv2.VideoWriter(f"outputs/{save_name}.mp4",
                        cv2.VideoWriter_fourcc(*'mp4v'), 30,
                        (frame_width, frame_height))

    player_images = [[file, cv2.imread(file, 0)]
                    for file in glob.glob('{}/*'.format(args['template'] + '/' + str(f)))]

    for i in player_images:
        print(i[0], i[0].split(os.path.sep)[-2].lstrip('0'))
        # Read the template in grayscale format.
        template = i[1]
        w, h = template.shape[::-1]
        frame_count = 0  # To count total frames.
        total_fps = 0  # To get the final frames per second.

        # get player id
        # player_id = int(i[0].split(os.path.sep)[-2].lstrip('0'))
        player_id = int(f.lstrip('0'))

        # get the frame number from filename
        number = int(i[0].split('_')[-2].lstrip("0"))
        print(f'Player {player_id}', f'Frame {number}')

        # Capture each frame of the video.
        cap.set(1, number)
        ret, frame = cap.read()
        if ret:
            image = frame.copy()
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            start_time = time.time()
            # Apply template Matching.
            result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
            end_time = time.time()

            # Get the current fps.
            fps = 1 / (end_time - start_time)
            # Add `fps` to `total_fps`.
            total_fps += fps
            # Increment frame count.
            frame_count += 1
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            # Top left x and y coordinates.
            x1, y1 = max_loc
            # Bottom right x and y coordinates.
            x2, y2 = (x1 + w, y1 + h)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
            cv2.putText(frame,
                        f"{fps:.1f} FPS", (15, 25),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1, (0, 255, 0),
                        2,
                        lineType=cv2.LINE_AA)
            cv2.imshow('Result', frame)
            out.write(frame)


            # TrackEval result.
            temp = str(number) + ' ' + str(player_id) + ' ' + str(x1) + ' ' + str(y1) + ' ' + str(x2) + ' ' + str(y2) + ' -1, -1, -1, -1'
            res.append(temp)
            print(temp)

            # Press `q` to exit.
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        else:
            break

with open('{}.txt'.format(args['template']), 'w') as f:
    for line in res:
        f.write(f"{line}\n")


# Release VideoCapture() object.
cap.release()
# Close all frames and video windows.
cv2.destroyAllWindows()

# python main.py --input input/video_1.mp4 --template input/video_1_template.jpg

# Input: A directory containing the players & its template images.
# python main.py --input "C:\Users\LuckyVerma\Documents\TemplateMatching\data\2000_plays\plays\run_right_seah_49s113.mp4" -t data\reid-v1\successful\labels-only\run_right_seah_49s113