# import
# ------------------------------------------------------------------------------------------
import os, sys, math, json
import matplotlib.pyplot as plt
import numpy as np
import easygui

# miam import
from . import WFNode, WFProcess, WFConnector
import miam.image.Image as MIMG
import miam.histogram.Histogram as MHIST
import miam.image.channel
import miam.utils

from miam.processing import (ColorSpaceTransform, ContrastControl, Duplicate, ExposureControl, NoOp,
                             TMO_CCTF, TMO_Lightness, TMO_Linear, Ymap, Blend, MaskSegmentPercentile,
                             Fuse
                             )

# ------------------------------------------------------------------------------------------
# MIAM project 2020
# ------------------------------------------------------------------------------------------
# author: remi.cozot@univ-littoral.fr
# ------------------------------------------------------------------------------------------
class WFWorkflow(object):
    """container of workflow"""

    def __init__(self, name='workflow:noname'):
        # attributes
        self.name = name

        # list of processes
        self.processes =  []
        # list of images
        self.connectors = []
        # root
        self.root = None
        # end
        self.leafs = []

    # methods
    # ---------------------------------------------
    # add process, connect, root and leaf
    # ---------------------------------------------
    def addProcess(self,p):
        self.processes.append(p)

        return p

    # connecting processes
    def connect(self, outputProcess=None, inputProcess=None, name=None):
        if (not outputProcess) or (not inputProcess):
            if not outputProcess: print("[ERROR] WFWorkflow.connect(): unknown output !")
            if not inputProcess: print("[ERROR] WFWorkflow.connect(): unknown input !")
            return None
        else: 
            # outputProcess -- WFImage --> inputProcess
            outputProcess = self.getByName(outputProcess) if isinstance(outputProcess,str) else outputProcess
            inputProcess = self.getByName(inputProcess) if isinstance(inputProcess,str) else inputProcess

        if not outputProcess: print("[ERROR] WFWorkflow.connect(): unknown output !")
        if not inputProcess: print("[ERROR] WFWorkflow.connect(): unknown input !")

        if not name:
            name = name=outputProcess.name+"->"+inputProcess.name
        connector = WFConnector.WFConnector(name=name)
        # link components
        outputProcess.outputs.append(connector)
        inputProcess.inputs.append(connector)
        connector.outputOf = outputProcess
        connector.inputOf = inputProcess

        self.connectors.append(connector)

        return connector

    def setRoot(self,p):
        p = self.getByName(p) if isinstance(p,str) else p

        if not p: print("[ERROR] WFWorkflow.setRoot(): unknown process !")

        connector = WFConnector.WFConnector(name="root"+"->"+p.name)
        connector.isRoot = True
        self.connectors.append(connector)
        self.root = connector
        # link components
        p.inputs.append(connector)
        connector.inputOf = p

        return connector

    def setLeaf(self,p):
        p = self.getByName(p) if isinstance(p,str) else p

        if not p: print("[ERROR] WFWorkflow.setLeaf(): unknown process !")

        connector = WFConnector.WFConnector(name=p.name+"->"+"leaf")
        connector.isLeaf = True
        self.connectors.append(connector)
        self.leafs.append(connector)
        # link components
        p.outputs.append(connector)
        connector.outputOf = p

        return connector

    # ---------------------------------------------
    # utils method 
    # ---------------------------------------------
    def checkName(self, name):

        res = []

        for p in self.processes:
            if p.name == name: res.append(p)

        for i in self.connectors:
            if i.name == name: res.append(i)

        return res

    def getByName(self, name):
        res = None

        asName = self.checkName(name)

        if len(asName)>0: res = asName[0]

        return res

    # ---------------------------------------------
    # computing
    # ---------------------------------------------
    def compute(self, input=None):
        #
        #
        if not input:
            filename = easygui.fileopenbox(msg="select  image.")
            print("selected image:", filename)
            # reading image
            input= MIMG.Image.readImage(filename)
            if input.isHDR(): input = input.removeZeros(0.5)
        self.root.image=input

        # main loop: 
        # -------------------------------------------------
        awaitingProcess = []
        # reset ready for compute
        for con in self.connectors:
            con.ready = False
        # start with root
        self.root.ready = True
        # put all processes in awaiting Process
        awaitingProcess = self.processes
        # -------------------------------------------------
        while len(awaitingProcess)>0 : 
            # ---------------------------------------------
            # looking for first process in list that all inputs are ready
            pReady = None
            pNotReady = []
            for p in awaitingProcess:
                if p.isReady() and (not pReady):
                    # take first ready for computation
                    pReady = p
                else: # not ready or firest ready already find
                    pNotReady.append(p)
            awaitingProcess = pNotReady
            pReady.compute()

        
    def compile(self):
        # check process input and output
        print("+-----------------------------------")
        print("| COMPILE WORKFLOW")        
        print("+-----------------------------------")
        print("| check process input and output")
        print("+-----------------------------------")

        noInput =[]
        noOutput =[]

        for p in self.processes:
            print("| process:",p.name)
            print("|-----------------------------------")
            print("| input(s)  -> ")

            if len(p.inputs) == 0:
                # no input
                noInput.append(p)
                print("|              ", "no input -> root")

            for input in p.inputs:
                print("|              ",input.name)
            print("| output(s) -> ")

            if len(p.outputs) == 0:
                # no output
                noOutput.append(p)
                print("|              ", "no output -> leaf")
    
            for output in p.outputs:
                print("|              ",output.name)            
            print("+-----------------------------------")

        # root
        if self.root and (len(noInput)== 0):
            print("| root:",self.root.name, ": OK")
            print("+-----------------------------------")
        if (not self.root) and (len(noInput)== 1):
            # add Root
            self.setRoot(noInput[0])
            print("| root -> pocess:",noInput[0].name)
            print("| root:",self.root.name, ": OK")
            print("+-----------------------------------")
            noInput.pop()

        if not self.root or (len(noInput)>1):
            print("[ERROR] root error !")

        # leaf
        for p in noOutput:
            self.setLeaf(p)
            print("| process:",p.name, "- > leaf")
        if len(self.leafs)>0:
            for c in self.leafs:
                print("+-----------------------------------")
                print("| leaf:",c.name)
            print("| leaf: OK")
            print("+-----------------------------------")
        else:
            print("[ERROR] leaf error !")

        # simulate processing
        # -------------------------------------------------
        print("| check workflow")
        print("+-----------------------------------")
        awaitingProcess = []
        # reset ready for compute
        for con in self.connectors:
            con.ready = False
        # start with root
        self.root.ready = True
        # put all processes in awaiting Process
        awaitingProcess = self.processes
        # -------------------------------------------------
        while len(awaitingProcess)>0 : 
            # ---------------------------------------------
            # looking for first process in list that all inputs are ready
            pReady = None
            pNotReady = []
            for p in awaitingProcess:
                if p.isReady() and (not pReady):
                    # take first ready for computation
                    pReady = p
                else: # not ready or firest ready already find
                    pNotReady.append(p)
            awaitingProcess = pNotReady

            # 'compute'
            print("+-----------------------------------")
            print("| computing:", pReady.name)        
            # set ouputs to ready
            for con in pReady.outputs:
                con.ready = True
                print("|    +---->", con.name, "ready")        

        print("+-----------------------------------")
        print("| check workflow: end reached: OK")
        print("+-----------------------------------")

        return self

    # ---------------------------------------------
    # read
    # ---------------------------------------------
    def readWorkflow(filename):

        # read filename
        with open(filename) as json_file: 
            jsonData = json.load(json_file)

        # recover data
        # name
        wf_name = jsonData['name'] if 'name' in jsonData else filename
        # processes and connectors
        processesJSON = jsonData['processes'] if 'processes' in jsonData else []
        connectorsJSON = jsonData['connectors'] if 'connectors' in jsonData else []

        # build process and connectors
        error = False

        # create wf
        wf = WFWorkflow(name=wf_name)

        # build process
        for processJSON in processesJSON:
            # type, name, params
            if 'type' in processJSON :
                type= processJSON['type']
            else:
                print("ERROR[miam.workflow.WFWorflow.readWorkflow(",processJSON,"): no type found !]")
                error = True
            if 'name' in processJSON:
                process_name= processJSON['name']
            else:
                print("ERROR[miam.workflow.WFWorflow.readWorkflow(",processJSON,"): no name found !]")
                error = True
            if 'parameters' in processJSON:
                parameters= processJSON['parameters']
            else:
                print("ERROR[miam.workflow.WFWorflow.readWorkflow(",processJSON,"): no parameters found !]")
                error = True
            # build process
            if not error:

                # add here all new processing class
                # --------------------------------------------
                # ColorSpaceTransform
                # ContrastControl
                # Duplicate
                # ExposureControl
                # NoOp
                # TMO_CCTF
                # TMO_Lightness
                # TMO_Linear
                # Ymap
                # Blend
                # MaskSegmentPercentile
                # Fuse
                # --------------------------------------------

                if type =='ColorSpaceTransform':
                   wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=ColorSpaceTransform.ColorSpaceTransform()).setParameters(eval(parameters))
                    )
                elif type == 'ContrastControl':
                   wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=ContrastControl.ContrastControl()).setParameters(eval(parameters))
                    )
                elif type == 'Duplicate':
                   wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=Duplicate.Duplicate()).setParameters(eval(parameters))
                    )
                elif type == 'ExposureControl':
                   wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=ExposureControl.ExposureControl()).setParameters(eval(parameters))
                    )
                elif type == 'NoOp':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=NoOp.NoOp())
                    )
                elif type == 'TMO_CCTF':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=TMO_CCTF.TMO_CCTF()).setParameters(eval(parameters))
                    )
                elif type == 'TMO_Lightness':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=TMO_Lightness.TMO_Lightness()).setParameters(eval(parameters))
                    )
                elif type == 'TMO_Linear':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=TMO_Linear.TMO_Linear()).setParameters(eval(parameters))
                    )
                elif type == 'Ymap':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=Ymap.Ymap()).setParameters(eval(parameters))
                    )
                elif type == 'Blend':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=Blend.Blend()).setParameters(eval(parameters))
                    )
                elif type == 'MaskSegmentPercentile':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=MaskSegmentPercentile.MaskSegmentPercentile()).setParameters(eval(parameters))
                    )
                elif type == 'Fuse':
                    wf.addProcess(WFProcess.WFProcess(
                        name = process_name, 
                        process=Fuse.Fuse()).setParameters(eval(parameters))
                    )
                else:
                    print("ERROR[miam.workflow.WFWorflow.readWorkflow(",filename,"): unkown process type: ",type," !]")
                    error = True

        # build connector
        for connectorJSON in connectorsJSON:
            # outputProcess -> inputProcess
            if 'outputProcess' in connectorJSON :
                outputProcess = connectorJSON['outputProcess']
            else:
                print("ERROR[miam.workflow.WFWorflow.readWorkflow(",connectorJSON,"): no outputProcess found !]")
                error = True

            if 'inputProcess' in connectorJSON:
                inputProcess= connectorJSON['inputProcess']
            else:
                print("ERROR[miam.workflow.WFWorflow.readWorkflow(",connectorJSON,"): no inputProcess found !]")
                error = True

            # find output and input
            pout = wf.getByName(outputProcess)
            pin = wf.getByName(inputProcess)

            wf.connect(outputProcess=pout, inputProcess=pin, name=outputProcess+"->"+inputProcess)

        # compile and return
        return  wf.compile()

    # ---------------------------------------------
    # display
    # ---------------------------------------------
    def display(self, withHistogram=True):
        # number of leaf image
        nbRes = len(self.leafs)
        # dual display
        dd = 2 if withHistogram else 1

        fig, ax = plt.subplots(2,nbRes*dd)

        # display input
        self.root.image.plot(ax[0,0])
        if withHistogram:
            if self.root.image.isHDR():
                ch = miam.image.channel.channel.Y
            else:
                ch = miam.image.channel.channel.L
            MHIST.Histogram.build(self.root.image, ch).plot(ax[0,1])
        
        # display results
        for i, leaf in enumerate(self.leafs):
            leaf.image.plot(ax[1,i*dd])
            if withHistogram:
                if leaf.image.isHDR():
                    ch = miam.image.channel.channel.Y
                else:
                    ch = miam.image.channel.channel.L
            MHIST.Histogram.build(leaf.image, ch).plot(ax[1,i*dd+1])

        plt.show(block=True)