from tkinter import *
from R2Graph import *
import math
import numpy as np
import itertools
from copy import deepcopy
from scipy.optimize import minimize
from sklearn.svm import SVC

SCALEX = 40.
SCALEY = SCALEX

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

polynomialDegree = 2
numFeatures = 6         # For quadratic polynomial
# numFeatures = 10      # For cubic polynomial
hyperPlane = np.array([0.]*(numFeatures+1))
coef0 = 1.

classifier = SVC(kernel="poly", degree=polynomialDegree, coef0=coef0)
useKernelClassifier = False

def classifierValue(p):
    if not useKernelClassifier:
        w = hyperPlane[:-1]
        b = hyperPlane[-1]
        features = extendedFeatures(p, polynomialDegree)
        return np.dot(w, features) - b
    else:
        x = np.array([[p.x, p.y]])
        # return classifier.predict(x)
        return classifier.decision_function(x)

def levelFunction(x, y):
    if not useKernelClassifier:
        return classifierValue(R2Point(x, y))
    else:
        p = np.array([[x, y]])
        # p = p.reshape(1, -1)
        # return classifier.predict(p)
        return classifier.decision_function(p)

def monomials(x, d):
    '''Return an array of all monomials in variables x[0], ..., x[n-1]
    of degree <= d'''
    res = [1.]
    for deg in range(1, d+1):
        for c in itertools.combinations_with_replacement(x, deg):
            res.append(np.prod(c))
    return np.array(res)

def extendedFeatures(p, d):
    return monomials(p, d)

def extendFeatures(p):
    '''p is a point: p = (x, y)
    Add the polynomial features:
    (1, x, y, x^2, xy, y^2)'''
    return extendFeaturesCubic(p)     ### For cubic polynomial

    x = p[0]; y = p[1]
    return np.array([
        1.,
        x, y, 
        x*x, x*y, y*y
    ])

def extendFeaturesCubic(p):
    '''p is a point: p = (x, y)
    Add the polynomial features:
    (1, x, y, x^2, xy, y^2, 
    x^3, x^2y, xy^2, y^3)'''
    x = p[0]; y = p[1]
    return np.array([
        1., 
        x, y, 
        x*x, x*y, y*y,
        x*x*x, x*x*y, x*y*y, y*y*y
    ])

def main():
    points = []
    mouseButtons = []
    objectIDs = []
    levelLineIDs = []
    
    scaleX = SCALEX; scaleY = SCALEY
    root = Tk()
    root.title("Support Vector Machine")
    root.geometry("900x600")
    panel = Frame(root)
    panel2 = Frame(root)
    drawButton = Button(panel, text="Draw")
    clearButton = Button(panel, text="Clear")
    drawArea = Canvas(root, bg="white")
    panel.pack(side=TOP, fill=X)
    panel2.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()

    polynomialLabel = Label(panel2, text="Polynomial degree:")
    polynomialScale = Scale(
        panel2, from_=1, to=5, resolution=1,
        orient=HORIZONTAL, length=200
    )
    polynomialScale.set(2)
    polynomialLabel.pack(side=LEFT, padx=4, pady=4)
    polynomialScale.pack(side=LEFT, padx=4, pady=4)

    useKernel = IntVar(value = 0)
    useKernelButton = Checkbutton(panel2, text="Use kernel", variable=useKernel)
    useKernelButton.pack(side=LEFT, padx=4, pady=4)

    kernelIdx = IntVar(value = 0)

    kernel = "poly"
    def setKernel():
        nonlocal kernel
        idx = kernelIdx.get()
        if idx == 0:
            kernel = "poly"
            print("Using 'poly' kernel")
        elif idx == 1:
            kernel = "rbf"
            print("Using 'rbf' kernel")
        elif idx == 2:
            kernel = "sigmoid"
            print("Using 'sigmoid' kernel")

        # if len(points) > 0:
        #     onDraw()

    kernelPolyRadio = Radiobutton(
        panel2,
        text = "poly", variable=kernelIdx, value=0,
        command = setKernel
    )
    kernelRbfRadio = Radiobutton(
        panel2,
        text = "rbf", variable=kernelIdx, value=1,
        command = setKernel
    )
    kernelSigmoidRadio = Radiobutton(
        panel2,
        text = "sigmoid", variable=kernelIdx, value=2,
        command = setKernel
    )

    def onChangeDegree():
        # if len(points) > 0:
        #    onDraw()
        pass

    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)

    kernelPolyRadio.pack(side=LEFT, padx=4, pady=4)
    kernelRbfRadio.pack(side=LEFT, padx=4, pady=4)
    kernelSigmoidRadio.pack(side=LEFT, padx=4, pady=4)

    coef0Label = Label(panel2, text="Coef0:")
    scaleCoef0 = Scale(
        panel2, from_=0., to=10., resolution=0.1,
        orient=HORIZONTAL, length=100
    )
    scaleCoef0.set(coef0)
    coef0Label.pack(side=LEFT, padx=4, pady=4)
    scaleCoef0.pack(side=LEFT, padx=4, pady=4)

    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 onDraw():
        nonlocal levelLineIDs, lossFunction, useKernel, kernel
        global hyperPlane, numFeatures, polynomialDegree, coef0
        global useKernelClassifier, classifier
        if len(points) == 0:
            return

        deleteLevelLine()

        polynomialDegree = polynomialScale.get()
        c = scaleC.get()
        coef0 = scaleCoef0.get()
        print(
            "C =", c, "kernel =", kernel, 
            "degree =", polynomialDegree, "coef0 =", coef0
        )

        use_kernel = useKernel.get()
        if use_kernel != 0:
            useKernelClassifier = True
            classifier = SVC(
                C=c, kernel=kernel, degree=polynomialDegree, coef0=coef0
            )
            data = [
                np.array(
                   [points[i].x, points[i].y],
                )
                for i in range(len(points))
            ]
            # data = np.array(data).reshape(len(data), -1)
            data = np.array(data)
            y = [
                (1. if mouseButtons[i] == 1 else (-1.))
                for i in range(len(points))
            ]
            y = np.array(y)
            classifier.fit(data, y)
        else:
            useKernelClassifier = False
            data = [
                np.array([
                    # [points[i].x, points[i].y],
                    extendedFeatures(points[i], polynomialDegree),
                    1. if mouseButtons[i] == 1 else (-1.)
                ])
                for i in range(len(points))
            ]
            numFeatures = len(data[0][0])
            print("numFeatures =", numFeatures)
            hyperPlane = np.array([0.]*(numFeatures+1))

            f = functor(data, lossFunc=lossFunction, C=c)

            # Initial approximation
            x0 = [0.]*(numFeatures + 1)
            # print("Initial approximation: x0 =", x0)

            res = minimize(f, x0 = x0)
            print("minimized: w =", res.x)
            hyperPlane = res.x

        drawLevelLine(levelFunction)

    def deleteLevelLine():
        nonlocal levelLineIDs
        for i in levelLineIDs:
            drawArea.delete(i)
        levelLineIDs.clear()

    def clearPicture():
        for i in objectIDs:
            drawArea.delete(i)
        objectIDs.clear()
        deleteLevelLine()
            
    def onClear():
        clearPicture()
        points.clear()
        mouseButtons.clear()
        
    def onConfigure(e):
        drawArea.delete("all")
        drawGrid()
        drawPoints()        

    def drawLevelLine(f, level = 0., color="blue"):
        nonlocal levelLineIDs
        levelLineIDs.clear()

        x0 = xMin(); x1 = xMax()
        y0 = yMin(); y1 = yMax()
        y = y0;
        v1 = []
        while y < y1:
            if y > y0:
                v0 = v1.copy()
            else:
                v0 = []
            v1.clear()
            x = x0;
            while x <= x1 + STEPX:
                if y <= y0:
                    z = f(x, y)
                    v0.append(z)
                z = f(x, y + STEPY)
                v1.append(z)
                x += STEPX;

            x = x0;
            ix = 0
            while x < x1:
                # z0--z2
                # |  /
                # z1
                z0 = v0[ix]
                z2 = v0[ix + 1]
                z1 = v1[ix]
                z3 = v1[ix + 1]

                # print("z0=", z0, "z1=", z1, "z2=", z2, "z3=", z3)
                # print("level=", level)

                points = [False]*5
                # z0--0--z2
                # |     /|
                # 1   2  4
                # |  /   |
                # z1--3--z3
                if (z0 <= level and level < z2) or (z0 >= level and level > z2):
                    xx = x + STEPX*abs(level - z0)/abs(z2 - z0)
                    p02 = R2Point(xx, y)
                    points[0] = True
                if (z0 <= level and level < z1) or (z0 >= level and level > z1):
                    yy = y + STEPY*abs(level - z0)/abs(z1 - z0)
                    p01 = R2Point(x, yy)
                    points[1] = True
                if (z1 <= level and level < z2) or (z1 >= level and level > z2):
                    xx = x + STEPX*abs(level - z1)/abs(z2 - z1)
                    yy = (y + STEPY) - STEPY*abs(level - z1)/abs(z2 - z1)
                    p12 = R2Point(xx, yy)
                    points[2] = True
                #    z2
                #   / |
                # z1--z3
                if (z1 <= level and level < z3) or (z1 >= level and level > z3):
                    xx = x + STEPX*abs(level - z1)/abs(z3 - z1)
                    p13 = R2Point(xx, y + STEPY)
                    points[3] = True
                if (z2 <= level and level < z3) or (z2 >= level and level > z3):
                    yy = y + STEPY*abs(level - z2)/abs(z3 - z2)
                    p23 = R2Point(x + STEPX, yy)
                    points[4] = True

                # z0--0--z2
                # |     /|
                # 1   2  4
                # |  /   |
                # z1--3--z3
                # Upper triangle:
                if points[0] and points[1]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p02), map(p01), width=2, fill=color)
                    )
                if points[0] and points[2]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p02), map(p12), width=2, fill=color)
                    )
                if points[1] and points[2]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p01), map(p12), width=2, fill=color)
                    )
                # Lower triangle:
                if points[2] and points[4]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p12), map(p23), width=2, fill=color)
                    )
                if points[2] and points[3]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p12), map(p13), width=2, fill=color)
                    )
                if points[3] and points[4]:
                    levelLineIDs.append(
                        drawArea.create_line(map(p13), map(p23), width=2, fill=color)
                    )

                x += STEPX;
                ix += 1

            y += STEPY;


    drawButton.configure(command = onDraw)
    clearButton.configure(command = onClear)
    # polynomialScale.configure(command = lambda e: onDraw())
    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.):
        # print("data:", data)

        self.data = deepcopy(data)
        self.lossFunc = lossFunc
        self.C = C

    def __call__(self, x):
        # print("Function call, x =", x)
        w = x[:-1]
        b = x[-1]
        # print("w =", w, "b =", b)
        loss = 0.
        for d in self.data:
            loss += self.lossFunc(d[1] * (np.dot(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()