# import
# ------------------------------------------------------------------------------------------
import os, sys
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats
import multiprocessing as mp

import easygui
import imageio

# miam import
import miam.image.Image as MIMG
import miam.image.imageType as MTYPE
import miam.image.channel as MCHN

import miam.processing.ColorSpaceTransform as MCST
import miam.processing.TMO_Lightness as TMO_L
import miam.processing.ContrastControl as PCC
import miam.histogram.Histogram  as MHIST
import miam.aesthetics.LightnessAesthetics as MLAC
import miam.aesthetics.Palette as MPAL
import miam.imageDB.ImageDB
import miam.imageDB.POGChecker
import miam.imageDB.HwHDRBuilder
import miam.html.generator
import miam.utils
import miam.pointcloud.PointCloud2D

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



def pCompPalette3(filename):
    
    # cut filename
    path, name, ext = miam.utils.splitFileName(filename)

    # read image
    img = MIMG.Image.read(filename,exif=False)

    # compute palette      
    nbColors = 3
    palette = MPAL.Palette.build(img, nbColors, fast=True, method='kmean-Lab',removeBlack=True)
    # get colors
    colors = palette.colors

    #print("█",end='')
    print("3.",end='')
    sys.stdout.flush()

    return colors

def pCompPalette2(filename):
    
    # cut filename
    path, name, ext = miam.utils.splitFileName(filename)

    # read image
    img = MIMG.Image.read(filename,exif=False)

    # compute palette      
    nbColors = 2
    palette = MPAL.Palette.build(img, nbColors, fast=True, method='kmean-Lab',removeBlack=True)
    # get colors
    colors = palette.colors

    #print("█",end='')
    print("2.",end='')
    sys.stdout.flush()

    return colors

if __name__ == '__main__':
    print("MIAM[POG Cluster Palette]")

    # ------------------------------------------------------------------------------------------
    # what to do !
    # ------------------------------------------------------------------------------------------
    # config: configfile
    jsonfile = "../DB/config_POG_DB.json"


    # all palettes file
    allPalettes_2colorsNoBlack = "../DB/POG_DB_Lab_kmeans_Palette_2colorsNoBlack.npy"
    allPalettes_3colorsNoBlack = "../DB/POG_DB_Lab_kmeans_Palette_3colorsNoBlack.npy"
    allPalettes_4colorsNoBlack = "../DB/POG_DB_Lab_kmeans_Palette_4colorsNoBlack.npy"
    allPalettes_5colorsNoBlack = "../DB/POG_DB_Lab_kmeans_Palette_5colorsNoBlack.npy"

    # ------------------------------------------------------------------------------------------

    # pog
    pogDB = miam.imageDB.ImageDB.ImageDB(jsonConfigFile ="../DB/config_POG_DB.json")
    pogDB.check(miam.imageDB.POGChecker.POGChecker(local=True))

    print("MIAM[",pogDB.name,":(csvFileName:: ",pogDB.csvFileName,", imagePATH::",pogDB.imagePATH,")]")

    # POG: compute Lightness Histograms
    # ------------------------------------------------------------------------------------------
    if not os.path.isfile(allPalettes_3colorsNoBlack):
        print("MIAM[",pogDB.name," palette 3 colors not found !]")

        nbCpu = mp.cpu_count()

        print("MIAM[", nbCpu, "cpu cores]")
        print("MIAM[",pogDB.name,"  number of images:",len(pogDB.db),"]")
        print("MIAM[",pogDB.name," launch parallel computation of palettes (3) !]")

        _pool = mp.Pool()
        result = _pool.map(pCompPalette3, pogDB.db)
        allColors = np.asarray(result)
        print("")
        print("[ ---------------------------------------------------------- ]")
        print("MIAM[",pogDB.name," parallel computation of palettes (3) done !]")
        print("MIAM[",pogDB.name," all palettes (3):",allColors.shape," ]")
        np.save(allPalettes_3colorsNoBlack,allColors)
        print("MIAM[",pogDB.name," all palettes saved !]")
    else:
        print("")
        print("[ --------------------------------------------------------]")
        print("MIAM[",pogDB.name," color palettes (3) found !]")

    if not os.path.isfile(allPalettes_2colorsNoBlack):
        print("MIAM[",pogDB.name," palette 2 colors not found !]")

        nbCpu = mp.cpu_count()

        print("MIAM[", nbCpu, "cpu cores]")
        print("MIAM[",pogDB.name,"  number of images:",len(pogDB.db),"]")
        print("MIAM[",pogDB.name," launch parallel computation of palettes (2) !]")

        _pool = mp.Pool()
        result = _pool.map(pCompPalette2, pogDB.db)
        allColors = np.asarray(result)
        print("")
        print("[ ---------------------------------------------------------- ]")
        print("MIAM[",pogDB.name," parallel computation of palettes (2) done !]")
        print("MIAM[",pogDB.name," all palettes  (2):",allColors.shape," ]")
        np.save(allPalettes_2colorsNoBlack,allColors)
        print("MIAM[",pogDB.name," all palettes saved !]")
    else:
        print("")
        print("[ --------------------------------------------------------]")
        print("MIAM[",pogDB.name," color palettes (2) found !]")