import yaml
import cv2
import os
import shutil
import deeplabcut


def generatetempyaml(yamlfile,videolist):
    try:
        #copy yaml and rename
        tempyaml = os.path.dirname(yamlfile) +'\\temp.yaml'
        shutil.copy(yamlfile,tempyaml)

        #adding new videos to tempyaml
        deeplabcut.add_new_videos(tempyaml,[videolist],copy_videos=True)

        with open(tempyaml) as f:
            read_yaml = yaml.load(f, Loader=yaml.FullLoader)

        original_videosets = read_yaml['video_sets'].keys()

        keys=[]
        for i in original_videosets:
            keys.append(i)
        #remove the original video to get only newly added videos
        read_yaml['video_sets'].pop(keys[0],None)

        with open(tempyaml, 'w') as outfile:
            yaml.dump(read_yaml, outfile, default_flow_style=False)

        print('temp.yaml generated.')
    except FileNotFoundError:
        print('Please select a video file.')

def generatetempyaml_multi(yamlfile,videolist):

    #copy yaml and rename
    tempyaml = os.path.dirname(yamlfile) +'\\temp.yaml'
    shutil.copy(yamlfile,tempyaml)


    deeplabcut.add_new_videos(tempyaml,videolist,copy_videos=True)

    with open(tempyaml) as f:
        read_yaml = yaml.load(f, Loader=yaml.FullLoader)

    original_videosets = read_yaml['video_sets'].keys()

    keys=[]
    for i in original_videosets:
        keys.append(i)


    read_yaml['video_sets'].pop(keys[0],None)

    with open(tempyaml, 'w') as outfile:
        yaml.dump(read_yaml, outfile, default_flow_style=False)


def updateiteration(yamlfile,iteration):
    yamlPath = yamlfile
    with open(yamlPath) as f:
        read_yaml = yaml.load(f, Loader=yaml.FullLoader)

    read_yaml["iteration"] = int(iteration)

    with open(yamlPath, 'w') as outfile:
        yaml.dump(read_yaml, outfile, default_flow_style=False)
    print('Iteration set to',iteration)

def update_init_weight(yamlfile,initweights):
    yamlPath=yamlfile
    initweights,initw_filetype = os.path.splitext(initweights)

    with open(yamlPath) as f:
        read_yaml = yaml.load(f, Loader=yaml.FullLoader)

    iteration = read_yaml['iteration']

    yamlfiledirectory = os.path.dirname(yamlfile)
    iterationfolder = yamlfiledirectory +'\\dlc-models\\iteration-' +str(iteration)
    projectfolder = os.listdir(iterationfolder)
    projectfolder = projectfolder[0]


    posecfg = iterationfolder + '\\' + projectfolder +'\\train\\' + 'pose_cfg.yaml'

    with open(posecfg) as g:
        read_cfg = yaml.load(g, Loader=yaml.FullLoader)

    read_cfg['init_weights'] = str(initweights)

    with open(posecfg, 'w') as outfile:
        yaml.dump(read_cfg, outfile, default_flow_style=False)
    print(os.path.basename(initweights),'selected')

def select_numfram2pick(yamlfile,numframe):
    try:
        yamlPath = yamlfile
        with open(yamlPath) as f:
            read_yaml = yaml.load(f, Loader=yaml.FullLoader)

        read_yaml["numframes2pick"] = int(numframe)

        with open(yamlPath, 'w') as outfile:
            yaml.dump(read_yaml, outfile, default_flow_style=False)
    except:
        print('Please load .yaml file and enter the number of frames to pick to proceed')
def add_multi_video_yaml(yamlfile,directory):
    filesFound = []

    ########### FIND FILES ###########
    for i in os.listdir(directory):
        if i.__contains__(".mp4"):
            a=os.path.join(directory,i)
            print(a)
            filesFound.append(a)


    videoFilePathsToAdd = filesFound
    yamlPath = yamlfile
    cropLineList = []

    with open(yamlPath) as f:
        read_yaml = yaml.load(f, Loader=yaml.FullLoader)

    ####################### GET DIMENSIONS #############################
    for i in videoFilePathsToAdd:
        cap = cv2.VideoCapture(i)
        width = int(cap.get(3))  # float
        height = int(cap.get(4))  # float
        cropLine = [0, width, 0, height]
        cropLineList.append(cropLine)


    for i in range(len(videoFilePathsToAdd)):
        currVidPath = videoFilePathsToAdd[i]
        currCropLinePath = str(cropLineList[i])
        currCropLinePath = currCropLinePath.strip("[]")
        currCropLinePath = currCropLinePath.replace("'","")
        read_yaml["video_sets"].update({currVidPath: {'crop': currCropLinePath}})

    with open(yamlPath, 'w') as outfile:
        yaml.dump(read_yaml, outfile, default_flow_style=False)

    print(len(filesFound),' videos added to config.yaml')

def add_single_video_yaml(yamlfile,videofile):
    yamlPath = yamlfile
    cap = cv2.VideoCapture(videofile)
    width = int(cap.get(3))  # float
    height = int(cap.get(4))  # float
    cropLine = [0, width, 0, height]
    cropLine = str(cropLine)
    currCropLinePath = cropLine.strip("[]")
    currCropLinePath = currCropLinePath.replace("'", "")
    with open(yamlPath) as f:
        read_yaml = yaml.load(f, Loader=yaml.FullLoader)

    read_yaml["video_sets"].update({videofile: {'crop': currCropLinePath}})

    with open(yamlPath, 'w') as outfile:
        yaml.dump(read_yaml, outfile, default_flow_style=False)


def copycsv(inifile,source):
    source = str(source)+'\\'
    dest = str(os.path.dirname(inifile))
    dest1 = str((dest)+ '\\' + 'csv' + '\\' + 'input_csv')
    files = []
    print(dest1)
    print(source)
    ########### FIND FILES ###########
    for i in os.listdir(source):
        if i.__contains__(".csv"):
            files.append(i)

    for f in files:
        try:
            shutil.copy(source+f,dest1)
            print(f,'copied to',dest1)
        except:
            print(f,'already exist in',dest1)


    print('DONE!',len(files),'copied to',dest1)