#! /usr/bin/env python

from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State
import sys, os, datetime
import shutil
import sqlite3
import zipfile
import utils
import cvutils
import cv2


def zipFolder(inputFolder, outputFile):
    '''Method to compress the content of the inputFolder in the outputFile'''
    zip = zipfile.ZipFile(outputFile, 'w')
    for root, dirs, files in os.walk(inputFolder):
        for file in files:
            zip.write(root+file, file)
    zip.close()




def getTypeDict(cursor):
    '''Return a dictionnary with integer key and associated type string
    i.e.: "0"  -> "unknown"
          "1"  -> "car"
          "2"  -> "pedestrians"
          "3"  -> "motorcycle"
          "4"  -> "bicycle"
          "5"  -> "bus"
          "6"  -> "truck"
          ... and other type if the objects_type table is defined in SQLite'''
    typeDict = dict()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='objects_type'")
    data = cursor.fetchone()

    if(data == None):
        typeDict["0"] = "unknown"
        typeDict["1"] = "car"
        typeDict["2"] = "pedestrians"
        typeDict["3"] = "motorcycle"
        typeDict["4"] = "bicycle"
        typeDict["5"] = "bus"
        typeDict["6"] = "truck"
        
    else:
        cursor.execute("SELECT road_user_type, type_string FROM objects_type")
        for row in cursor:
            typeDict[row[0]]= row[1]

    return typeDict

def extractFrames(videoFile, framePath, fps, time, startFrame = 0, endFrame = None):
    '''Method to extract all the frames of the video'''
    print('Extracting frame')
    deltaTimestamp = 1000.0/float(fps);
    time+=datetime.timedelta(microseconds=startFrame*deltaTimestamp*1000)
    
    inc = 1000 #How many frame we fetch in the video at a time

    if endFrame != None:
        delta = endFrame-startFrame
        if delta < inc:
            inc = delta
    
    currentIdx = startFrame
    frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
    
    while len(frameList) == inc and inc > 0:
        
        for f in frameList:
            cv2.imwrite(os.path.join(framePath,time.strftime("%Y%m%d-%H%M%S.%f")[:-3]+'.jpg'), f)
            time += datetime.timedelta(microseconds=deltaTimestamp*1000)
        currentIdx = currentIdx + inc

        if endFrame != None:            
            delta = endFrame-currentIdx
            if delta < inc:
                inc = delta        
        if inc:
            frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
        print('Extracting frame ' + str(currentIdx))
    return len(frameList) > 0

    

def convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = None, videoFolderExist = False, startFrame = 0, endFrame = None):
    '''
    Method to convert database from polytrack to PDTV:
    workDirname  is the current working directory
    sceneFilename is the path to the .cfg file
    sectionName is the name of the section we want to process in this file
    videoFolderExist specifiy if we want to reextract the video frame or if they already exist at workdir/videoframes/
    startFrame is the first frame we want to extract
    endFrame is the last frame we want to extract (or None if we want to extract everything)
    '''

    scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename))
    if databaseFilename != None:
        inputDb = os.path.join(workDirname, databaseFilename)
    else:
        inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename)
    videoFile = os.path.join(workDirname, scene[sectionName].videoFilename)

    error = False

    if not os.path.exists(inputDb):
        print('Database path is not valid: ' + inputDb)
        error = True
    if not os.path.exists(videoFile):
        print('Video path is not valid: ' + videoFile)
        error = True

    time = scene[sectionName].date    

    
    videoFolderPath = os.path.join(workDirname, "videoframes/")
    fileName = scene[sectionName].sitename
    
    fps = cvutils.getFPS(videoFile)
    print('Video should run at ' + str(fps) + ' fps')
    deltaTimestamp = 1000.0/float(fps);
    if videoFolderExist == False:
        if os.path.exists(videoFolderPath):
            shutil.rmtree(videoFolderPath)
        utils.mkdir(videoFolderPath)
        if extractFrames(videoFile, videoFolderPath, fps, time, startFrame, endFrame) == 0:
            print("Error. Frame were not extracted")
            error = True
    
    if not error:
        masterTimestamp = 0.0
        masterTimestampList = list()
        video_timestringsList = list()
        frameNumberToMasterTimestamp = dict()
        for r,d,f in os.walk(videoFolderPath):
            i = startFrame
            for files in f:
                name, ext = os.path.splitext(files)           
                video_timestringsList.append(name)
                masterTimestampList.append(masterTimestamp)            
                frameNumberToMasterTimestamp[i] = masterTimestamp
                i = i +1
                masterTimestamp = masterTimestamp+deltaTimestamp
        
        inputZipVideoName = fileName+".zip"
        print('Zipping files...')
        if not os.path.exists(inputZipVideoName) or not videoFolderExist:
            zipFolder(videoFolderPath, inputZipVideoName)
        print('Zipping files...Done.')
        #We generate the structure for ZipVideo
        zipVideo = ZipVideo(video_zip_file=inputZipVideoName,
                        time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file='calib.json')
        zipVideo.save(os.path.join(workDirname, 'video.json'))

        print('ZipVideo saved')
        obj = SyncedVideos(master_timestamps = [masterTimestamp],
                               video_timestrings = [video_timestringsList],
                               video_files = ['video.json'],
                               fps=fps)
        obj.save(os.path.join(workDirname, 'syncedvideo.json'))

        print('SyncedVideos saved')

        print('Opening db')
        connection = sqlite3.connect(inputDb)
        cursor = connection.cursor()
      

        #Tracket database
        trackset = TrackSet(synced_file = ['syncedvideo.json'])

        
        #1) We build the type index dictionnary
        typeDict = getTypeDict(cursor)
        
        idToTrackDict = dict()
        #2) We read the object database    
        cursor.execute("SELECT object_id, road_user_type FROM objects")
        for row in cursor:
            objectId = row[0]
            objectType = row[1]
            t = Track(type=typeDict.get(objectType, "unknown"))
            idToTrackDict[objectId] = t;
            trackset.append(t)
        print('Reading boundingbox table')
        #3) We read the bounding box table    
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='bounding_boxes'")
        data = cursor.fetchone()
        if data == None:
            print('No bounding box table. Maybe it was not generated ?')
        else:
            cursor.execute("SELECT object_id, frame_number, x_top_left, y_top_left, x_bottom_right, y_bottom_right FROM bounding_boxes")
            for row in cursor:
                objectId = row[0]
                frameNumber = row[1]
                x_top_left = row[2]
                y_top_left = row[3]
                x_bottom_right = row[4]
                y_bottom_right = row[5]
                
                idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.], 
                                                 master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)],
                                                 bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]]))
        print('Saving db in json')
        trackset.save(os.path.join(workDirname, 'roaduser.json'))
        connection.close()
        print('Done.')

    
    

if __name__ == "__main__":
    workDirname = "./"
    sceneFilename = "scene.cfg"
    sectionName = "amherst-11"
    # videoFrameAlreadyExtracted Set this to true if you have the frame already extracted in the folder workDirname/videoframes/
    convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = 'amherst-11_gt.sqlite', videoFolderExist = True, startFrame = 0, endFrame = 4000)