# Given pickle files as generated by GA.py or GA_xor.py, this will generate videos of the experiments.

import pickle, sys, glob, multiprocessing
import cv2
import numpy as np


def calculateFitness(experiment):
    results = []

    for i in range(4):
        total_sum = np.sum(experiment[i])
        center_sum = np.sum(experiment[i][:,12])

        if i == 1 or i == 2: # center must be on (10, 01)
            fitness = center_sum * 25 - total_sum
            results.append(fitness)

        if i == 0 or i == 3: # center must be off (00, 11)
            fitness = total_sum - center_sum * 25 
            results.append(fitness)
    
    return results


def generate_singlevideo(path, processLimiter=multiprocessing.Lock()):

    font = cv2.FONT_HERSHEY_SIMPLEX

    with processLimiter:
        print("Processing p file "+path)
        data = pickle.load( open( path, "rb"))
        motors, predictions = data
        num_preds = len(predictions[0])

        for pred in range(num_preds):
            ps = [ predictions[i][pred] for i in range(len(predictions)) ]
            fitness = calculateFitness( ps ) 

            for xor in range(4):

                fourcc = cv2.VideoWriter_fourcc(*'X264')
                outvideo = path.split(".")[0]+"_"+str(pred)+"_"+str(fitness[xor])
                outvideo += "_" + str(xor) + ".avi"
                outvideo = cv2.VideoWriter(outvideo, fourcc, 20.0, (500,500))

                accumulator = np.zeros(25)

                for entry in range(len(predictions[xor][pred])):
                    blank_frame = np.zeros((500,500,3), np.uint8)
                    datarow = predictions[xor][pred][entry]
                    motor_row = motors[pred]
                    accumulator += datarow

                    for i in range(5):
                        for j in range(5):
                            point1 = (i*100, j*100)
                            point2 = (i*100+100, j*100+100)
                            colour = int(datarow[i*5+j] * 255)
                            cv2.rectangle(blank_frame, point1, point2, 
                                    [colour,colour,colour], -1) 

                            #speed = (motor_row[i*5+j] * 2) -1
                            #cv2.putText(blank_frame, f'{ int(speed*10) }',
                            #        (point1[0]+25, point1[1]+55), font, 1, (255, 255, 255))

                            fit = (accumulator[i*5+j] * 2) -1
                            cv2.putText(blank_frame, f'{ int(fit) }',
                                    (point1[0]+25, point1[1]+55), font, 1, (255, 255, 255))
                    
                    outvideo.write(blank_frame)

                outvideo.release()
                cv2.destroyAllWindows()

def processfolder(pathtofolder):
    ''' This function will execute the previous 1 in all the files of a folder.
    '''

    # number of operations in parallel. auchentoshan has 12 cores
    s = multiprocessing.Semaphore(12)
    # we only want to process the p files
    allpfiles = glob.glob(pathtofolder+'*.p')
    
    for pfile in allpfiles:
        p = multiprocessing.Process(target=generate_singlevideo, args=(pfile,s))
        p.start()


if __name__ == '__main__':

    processfolder(sys.argv[1])
