import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn



import os
import sys
import numpy as np
import time
import yaml
import cv2 
import matplotlib.pyplot as plt

"""
Author: Jenö Faist, Paul Judis 
"""

"""
With this File you can manual can mark a video of a 3D print to make a fail detection dataset for your printer
you can use this dataset to mix it with other Dataset like we did with our mixed Dataset for domain shift.
"""

"""
Depending on the Video you must crop it so there a no black lines in it.
You can find out the cropping by just using try and error.
(TODO) IMPLEMENT ADAPTIVE CROPING 
"""
def crop_frame(image):
    cromped_image = image[50:1080-200,550:1920-600,:]
    return cromped_image


if __name__ == '__main__':
    """
    Saving every frame would resulte in huge dataset with pretty reduandend frames so use this variable to 
    specifice how many frames should be skiped before saving the image in the dataset
    """

    frame_save_skip = 30

    """
    Specifice here the Directory of the new Dataset and what Video you using.
    """
    absolutepath = os.path.dirname(__file__)
    dataset_PATH = absolutepath+'/DATASETS/own_SET'
    video_PATH = absolutepath+'/TRAIN_VIDEOS/rot_fail.mp4'

    print("STARTING DATASET CREATER FOR 3D PRINT ERROR DETECTION")
    print("--------------------")
    print("Video PATH:         "+video_PATH)
    print("DATASET PATH:       "+dataset_PATH)
    print("--------------------")
    print("Loading Model...")

    """
    Making a Name for the new Files 
    """

    temp = video_PATH.replace(".", "_")
    Images_Name = temp.split('/')[-1].replace("mp4","")

    print("Loading VIDEO ...")
    """
    Searching for the Video if the Video cant be openend or found close the programm
    """
    cap = None
    try:
        cap = cv2.VideoCapture(video_PATH)
        print("[✓] Video Loaded [✓]")
    except:
        print("[!?] Video couldn't be found [?!]")
        exit(1)


    loop_var = True
    count = 0

    while(loop_var):    

        ret, frame = cap.read()        
        frame = crop_frame(frame)
        video_frame = cv2.resize(frame, (250, 250)) 
        cv2.imshow('PRESS "F" TO MARK AS FAIL',video_frame)
        k =cv2.waitKey(20)
        FAILED_PRINT = False

        if(k==ord("f")):
            FAILED_PRINT = True

        if(count % frame_save_skip == 0):
            if FAILED_PRINT:
                cv2.imwrite(dataset_PATH+"/1/"+Images_Name+str(int(count/frame_save_skip))+".png", frame)
                sys.stdout.write("\033[K")
                print("SAVE AS FAILED PRINT!!! : "+str(int(count/frame_save_skip)), end='\r')
            else:
                cv2.imwrite(dataset_PATH+"/0/"+Images_Name+str(int(count/frame_save_skip))+".png", frame)
                sys.stdout.write("\033[K")
                print("SAVE AS GOOD PRINT!!! : "+str(int(count/30)), end='\r')

        count +=1

        if(k == ord('q')):
            loop_var = False  


sys.stdout.write("\033[K")       
print("MAKE DATASET FOR VIDEO "+ video_PATH + " DONE ! "+"MARKED "+str(int(count/frame_save_skip))+" FRAMES !", end='\r')