# import
# ------------------------------------------------------------------------------------------
import os
import multiprocessing as mp
import matplotlib.pyplot as plt
import numpy as np
import easygui
import colour

# miam import
import miam.image.Image as MIMG
import miam.processing.ContrastControl as PCC
import miam.utils

# ------------------------------------------------------------------------------------------
# MIAM project 2020
# ------------------------------------------------------------------------------------------
# author: remi.cozot@univ-littoral.fr
# ------------------------------------------------------------------------------------------

def pfun(f):
    path, name, ext = miam.utils.splitFileName(f)

    current = mp.current_process()
    print('core:', current.name,"(",current._identity,")","::",name)

    # load image and remove zeros
    img = MIMG.Image.readImage(f)
    img.name = name   

    # tonemap
    leqImg = img.process(PCC.ContrastControl(),method='localEqualization', size=1500)
    leqImg.name = img.name+"_localEQ"

    # save tone mapped image
    leqImg.writeImage("../local_EQ/"+leqImg.name)

    print('core:', current.name,"(",current._identity,")","::",leqImg.name,":: done n saved")

    return None

if __name__ == '__main__':

    print("MIAM: local equalization of all images")

    nbCpu = mp.cpu_count()
    print(" ---- found:", nbCpu, "cpu cores")

    srcPath = "../../../photos/HDR-DB/HDR/"
    hdrfiles = os.listdir(srcPath)

    todo = []

    tmopath = "../lightness_TMO/"
    leqpath = "../local_EQ/"
    for h in hdrfiles:
        path, name, ext = miam.utils.splitFileName(h)

        tmofile = tmopath+name+"_Lightness_TMO.jpg"
        leqfile = leqpath+name+"_Lightness_TMO"+"_localEQ.jpg"

        tmofileOK = os.path.isfile(tmofile)
        leqfileOK = os.path.isfile(leqfile)

        #print("source:", h,
        #      "tmofile::",tmofile,"::", tmofileOK,
        #      "leqfile::",leqfile,"::", leqfileOK,
        #      )


        if (tmofileOK and (not leqfileOK)): todo.append(tmofile)

    listFiles = todo


    print(" ---- number of images:",len(listFiles))
    for f in listFiles: print(" ---- file:", f)
    print(" ---- launch local equalization !")

    _pool = mp.Pool()
    result = _pool.map(pfun, listFiles)
    print(" ---- local equalize all done !")