from tkinter import *
from tkinter import filedialog
from PIL import Image, ImageTk
import os
import random
import numpy as np
from sklearn.tree import DecisionTreeRegressor


def main():
    image = None
    imageFilePath = None
    scaledImage = None
    photoImage = None
    imageSize = (0, 0)
    zoom = 1.

    root = Tk()
    root.title("Transform Image with Decision Trees")
    root.geometry("900x600")

    panel = Frame(root)
    panel.pack(side=TOP, fill=X)

    browseLabel = Label(panel, text="Image:")
    imagePath = StringVar()
    imageEntry = Entry(
        panel, bg="white", textvariable=imagePath, width=40
    )
    browseButton = Button(panel, text="Browse...")
    loadButton = Button(panel, text="Reload image")

    browseLabel.pack(side=LEFT, padx=4, pady=4)
    imageEntry.pack(side=LEFT, fill=X, expand=True, padx=4, pady=4)
    browseButton.pack(side=LEFT, padx=4, pady=4)
    loadButton.pack(side=LEFT, padx=4, pady=4)

    panel2 = Frame(root)
    panel2.pack(side=TOP, fill=X)

    depthLabel = Label(panel2, text="Maximal tree depth:")
    depthVar = IntVar(value=5)
    depthEntry = Entry(
        panel2, bg="white", textvariable=depthVar, width=10
    )
    transformButton = Button(panel2, text="Transform")

    saveButton = Button(panel2, text="Save image")

    depthLabel.pack(side=LEFT, padx=4, pady=4)
    depthEntry.pack(side=LEFT, padx=4, pady=4)
    transformButton.pack(side=LEFT, padx=4, pady=4)
    saveButton.pack(side=RIGHT, padx=4, pady=4)

    drawArea = Canvas(root, bg="white")
    drawArea.pack(side=TOP, fill=BOTH, expand=True, padx=4, pady=4)

    def browse():
        nonlocal imageFilePath
        path = filedialog.askopenfilename(
            parent=root, title="Select an image file:"
        )
        # print("Selected:", imageFile)
        if path != None:
            imageFilePath = path
            imagePath.set(path)
            loadImage()

    browseButton.configure(command=browse)

    def loadImage():
        nonlocal image, scaledImage, photoImage, imageSize, zoom
        path = imagePath.get()
        if path == None or path == "":
            return
        im = Image.open(path).convert("RGB")
        if im == None:
            print("Cannot open an image file")
            return
        image = im
        imageSize = image.size

        scaleImage()
        redraw()

    loadButton.configure(command=loadImage)
    imageEntry.bind("<Return>", lambda e: loadImage())

    def scaleImage():
        nonlocal scaledImage, photoImage, zoom
        w = drawArea.winfo_width()
        h = drawArea.winfo_height()
        if w == 0:
            w = 1
        if h == 0:
            h = 1

        # Scale the big image to fit in the window
        zoom = 1.
        image_width = float(imageSize[0]); image_height = float(imageSize[1])
        if image_width > w:
            zoom = w/image_width
            image_width *= zoom
            image_height *= zoom
        if image_height > h:
            zoom2 = h/image_height
            image_width *= zoom2
            image_height *= zoom2
            zoom *= zoom2
        if zoom < 1.:
            scaledImage = image.resize(
                (int(image_width), int(image_height))
            )
            photoImage = ImageTk.PhotoImage(scaledImage)
        else:
            photoImage = ImageTk.PhotoImage(image)

    def redraw():
        drawArea.delete("all")
        if photoImage != None:
            w = drawArea.winfo_width()
            h = drawArea.winfo_height()
            drawArea.create_image(w/2, h/2, anchor=CENTER, image=photoImage)

    def onConfigure(e):
        if image == None:
            return
        scaleImage()
        redraw()

    drawArea.bind("<Configure>", onConfigure)

    def onTransform():
        '''Transform the image using Decision Trees'''
        if image == None:
            return
        depth = depthVar.get()
        if depth < 2:
            depth = 2
            depthVar.set(value = depth)

        showWatchCursor();

        redModel = DecisionTreeRegressor(max_depth = depth)
        greenModel = DecisionTreeRegressor(max_depth = depth)
        blueModel = DecisionTreeRegressor(max_depth = depth)

        imageMatrix = list(image.getdata())
        (imageWidth, imageHeight) = imageSize
        X = np.array(
            [[0., 0.] for i in range(imageWidth*imageHeight)]
        )
        YRed = np.array(
            [0. for i in range(imageWidth*imageHeight)]
        )
        YGreen = np.array(
            [0. for i in range(imageWidth*imageHeight)]
        )
        YBlue = np.array(
            [0. for i in range(imageWidth*imageHeight)]
        )

        i = 0
        for y in range(imageHeight):
            for x in range(imageWidth):
                # srcPixel = imageMatrix[y*imageWidth + x]
                srcPixel = imageMatrix[i]
                X[i][0] = float(y)
                X[i][1] = float(x)
                YRed[i] = srcPixel[0]
                YGreen[i] = srcPixel[1]
                YBlue[i] = srcPixel[2]
                i += 1

        redModel.fit(X, YRed)
        greenModel.fit(X, YGreen)
        blueModel.fit(X, YBlue)

        YRed_predicted = redModel.predict(X)
        YGreen_predicted = greenModel.predict(X)
        YBlue_predicted = blueModel.predict(X)

        # We use the white background
        newPixels = [(255, 255, 255) for i in range(imageWidth*imageHeight)]
        i = 0
        for y in range(imageHeight):
            for x in range(imageWidth):
                newPixels[i] = (
                    round(YRed_predicted[i]),
                    round(YGreen_predicted[i]),
                    round(YBlue_predicted[i])
                )
                i += 1

        image.putdata(newPixels)

        scaleImage()
        redraw()
        restoreCursor()

    transformButton.configure(command = onTransform)
    depthEntry.bind("<Return>", lambda e: onTransform())

    def onSave():
        if image == None:
            return

        # Add the "_edited" suffix to the file name
        # To do this, parse the file path

        # 1. Find the file extension
        pos = len(imageFilePath) - 1
        while ( 
            pos > 0 and imageFilePath[pos] != '.'
            and imageFilePath[pos] != '/'
            and imageFilePath[pos] != '\\'
        ):
            pos -= 1
        if imageFilePath[pos] == '/' or imageFilePath[pos] == '\\':
            pos += 1
        if imageFilePath[pos] != '.':
            dotPos = pos
            fileExt = ""
            pos = len(imageFilePath) - 1
        else:
            dotPos = pos
            fileExt = imageFilePath[dotPos:]
            pos -= 1

        # 2. Find the file name and directory
        while ( 
            pos > 0 and
            imageFilePath[pos] != '/' and
            imageFilePath[pos] != '\\'
        ):
            pos -= 1
        if imageFilePath[pos] == '/' or imageFilePath[pos] == '\\':
            namePos = pos + 1
            fileName = imageFilePath[namePos:dotPos]
            directory = imageFilePath[:(namePos-1)]
        else:
            fileName = imageFilePath[:dotPos]
            directory = ""

        print(
            "directory:", directory, 
            "fileName:", fileName, 
            "fileExt:", fileExt
        )

        dstFileName = fileName + "_edited" + fileExt

        path = filedialog.asksaveasfilename(
            parent=root, title="Save an image file:",
            initialdir = directory,
            initialfile = dstFileName
            # defaultextension = ".png"
        )
        print("Selected to save as:", path)
        image.save(path)

    saveButton.configure(command = onSave)

    def showWatchCursor():
        try:
            root.configure(cursor="watch") # Wait cursor in all system,
        except:                            #      but NOT in MS Windows
            try:
                root.configure(cursor="wait") # Wait cursor in MS Windows
            except:
                pass
        root.update()

    def restoreCursor():
        root.configure(cursor="")

    root.mainloop()

if __name__ == "__main__":
    main()