'''
Functions for reaching 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
from imutils.video import VideoStream
import statistics
from pathlib import Path
from cv_utils import *
import csv

def json_serial(obj):
    import datetime as dt
    if isinstance(obj, (dt.datetime, dt.date)):
        return obj.isoformat()
    raise TypeError ("Type %s not serializable" % type(obj))
    
def getLabelDict(labelsPath):
    with open(labelsPath, newline='') as csvfile:
        myReader = csv.reader(csvfile)
        rowNum = 0
        for row in myReader:
            if rowNum==1:
                bodyparts = []
                for r in row[1:]:
                    if not r in bodyparts:
                        bodyparts = bodyparts + [r]
                print(bodyparts)
                numparts = len(bodyparts)
                colors = np.linspace(0, 1, numparts)
                partsDict = {}
                for cidx, part in enumerate(bodyparts):
                    partsDict.update({part:{
                                    'x':[],
                                    'y':[],
                                    'likelihood':[],
                                    'coloridx':colors[cidx]
                                    }
                            })
                xinds = [x for x in range(1,3*numparts,3)]
                yinds = [y for y in range(2,3*numparts,3)]
                linds = [l for l in range(3,3*numparts+1,3)]
                
            if rowNum > 2:
                for partnum, x in enumerate(xinds):
                    part = bodyparts[partnum]
                    oldList = partsDict[part]['x']
                    newList = oldList + [np.float64(row[x])]
                    partsDict[part].update({'x':newList})
                for partnum, y in enumerate(yinds):
                    part = bodyparts[partnum]
                    oldList = partsDict[part]['y']
                    newList = oldList + [np.float64(row[y])]
                    partsDict[part].update({'y':newList})
                for partnum, l in enumerate(linds):
                    part = bodyparts[partnum]
                    oldList = partsDict[part]['likelihood']
                    newList = oldList + [np.float64(row[l])]
                    partsDict[part].update({'likelihood':newList})
                    
            rowNum+=1
            
        return partsDict

#############################################################################
#
# 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 makeReachingMovieLabeled(topMovPath, topLabelPath, bottomMovPath, bottomLabelPath, saveToFolder, codec='MJPG',ext='.avi'):
    plt.close('all')
    bn = os.path.basename(topMovPath)
    mn = bn[:len(bn)-4]
    labeledMovieName = mn + '_labeled'
    labeledMoviePath = os.path.join(saveToFolder, labeledMovieName + ext)
    print('Save to:', labeledMoviePath)
    
    h = 1024
    w = 1280

    (hf, wf) = (h*2, int(w*1.3))

    outputFrame = np.zeros((hf, wf, 3)).astype(np.uint8)

    fontSize = 50
    font = {'family' : 'sans-serif',
            'weight' : 'light',
            'size'   : 18}

    mpl.rc('font', **font)
    font = ImageFont.truetype(os.path.join(Path().absolute(), 'Roboto', 'Roboto-Black.ttf'), fontSize)
    
    yinc = fontSize
    ypos = 0
    capTop = cv2.VideoCapture(topMovPath)
    capBottom = cv2.VideoCapture(bottomMovPath)
    
    fps = capTop.get(cv2.CAP_PROP_FPS)
    print('bottomMovPath:', bottomMovPath)
    print('fps:', fps)
    # initialize the FourCC, video writer, dimensions of the frame, and
    fourcc = cv2.VideoWriter_fourcc(*codec)
    writer = cv2.VideoWriter(labeledMoviePath, fourcc, fps, (wf, hf))

    lcutoffT = 0.2
    lcutoffB = 0.2
    markerRadius = 15
    
    topPartsDict = getLabelDict(topLabelPath)
    bottomPartsDict = getLabelDict(bottomLabelPath)
    
    topKeyPoints = list(topPartsDict.keys())
    bottomKeyPoints = list(bottomPartsDict.keys())
    
    numTopParts = len(topKeyPoints)
    numBottomParts = len(bottomKeyPoints)
    
    topcmap = mpl.cm.get_cmap('hsv', numTopParts)
    bottomcmap = mpl.cm.get_cmap('hsv', numBottomParts)
    
    bodyPartsColsT = {}
    for part in topPartsDict.keys():
        colidx = topPartsDict[part]['coloridx']
        col = tuple(np.round(np.multiply(255, topcmap(colidx))).astype(int))[0:3]
        bodyPartsColsT.update({part:col})
    print(bodyPartsColsT)
    bodyPartsColsB = {}
    for part in bottomPartsDict.keys():
        colidx = bottomPartsDict[part]['coloridx']
        col = tuple(np.round(np.multiply(255, bottomcmap(colidx))).astype(int))[0:3]
        bodyPartsColsB.update({part:col})
    
    topLabelIm = Image.new('RGB', (int(w*0.3), h), color=(50, 50, 50))
    bottomLabelIm = Image.new('RGB', (int(w*0.3), h), color=(50, 50, 50))
    drawObjTLI = ImageDraw.Draw(topLabelIm)
    drawObjBLI = ImageDraw.Draw(bottomLabelIm)
    
    for part in bodyPartsColsT.keys():
        drawObjTLI.text((0, ypos), part, fill=bodyPartsColsT[part], font=font)
        ypos += yinc
    drawObjTLI.text((0, ypos), 'likelihood>%04.2f' % lcutoffT, fill=(255,255,255), font=font)
    ypos = 0
    for part in bodyPartsColsB.keys():
        drawObjBLI.text((0, ypos), part, fill=bodyPartsColsB[part], font=font)
        ypos += yinc
    drawObjBLI.text((0, ypos), 'likelihood>%04.2f' % lcutoffB, fill=(255,255,255), font=font)
    
    topLabelOut = np.array(topLabelIm).astype(np.uint8)
    bottomLabelOut = np.array(bottomLabelIm).astype(np.uint8)
    outputFrame[0:h, w:, :] = np.flip(topLabelOut,axis=2)
    outputFrame[h:, w:, :] = np.flip(bottomLabelOut,axis=2)
    ret = True
    frameNum = 0
    while ret:
        ret, frameTopRGB = capTop.read()
        if not ret:
            break
        else:
            ret, frameBottomRGB = capBottom.read()
            topIm = Image.fromarray(frameTopRGB, mode='RGB')
            bottomIm = Image.fromarray(frameBottomRGB, mode='RGB')
            #topIm = Image.new('RGB', (w, h), color=(50, 50, 50))
            #bottomIm = Image.new('RGB', (w, h), color=(50, 50, 50))
            
            drawObjT = ImageDraw.Draw(topIm)
            drawObjB = ImageDraw.Draw(bottomIm)
            #draw Labels
            for part in topPartsDict.keys():
                x = topPartsDict[part]['x'][frameNum]
                y = topPartsDict[part]['y'][frameNum]
                l = topPartsDict[part]['likelihood'][frameNum]
                #ladd = np.log10(l)
                #print('ladd:', ladd)
                #xyT = [int(x-(markerRadius+ladd)), int(y-(markerRadius+ladd)), 
                #      int(x+markerRadius+ladd), int(y+(markerRadius+ladd))]
                
                if l>lcutoffT:
                    xyT = [int(x-(markerRadius*l)), int(y-(markerRadius*l)), 
                      int(x+markerRadius*l), int(y+(markerRadius*l))]
                    drawObjT.ellipse(xyT,
                                  fill=bodyPartsColsT[part], outline=None)
            for part in bottomPartsDict.keys():
                x = bottomPartsDict[part]['x'][frameNum]
                y = bottomPartsDict[part]['y'][frameNum]
                l = bottomPartsDict[part]['likelihood'][frameNum]
                
                if l>lcutoffB:
                    xyB = [int(x-(markerRadius*l)), int(y-(markerRadius*l)), 
                      int(x+markerRadius*l), int(y+(markerRadius*l))]
                    drawObjB.ellipse(xyB,
                                  fill=bodyPartsColsB[part], outline=None)
                
        topOut = np.array(topIm)#.astype(np.uint8)
        bottomOut = np.array(bottomIm)#.astype(np.uint8)
        outputFrame[0:h, 0:w, 0] = topIm.getchannel('B')#np.flip(topOut,axis=2)
        outputFrame[0:h, 0:w, 1] = topIm.getchannel('G')
        outputFrame[0:h, 0:w, 2] = topIm.getchannel('R')
        #outputFrame[0:h, 0:w, :] = np.flip(topOut, axis=2)
        outputFrame[h:, 0:w, :] = np.flip(bottomOut,axis=2)
        
        
        cv2.imshow("OutputFrame", outputFrame)
        #cv2.imshow("topIm", np.array(topIm).astype(np.uint8))
        writer.write(outputFrame)
        # show the frames
        key = cv2.waitKey(30)
        # if the `q` key was pressed, break from the loop
        if key == 27:
            break
        frameNum += 1
    # do a bit of cleanup
    print("[INFO] cleaning up...")
    cv2.destroyAllWindows()
    writer.release()
