from tkinter import *
from R2Graph import *
import math
import numpy as np
from copy import deepcopy
from scipy.optimize import minimize

# pip3 install scikit-image
from skimage import measure

SCALEX = 40.
SCALEY = SCALEX

STEPX = 3./SCALEX       # 3 pixels
STEPY = 3./SCALEX

def main():
    points = []
    mouseButtons = []
    objectIDs = []
    levelLineIDs = []

    # Parameters of linear classifier:
    w = np.array([0., 0.])  # Normal vector to the hyperplane
    b = 0.                  # Intercept
    
    def linearClassifier(x):
        """Value of the linear classifier
        at the point x"""
        return w @ x - b

    scaleX = SCALEX; scaleY = SCALEY
    root = Tk()
    root.title("Support Vector Machine")
    root.geometry("860x600")
    panel = Frame(root)
    drawButton = Button(panel, text="Draw")
    clearButton = Button(panel, text="Clear")
    drawArea = Canvas(root, bg="white")
    panel.pack(side=TOP, fill=X)
    drawButton.pack(side=LEFT, padx=4, pady=4)    
    clearButton.pack(side=LEFT, padx=4, pady=4)
    drawArea.pack(side=TOP, fill=BOTH, expand=True, padx=4, pady=4)
    cLabel = Label(panel, text="C:")
    scaleC = Scale(
        panel, from_=0.1, to=20., resolution=0.1,
        orient=HORIZONTAL, length=200
    )
    scaleC.set(10.)
    cLabel.pack(side=LEFT, padx=4, pady=4)
    scaleC.pack(side=LEFT, padx=4, pady=4)

    lossFuncLabel = Label(panel, text="Loss func:")
    lossFuncLabel.pack(side=LEFT, padx=4, pady=4)

    lossFuncIdx = IntVar()  # Control variable for the group of radio buttons
    lossFuncIdx.set(1)
    lossFunction = hingeLoss2
    def setLossFunc():
        nonlocal lossFunction
        idx = lossFuncIdx.get()
        if idx == 0:
            lossFunction = hingeLoss
            print("Using Hinge Loss")
        elif idx == 1:
            lossFunction = hingeLoss2
            print("Using Hinge Loss Square")
        elif idx == 2:
            lossFunction = logisticLoss
            print("Using Logistic Loss")
        # if len(points) > 0:
        #    onDraw()

    hingeLossRadio = Radiobutton(
        panel,
        text = "Hinge loss", variable=lossFuncIdx, value=0,
        command = setLossFunc
    )
    hingeLoss2Radio = Radiobutton(
        panel,
        text = "Hinge loss square", variable=lossFuncIdx, value=1,
        command = setLossFunc
    )
    logisticLossRadio = Radiobutton(
        panel,
        text = "Logistic loss", variable=lossFuncIdx, value=2,
        command = setLossFunc
    )

    hingeLossRadio.pack(side=LEFT, padx=4, pady=4)
    hingeLoss2Radio.pack(side=LEFT, padx=4, pady=4)
    logisticLossRadio.pack(side=LEFT, padx=4, pady=4)

    def classifierFunction(x, y):
        p = np.array([x, y], dtype="float64")
        return linearClassifier(p)

    root.update()
    
    def map(t):
        w = drawArea.winfo_width()
        h = drawArea.winfo_height()
        centerX = w/2.
        centerY = h/2.
        x = centerX + t.x*scaleX
        y = centerY - t.y*scaleY
        return (x, y)
        
    def invmap(p):
        w = drawArea.winfo_width()
        h = drawArea.winfo_height()
        centerX = w/2.
        centerY = h/2.
        x = (p[0] - centerX)/scaleX
        y = (centerY - p[1])/scaleY
        return R2Point(x, y)
        
    def xMin():
        w = drawArea.winfo_width()
        return (-(w/scaleX)/2.)
    def xMax():
        return (-xMin())
    def yMin():
        w = drawArea.winfo_height()
        return (-(w/scaleY)/2.)
    def yMax():
        return (-yMin())
        
    def drawGrid():
        ix0 = int(xMin())
        ix1 = int(xMax())
        x = ix0
        while x <= ix1:
            if x != 0:
                p0 = map(R2Point(x, yMin()))
                p1 = map(R2Point(x, yMax()))
                drawArea.create_line(p0, p1, fill="lightGray", width=1)
            x += 1

        iy0 = int(yMin())
        iy1 = int(yMax())
        y = iy0
        while y <= iy1:
            if y != 0:
                p0 = map(R2Point(xMin(), y))
                p1 = map(R2Point(xMax(), y))
                drawArea.create_line(p0, p1, fill="lightGray", width=1)
            y += 1
            
        # Draw x-axis
        drawArea.create_line(
            map(R2Point(xMin(), 0.)), map(R2Point(xMax(), 0.)),
            fill="black", width=2
        )    
        # Draw y-axis
        drawArea.create_line(
            map(R2Point(0., yMin())), map(R2Point(0., yMax())),
            fill="black", width=2
        )
        
    def onMouseRelease(e):
        # print("Mouse release event:", e)
        p = (e.x, e.y)
        t = invmap(p)
        points.append(t)
        mouseButtons.append(e.num)
        drawPoint(t, e.num)
        
    def drawPoint(t, mouseButton = 1):
        vx = R2Vector(0.3, 0.)
        vy = R2Vector(0., 0.3)
        color = "red"
        if mouseButton == 2:
            color = "green"
        elif mouseButton == 3:
            color = "magenta"
        lineID = drawArea.create_line(
            map(t - vx), map(t + vx), fill=color, width=3
        )
        objectIDs.append(lineID)    
        lineID = drawArea.create_line(
            map(t - vy), map(t + vy), fill=color, width=3
        )    
        objectIDs.append(lineID)
        
    def drawPoints():
        for i in range(len(points)):
            drawPoint(points[i], mouseButtons[i])
        
    def drawLevelLine(f, level = 0., color="blue"):
        nonlocal levelLineIDs
        levelLineIDs.clear()

        x0 = xMin(); x1 = xMax()
        y0 = yMin(); y1 = yMax()
        y = y0;
        a = []
        while y <= y1:
            row = []
            x = x0
            while x <= x1:
                z = f(x, y)
                row.append(z)
                x += STEPX;
            a.append(row)
            y += STEPY
        a = np.array(a)

        contours = measure.find_contours(a, level)
        for contour in contours:
            line = []
            for p in contour:
                x = x0 + p[1]*STEPX
                y = y0 + p[0]*STEPY
                line.append(
                    map(R2Point(x, y))
                )
            lineID = drawArea.create_line(line, width=2, fill=color)
            levelLineIDs.append(lineID)

    def onDraw():
        nonlocal lossFunction, w, b
        if len(points) == 0:
            return

        for i in levelLineIDs:
            drawArea.delete(i)
        levelLineIDs.clear()

        data = [
            [
                np.array([points[i].x, points[i].y]),
                1. if mouseButtons[i] == 1 else (-1.)
            ]
            for i in range(len(points))
        ]
        c = scaleC.get()
        print("C =", c)
        f = functor(data, lossFunc=lossFunction, C=c)

        # Compute an initial approximation
        p0 = R2Point(0., 0.); n0 = 0
        p1 = R2Point(0., 0.); n1 = 0
        for i in range(len(points)):
            if mouseButtons[i] == 1:
                p0 += points[i]
                n0 += 1
            else:
                p1 += points[i]
                n1 += 1
        if n0 > 0:
            p0 *= 1./n0
        if n1 > 0:
            p1 *= 1./n1
        v = p1 - p0
        t = p0 + v*0.5
        w0 = v.normalized()
        b0 = w0*(t - R2Point(0., 0.))
        print("w0 =", w0, " b0 =", b0)
        wb0 = np.array([w0.x, w0.y, b0])
        print("Initial approximation: ", wb0)

        res = minimize(f, x0 = wb0)
        print("minimized: wb =", res.x)
        w = res.x[:-1]
        b = res.x[-1]

        """
        a = R2Vector(res.x[0], res.x[1])
        b = res.x[-1]
        a2 = np.dot(a, a)
        p = R2Point(0., 0.) + a*(b/a2)
        v = a.normal().normalized()
        p0 = p - v*20.
        p1 = p + v*20.
        lineID = drawArea.create_line(
            map(p0), map(p1), width=3, fill="blue"
        )
        """

        drawLevelLine(classifierFunction, 0.)
        
    def clearPicture():
        for i in objectIDs:
            drawArea.delete(i)
        objectIDs.clear()

        for i in levelLineIDs:
            drawArea.delete(i)
        levelLineIDs.clear()

    def onClear():
        clearPicture()
        points.clear()
        mouseButtons.clear()
        
    def onConfigure(e):
        drawArea.delete("all")
        drawGrid()
        drawPoints()        

    drawButton.configure(command = onDraw)
    clearButton.configure(command = onClear)
    drawArea.bind("<ButtonRelease-1>", onMouseRelease)
    drawArea.bind("<ButtonRelease-2>", onMouseRelease)
    drawArea.bind("<ButtonRelease-3>", onMouseRelease)
    drawArea.bind("<Configure>", onConfigure)
    
    drawGrid()
    
    root.mainloop()
    
def hingeLoss(x):
    res = 1. - x
    if res < 0.:
        res = 0.
    return res

def hingeLoss2(x):
    return hingeLoss(x)**2

def logisticLoss(x):
    return math.log(1. + math.exp(-x))

class functor:
    def __init__(self, data, lossFunc = hingeLoss2, C = 10.):
        self.data = data
        self.lossFunc = lossFunc
        self.C = C

    def __call__(self, x):
        # print("Function call, x =", x)
        w = np.array(x[:-1])
        b = x[-1]
        # print("w =", w, "b =", b)
        loss = 0.
        for d in self.data:
            # print("d[0] =", d[0])
            loss += self.lossFunc(d[1] * (w @ d[0] - b))
        if len(self.data) > 0:
            loss /= len(self.data)
        wNorm = np.dot(w, w)
        err = wNorm/2. + self.C*loss
        # print("wNorm =", wNorm, "loss =", loss, "err =", err)
        return err

if __name__ == "__main__":
    main()