import os 
from sys import platform
if platform == "linux" or platform == "linux2":
    # linux
    import apriltag
elif platform == "darwin":
    # OS X
    import apriltag
elif platform == "win32":
    # Windows
    import pupil_apriltags as apriltag

import cv2
import matplotlib.pyplot as plt 
import pandas as pd

visualize = True

def process(file):

    VIDEO = file
    VIDEOOUT = VIDEO.split("/")[-1].split(".")[0]
    ROOT = "/".join(VIDEO.split("/")[:-1]) + "/"
    TMP_DIR = "/".join(VIDEO.split("/")[:-2]) + "/temp/"
    FRAMES = "%s%s_frames" % (TMP_DIR, VIDEOOUT)

    if not os.path.exists(FRAMES):
        print('WARNING: Could not find frame directory')
        return

    img_paths = [f for f in os.listdir(FRAMES) if 'jpg' in f]
    print('Number of frames: ', len(img_paths))

    if platform == "linux" or platform == "linux2" or platform == "darwin":
        # Circumvent error: too many borders in contour_detect (max of 32767!)
        options = apriltag.DetectorOptions(refine_edges=False, quad_contours=False)
        detector = apriltag.Detector(options)
    elif platform == "win32":
        print('WARNING: apriltag2 not supported on windows, running with pupil_apriltags...')
        detector = apriltag.Detector(refine_edges=False)

    detections = {}

    if visualize:
        fig = plt.Figure(figsize=(15, 10))
        path = os.path.join(FRAMES, img_paths[0])

        img = cv2.imread(path)
        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        result = detector.detect(img)

        for i in range(len(result)):
            tf = result[i].tag_family
            tag_id = result[i].tag_id
            cx, cy = result[i].center
            # print('Found tag: ', tag_id)
            img = cv2.circle(img, (int(cx), int(cy)), 50, (255, 255, 0), thickness=10)

        plt.imshow(image)
        plt.axis('off')
        #plt.savefig('./AprilTag_Detection_%s.jpg' % VIDEOOUT)
        plt.show()

    tags = dict()
    for frame, p in enumerate(img_paths):

        path = os.path.join(FRAMES, p)

        img = cv2.imread(path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        result = detector.detect(img)

        for i in range(len(result)):
            tf = result[i].tag_family
            cx, cy = result[i].center
            tag_id = result[i].tag_id
            tags[tag_id] = [cx, cy]

        print('Frame %i found %i tags' % (frame, len(result)))
        detections[frame] = tags

    df = pd.DataFrame.from_dict(detections, orient='index')

    path = './AprilTag_%s.pkl' % VIDEOOUT
    df.to_pickle(path)
    print('Saved AprilTag detections to %s' % path)


if __name__ == '__main__':
    process('./Data/ShowCase_3.mp4')