'''
Functions for movies
'''
import numpy as np
import os, sys
from datetime import datetime as dt
import importlib
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import json
import random
import re
from PIL import Image, ImageFont, ImageDraw
import cv2
import dlib
import argparse
import imutils
import time
import csv
from imutils.video import VideoStream
import statistics
from pathlib import Path

#############################################################################
#
# Finds sublist of elements containing a substring in an input list
#
# input arguments:
# searchList - (numpy array or list of strings) input list
# subString  - (string) string to look for in each element
#
# returns:
# subList    - (numpy array) elements from input list containing substring
#
#############################################################################
def subStringInList(searchList, subString):
    if type(searchList)==list:
        searchList = np.array(searchList)
    inds = np.flatnonzero(np.core.defchararray.find(searchList,subString)!=-1)
    subList = searchList[inds]
    return subList

def nestedSqueeze(workingDict):
    keys = workingDict.keys()
    for k, key in enumerate(keys):
        if type(workingDict[key])==dict:
            workingDict.update({key:nestedSqueeze(workingDict[key])})
        elif type(workingDict[key])==list:
            sqval = np.array(workingDict[key]).squeeze()
            workingDict.update({key:sqval})

    return workingDict


def movmean(arr, window) :
    ret = np.cumsum(arr)
    ret[window:] = ret[window:] - ret[:-window]
    return ret / window

def formatTimeStr(inTime, units='s'):
    if units=='s':
        inTime = inTime*1000
    else:
        print('units must be "s" or "ms"')

    hr = inTime // 3600000
    remainingSecs = inTime % 3600000
    mins = remainingSecs // 60000
    remainingSecs = remainingSecs % 60000
    sec = remainingSecs // 1000
    ms = remainingSecs % 1000

    formatted = '%02d:%02d:%02d.%03d' % (hr, mins, sec, ms)
    return formatted

#############################################################################
#
# Load labels from csv
#
#############################################################################
def loadLabels(labelCsvFile):
    with open(labelCsvFile, 'r', newline='') as csvfile:
        r = csv.reader(csvfile, delimiter=' ', quotechar='|')
        next(r)
        for rowNum, row in enumerate(r):
            if rowNum==0:
                bodyParts = {}
                partList = row[0].split(',')
                partList = partList[1:]
                parts = np.array([])
                for part in partList:
                    if part not in parts:
                        parts = np.append(parts, part)
                        bodyParts.update({part:{}})
            elif rowNum==1:
                line = row[0].split(',')
                partIdx = 0
                valIdx = 0
                valList = np.array([])
                for val in line[1:]:
                    if val not in valList:
                        valList = np.append(valList, val)
                    bodyParts[parts[partIdx]].update({val:np.array([])})
                    valIdx += 1
                    if valIdx>2:
                        valIdx = 0
                        partIdx += 1
            else:
                data = row[0].split(',')
                valIdx = 0 #0, 1, or 2 for x, y, likelihood
                partIdx = 0

                for point in data[1:]:
                    bp = parts[partIdx]
                    val = valList[valIdx]
                    added = np.append(bodyParts[bp][val], float(point))
                    bodyParts[bp].update({val:added})
                    valIdx += 1
                    if valIdx>2:
                        valIdx = 0
                        partIdx += 1
    return bodyParts

#############################################################################
#
# Make HoldBite Movie
#
# input arguments:
# movInfo      - (dict) data to make movie (generated by matlab)
# saveToFolder - (string) path to folder where finished movies are saved
#
#############################################################################
def makeSingleVideo(imagePaths, frameTimes, saveToPath, window=None, rotate=False, codec='MJPG',ext='.avi'):
    firstIm = Image.open(imagePaths[0])

    if window == None:
        wf, hf = firstIm.size
    else:
        hf = window[3] - window[1]
        wf = window[2] - window[0]
    print('hf:', hf)
    print('wf:', wf)
    timestampPos = (10, 10)
    fontSize = int(round(hf/20))

    drive, stuff = os.path.splitdrive(Path().absolute())
    fontPath = os.path.join(drive+os.sep, 'python_code', 'Roboto', 'Roboto-Black.ttf')
    font = ImageFont.truetype(fontPath, fontSize)
    timestampCol = (255)
    # initialize the FourCC, video writer, dimensions of the frame, and FrameRate
    fps = int(np.round(1/np.mean(np.diff(frameTimes))))
    print('FrameRate:', fps)
    fourcc = cv2.VideoWriter_fourcc(*codec)
    writer = cv2.VideoWriter(saveToPath, fourcc, fps, (wf, hf), False)
    print('Writer size:', wf, hf)


    for fidx, frameTime in enumerate(frameTimes):

        imObj = Image.open(imagePaths[fidx])
        if window != None:
            imObj = imObj.crop(window)
        if rotate:
            imObj = imObj.rotate(180)
        drawObj = ImageDraw.Draw(imObj)

        ftStr = formatTimeStr(frameTime)
        textWidth, textHeight = drawObj.textsize(ftStr, font=font)
        timestampPos = (wf-textWidth, 0)
        drawObj.text(timestampPos, ftStr, fill=timestampCol, font=font)

        outputFrame = np.array(imObj).astype(np.uint8)

        cv2.imshow("OutputFrame", outputFrame)
        writer.write(outputFrame)
        print('Frame size:', outputFrame.shape)
        
        # show the frames
        key = cv2.waitKey(30)
        # if the `q` key was pressed, break from the loop
        if key == 27:
            break

    # do a bit of cleanup
    print("[INFO] cleaning up...")
    cv2.destroyAllWindows()
    writer.release()
    
#############################################################################
#
# Takes list of paths to videos and creates synced video
#
# input arguments:
# imgFolderList - (list of strings) full paths to 2 or 3 videos
# saveTo        - (string) full path to folder where movie will be saved
# movieName     - (string) name of video file to be saved
# 
#############################################################################
def syncVideos(imgFolderList,saveTo, movieName, h=480,w=640,fps=20.0,codec='XVID',ext='.avi',fontSize=25):
    #Initialize variables
    maxTimeStamp = 0
    imgFolders = [os.path.basename(imgFolder) for imgFolder in imgFolderList]
    print(imgFolders)
    dirNames = [os.path.dirname(imgFolder) for imgFolder in imgFolderList]
    vidDict = {
                imgFolder:{
                    'dirName':'',
                    'imgNames':[],
                    'frameTimes':[], 
                    'frameRates':[]
                } for imgFolder in imgFolderList}
    for imf, fullPath in enumerate(imgFolderList):
        imgFolder = os.path.basename(fullPath)
        with open(os.path.join(fullPath, 'frame_times.txt'), 'r') as f:
            frameTimeStrs = f.readlines()
        
        #print('last ftstr:', frameTimeStrs[len(frameTimeStrs)-1])
        #print('last ftstr:', frameTimeStrs[len(frameTimeStrs)-2])
        #del frameTimeStrs[len(frameTimeStrs)-1]
        #del frameTimeStrs[len(frameTimeStrs)-2]
        numFT = len(frameTimeStrs)
        frameTimes = np.zeros(numFT+1)
        for idx in range(numFT):
            #print('fts:', frameTimeStrs[idx])
            frameTimes[idx+1] = int(frameTimeStrs[idx].strip())
        frameDiff = np.diff(frameTimes)
        frameRates = np.divide(1000, frameDiff)
        frameTimes = frameTimes[1:]
        imNames = os.listdir(fullPath)
        imNames.remove('frame_times.txt')
        try:
            imNames.remove('._frame_times.txt')
        except:
            noRM = True
        vidDict[fullPath]['dirName'] = dirNames[imf]
        vidDict[fullPath]['imgNames'] = imNames
        vidDict[fullPath]['frameTimes'] = frameTimes
        vidDict[fullPath]['frameRates'] = frameRates
        vidDict[fullPath]['medfps'] = statistics.median(frameRates)
        numIms = len(imNames)
        vidDict[fullPath]['numIms'] = numIms
        maxTimeStamp = max(maxTimeStamp, frameTimes[numIms-1])
        
    
    if saveTo==None:
        movieName = folderName + ext
        saveTo = os.path.join(saveToPath, movieName)
        print(saveTo)
    else:
        movieName = movieName + ext
    


    # initialize the FourCC, video writer, dimensions of the frame, and
    # zeros array
    fourcc = cv2.VideoWriter_fourcc(*codec)
    writer = None

    (hf, wf) = (None, None)
    
    #Main for loop
    #Opens image, adds text to image, adds image to movie writer
    largeFR = 2*fps
    msperFrame = 1000/largeFR
    numFrames = int(np.ceil(maxTimeStamp/msperFrame))
    print(numFrames)
    
        
    font = ImageFont.truetype(os.path.join(eatingBoxPath, 'Roboto', 'Roboto-Black.ttf'), fontSize)
    tsfont = ImageFont.truetype(os.path.join(eatingBoxPath, 'Roboto', 'Roboto-Black.ttf'), int(np.ceil(0.5*fontSize)))
    
    #draw object for info panel
    infoImg = Image.new('RGB', (w, h))
    drawObj = ImageDraw.Draw(infoImg)
    xpos = 10
    ypos = 10
    yinc = fontSize
    text_color = (255, 255, 255)
    if '-' in imgFolder:
        splitChar = '-'
    else:
        splitChar = '_'
    txtInfo = movieName.split(splitChar)
    for line in txtInfo:
        drawObj.text((xpos, ypos), line, fill=text_color, font=font)
        ypos += yinc
    
    outputInfoObj = infoImg.convert('L')
    topLeftPanel = np.array(outputInfoObj)
    hasThirdVideo = len(imgFolderList) > 2
    if not hasThirdVideo:
        bottomLeftPanel = np.zeros((h,w))
    
    frameTime = 0
    imgFolders = list(vidDict.keys())
    
    print('numFrameTimes:', len(vidDict[imgFolders[2]]['frameTimes']))
    print('numIms:', len(vidDict[imgFolders[2]]['imgNames']))
    for fidx in range(numFrames):
        closestFrameIdx = [0, 0, 0]
        
        #look for closest frame time for top right panel
        ftarrayTR = vidDict[imgFolders[0]]['frameTimes']
        closestFrameIdxTR = (np.abs(ftarrayTR - frameTime)).argmin()
        topRightPanelName = vidDict[imgFolders[0]]['imgNames'][closestFrameIdxTR]
        topRightPanelImg =Image.open(os.path.join(vidDict[imgFolders[0]]['dirName'], imgFolders[0], topRightPanelName))
        RGBTR = topRightPanelImg.convert('RGB')
        drawObjTR = ImageDraw.Draw(RGBTR)
        ftTR = ftarrayTR[closestFrameIdxTR]
        hrTR = ftTR // 3600000
        remainingSecs = ftTR % 3600000
        minTR = remainingSecs // 60000
        remainingSecs = remainingSecs % 60000
        secTR = remainingSecs // 1000
        msTR = remainingSecs % 1000
        
        drawObjTR.text((10, 10), '%02d:%02d:%02d.%03d' % (hrTR, minTR, secTR, msTR), fill=text_color, font=tsfont)
        grayTR = RGBTR.convert('L')
        topRightPanel = np.array(grayTR)
        
        #look for closest frame time for bottom right panel
        ftarrayBR = vidDict[imgFolders[1]]['frameTimes']
        closestFrameIdxBR = (np.abs(ftarrayBR - frameTime)).argmin()
        bottomRightPanelName = vidDict[imgFolders[1]]['imgNames'][closestFrameIdxBR]
        bottomRightPanelImg =Image.open(os.path.join(vidDict[imgFolders[1]]['dirName'], imgFolders[1], bottomRightPanelName))
        RGBBR = bottomRightPanelImg.convert('RGB')
        drawObjBR = ImageDraw.Draw(RGBBR)
        ftBR = ftarrayBR[closestFrameIdxBR]
        hrBR = ftBR // 3600000
        remainingSecs = ftBR % 3600000
        minBR = remainingSecs // 60000
        remainingSecs = remainingSecs % 60000
        secBR = remainingSecs // 1000
        msBR = remainingSecs % 1000
        drawObjBR.text((10, 10), '%02d:%02d:%02d.%03d' % (hrBR, minBR, secBR, msBR), fill=text_color, font=tsfont)
        grayBR = RGBBR.convert('L')
        bottomRightPanel = np.array(grayBR)
        
        #look for closest frame time for bottom left panel
        if hasThirdVideo:
            ftarrayBL = vidDict[imgFolders[2]]['frameTimes']
            closestFrameIdxBL = (np.abs(ftarrayBL - frameTime)).argmin()
            bottomLeftPanelName = vidDict[imgFolders[2]]['imgNames'][closestFrameIdxBL]
            bottomLeftPanelImg =Image.open(os.path.join(vidDict[imgFolders[2]]['dirName'], imgFolders[2], bottomLeftPanelName))
            RGBBL = bottomLeftPanelImg.convert('RGB')
            drawObjBL = ImageDraw.Draw(RGBBL)
            ftBL = ftarrayBL[closestFrameIdxBL]
            hrBL = ftBL // 3600000
            remainingSecs = ftBL % 3600000
            minBL = remainingSecs // 60000
            remainingSecs = remainingSecs % 60000
            secBL = remainingSecs // 1000
            msBL = remainingSecs % 1000
            drawObjBL.text((10, 10), '%02d:%02d:%02d.%03d' % (hrBL, minBL, secBL, msBL), fill=text_color, font=tsfont)
            grayBL = RGBBL.convert('L')
            bottomLeftPanel = np.array(grayBL)
            
        topCat = np.hstack([topLeftPanel, topRightPanel])
        botCat = np.hstack([bottomLeftPanel, bottomRightPanel])
                                           
        outputFrame = np.vstack([topCat, botCat])

    
        #frame = imutils.resize(frame, width=640)
        # check if the writer is None
        if writer is None:
            # store the image dimensions, initialize the video writer,
            # and construct the zeros array
            writer = cv2.VideoWriter(os.path.join(saveTo, movieName), fourcc, largeFR, (2*h, 2*w), False)

        # write the output frame to file

        writer.write(outputFrame)
        # show the frames
        #cv2.imshow("OutputText", outputFrame)
        frameTime += msperFrame
        key = cv2.waitKey(1) & 0xFF
        # if the `q` key was pressed, break from the loop
        if key == ord("q"):
            break

    # do a bit of cleanup
    print("[INFO] cleaning up...")
    cv2.destroyAllWindows()
    writer.release()
