#generate files for images and V-A values for SEMAINE dataset. landmark points generated by D. Aspandi because SEMAINE does not provide them

import os
import os.path as path
import zipfile
import cv2
import numpy as np

currentDir = os.getcwd()

#create the data folder that will hold all the data files
dataFolder = path.join(path.abspath(path.join(currentDir,'..')), 'data/SemData')
if not os.path.exists(dataFolder):
    os.makedirs(dataFolder)

#get 'data' path
currentDirLevelUp = path.abspath(path.join(currentDir,".."))
for files in os.listdir(currentDirLevelUp):
    if files == 'data':
        dataDir = path.join(currentDirLevelUp,files)
#choose how many sessions: all or just a subset
sessionsDirs = []

#get list of zip files paths
for file in os.listdir(dataDir):
    if file == 'SEMAINE':
        subDataDir = path.join(dataDir,file)
        for file2 in os.listdir(subDataDir):
            if file2 == 'semaineDB':
                subDataDir2 = path.join(subDataDir,file2)
                for file3 in os.listdir(subDataDir2):
                    if file3 == 'Sessions':
                        subDataDir3 = path.join(subDataDir2,file3)
                        for file4 in os.listdir(subDataDir3):
                            sessionsDirs.append(path.join(subDataDir3,file4))


testDirs = sessionsDirs

#only build samples where there i both valence and arousal
hasVSess = False
hasASess = False

#dirs list for each session participant (user and/or operator)
#each session may have more than one file for a participant's valence and arousal
userAroDirs = []
userValDirs = []
operAroDirs = []
operValDirs = []

#some data files have a great difference between total time of analysis files with viodeo videoLength
#a minimum difference is needed to not save lists which will require tp cut a lot of the video
minimumDifference = 230 #170 samples = 3.4 sec, 230samples = 4.6sec

for sessDirs in testDirs:
    subSessDirs = []
    userVidDir = ""
    userAroDirs = []
    userValDirs = []
    operVidDir = ""
    operAroDirs = []
    operValDirs = []

    personSubSessDirs = []
    #check if contains both valence and arousal info (files)
    for subFiles in os.listdir(sessDirs):
        if "DA.txt" in subFiles:
            hasASess = True
        if "DV.txt" in subFiles:
            hasVSess = True
    if hasASess and hasVSess:
        hasASess,hasVSess = False, False #reset to initial values
        for subFiles in os.listdir(sessDirs):
            #save videos dirs for user and operator (avi files)
            #save those file dirs that contain TU (user) or TO (oper) and DV,DA or DV2,DA2 substrings
            if "User" in subFiles and ".avi" in subFiles:
                userVidDir = path.join(sessDirs,subFiles)
            if "TU" in subFiles and ("DA.txt" in subFiles or "DA2.txt" in subFiles):
                userAroDirs.append(path.join(sessDirs,subFiles))
            if "TU" in subFiles and ("DV.txt" in subFiles or "DV2.txt" in subFiles):
                userValDirs.append(path.join(sessDirs,subFiles))
            if "Operator" in subFiles and ".avi" in subFiles:
                operVidDir = path.join(sessDirs,subFiles)
            if "TO" in subFiles and ("DA.txt" in subFiles or "DA2.txt" in subFiles):
                operAroDirs.append(path.join(sessDirs,subFiles))
            if "TO" in subFiles and ("DV.txt" in subFiles or "DV2.txt" in subFiles):
                operValDirs.append(path.join(sessDirs,subFiles))

        #if no data files for operator, remove its video dir
        if not operAroDirs and not operValDirs:
            operVidDir = ""

        subSessDirs = [userAroDirs,userValDirs,operAroDirs,operValDirs]
        #if file has 2 at the end, this is the useful vesion, discard the original ones
        #check for DV2 or DA2, take corresponding DV or DA out of file dirs
        for i,subPartiValueDir in enumerate(subSessDirs):
            for subDirs in subPartiValueDir:
                if "DV2.txt" in subDirs:
                    commonSubstring = subDirs[:-5] #this is the common part to both versions names
                    originalOne = commonSubstring+".txt"
                    for subDirs2 in subPartiValueDir:
                        #if there is an original version to be removed for the file 2
                        if originalOne in subDirs2:
                            subSessDirs[i].remove(originalOne) #remove dir from list
                            print("hey")
                if "DA2.txt" in subDirs:
                    commonSubstring = subDirs[:-5]
                    originalOne = commonSubstring+".txt"
                    for subDirs2 in subPartiValueDir:
                        if originalOne in subDirs2:
                            subSessDirs[i].remove(originalOne)

        #print(subSessDirs)
        #print(userVidDir,operVidDir)
        userSubSessDirs = subSessDirs[:2]
        operSubSessDirs = subSessDirs[2:]
        personSubSessDirs.append(userSubSessDirs)
        #only add operator lists variables if there is any file
        if not (len(subSessDirs[2]) == 0) or not (len(subSessDirs[3]) == 0):
            personSubSessDirs.append(operSubSessDirs)

        v = cv2.VideoCapture(userVidDir)
        v.set(cv2.CAP_PROP_POS_AVI_RATIO,1) #1 is to set pointer to the last frame
        videoLength = v.get(cv2.CAP_PROP_POS_MSEC)/1000 #get duration, /1000 because ir is given in milisec
        #only perform for correct files
        if videoLength > 0.0:
            #user or operator (if there is)
            for k,subPersonSubSessDirs in enumerate(personSubSessDirs):
                dimensionValuesLists = []
                #valence or arousal data files
                for subPartiValueDir2 in subPersonSubSessDirs:
                    emotionValuesLists = []
                    if len(subPartiValueDir2) > 1:
                        #subData files (different raters, same partcipant and data dimension)
                        for ratersSubFiles in subPartiValueDir2:
                            with open(ratersSubFiles, "r") as file:
                                listOfValues = [float(line.split()[1]) for line in file] #data samples for file
                            fileTotalTime = len(listOfValues)*0.02 #how long is the sample analysis time, a sample is taken every 0.02 sec
                            #discard data files that took much less time than video length
                            if len(listOfValues) > ((videoLength/0.02)-minimumDifference):
                                emotionValuesLists.append(listOfValues)
                                #print(ratersSubFiles)
                        dimensionValuesLists.append(emotionValuesLists)

                #the common minimum length to build the data lists that will be used to write the new data files
                minLength = min(list(map(len,dimensionValuesLists[0]))) #minArousal
                minValenceLen = min(list(map(len,dimensionValuesLists[1])))
                if (minValenceLen < minLength):
                    minLength = minValenceLen
                roundVideoLength = round(videoLength/0.02) #to int in python 3
                if (roundVideoLength < minLength):
                    minLength = roundVideoLength - 10 #-10 so not the exact ending of the video
                #print(list(map(len,dimensionValuesLists[0])), videoLength/0.02, minLength)

                #calculate average arousal and valence arrays
                averageLists = []
                for subDim in dimensionValuesLists:
                    averageValues = np.zeros(minLength)
                    for subRaters in subDim:
                        #a+=b does not work because types do not match
                        #the legth of all lists is aligned with limit of minLength
                        np.add(averageValues, np.array(subRaters[:minLength]), out=averageValues, casting='unsafe')
                    averageValues=averageValues/len(subDim)
                    averageLists.append(averageValues)
                    #print(len(averageValues))

                folderID = userVidDir.split('/')[-2] #separate dir substrings by '/' and get penultimate (semaine session subfolder number)

                #the second iteration means there is operator data, change video dir and folder ID
                if k == 1:
                    userVidDir = operVidDir
                    folderID = str(int(folderID)+140) #the top number of folders ID is 140, so the new ones for operator case can use +140 to not overlap

                #subfolders to store session data
                newSessionFolder = path.join(dataFolder,folderID)
                if not os.path.exists(newSessionFolder):
                    os.makedirs(newSessionFolder)
                imgSessionFolder = path.join(newSessionFolder,'img')
                if not os.path.exists(imgSessionFolder):
                    os.makedirs(imgSessionFolder)
                vaSessionFolder = path.join(newSessionFolder,'annot2')
                if not os.path.exists(vaSessionFolder):
                    os.makedirs(vaSessionFolder)

                #write data to new files
                counter = 0
                v2 = cv2.VideoCapture(userVidDir)
                print(minLength,'+++',folderID)
                #each sample is taken after 0.02 seconds, to obtain fewer samples, jump each 10 samples and cut every 0.2 sec (10*0.02) 
                for x in range(0,minLength,10):
                    counter=round(counter+0.2,2) #video steps to cut to images in sec
                    #create imaga files and store them
                    v2.set(cv2.CAP_PROP_POS_MSEC,counter*1000) #time pointer in milisec
                    success,image = v2.read()
                    if success:
                        cv2.imwrite(imgSessionFolder+'/'+str(folderID)+'_'+str(x).zfill(5)+'.png', image)
                        #write [arousal valence] to .aro files and store them
                        with open(vaSessionFolder+'/'+str(folderID)+'_'+str(x).zfill(5)+'.aro', "w") as file:
                            file.write(str(averageLists[0][x])+' '+str(averageLists[1][x]))


    print("__________")
