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

EPS = 1e-5
MAX_STEPS = 100

def main():
    points = []
    mouseButtons = []
    objectIDs = []
    lineID = None
    iterLineID = None
    currentPnt = (-1)
    currentW = None
    f = None
    
    scaleX = 40.; scaleY = 40.
    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")

    startButton = Button(panel2, text="Start")
    nextButton = Button(panel2, text="Step")
    go100Button = Button(panel2, text="100 steps")
    go1000Button = Button(panel2, text="1000 steps")
    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=100
    )
    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(0)
    lossFunction = hingeLoss

    def setLossFunc():
        nonlocal lossFunction
        idx = lossFuncIdx.get()
        if idx == 0:
            lossFunction = hingeLoss
            print("Using Hinge Loss")
        elif idx == 1:
            lossFunction = logisticLoss
            print("Using Logistic Loss")
        elif idx == 2:
            lossFunction = hingeLoss2
            print("Using Hinge Loss Square")
        # if len(points) > 0:
        #    onDraw()

    hingeLossRadio = Radiobutton(
        panel,
        text = "Hinge", variable=lossFuncIdx, value=0,
        command = setLossFunc
    )
    logisticLossRadio = Radiobutton(
        panel,
        text = "LogLoss", variable=lossFuncIdx, value=1,
        command = setLossFunc
    )

    ''' BAD Function! Values grow too fast, resulting in overfull...
    hingeLoss2Radio = Radiobutton(
        panel,
        text = "HingeSquare", variable=lossFuncIdx, value=2,
        command = setLossFunc
    )
    '''

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

    stepLabel = Label(panel2, text="Step\n(alpha):", fg="DarkBlue")
    stepScale = Scale(
        panel2, from_=0.001, to=1, resolution=0.001,
        orient=HORIZONTAL, fg="DarkBlue" # , length=200
    )
    stepScale.set(0.1)

    etaLabel = Label(panel2, text="Step diminish\nrate (eta):", fg="DarkMagenta")
    etaScale = Scale(
        panel2, from_=0., to=2., resolution=0.01,
        orient=HORIZONTAL, fg="DarkMagenta" # , length=200
    )
    etaScale.set(0.1)

    stepNumber = 0

    miniBatch = 0
    miniBatchSize = 4
    miniBatchPortion = 0.2

    miniBatchLabel = Label(panel2, text="Mini batch:", fg="DarkGreen")
    miniBatchScale = Scale(
        panel2, from_=1, to=16, resolution=1,
        orient=HORIZONTAL, fg="DarkGreen", length=100
    )
    miniBatchScale.set(4)

    def setMiniBatchSize():
        nonlocal miniBatchSize, miniBatchPortion
        n = len(points)
        miniBatchSize = int(n*miniBatchPortion)
        if miniBatchSize < 4:
            miniBatchSize = 4
        if miniBatchSize > n:
            miniBatchSize = n
        miniBatchScale.set(miniBatchSize)

    startButton.pack(side=LEFT, padx=4, pady=4)
    nextButton.pack(side=LEFT, padx=4, pady=4)
    go100Button.pack(side=LEFT, padx=4, pady=4)
    go1000Button.pack(side=LEFT, padx=4, pady=4)

    miniBatchLabel.pack(side=LEFT, padx=4, pady=4)
    miniBatchScale.pack(side=LEFT, padx=4, pady=4)

    stepLabel.pack(side=LEFT, padx=4, pady=4)
    stepScale.pack(side=LEFT, padx=4, pady=4)
    etaLabel.pack(side=LEFT, padx=4, pady=4)
    etaScale.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)
        setMiniBatchSize()
        
    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 lineID, iterLineID, lossFunction, f
        if len(points) == 0:
            return

        if lineID != None:
            drawArea.delete(lineID)
            lineID = None
            return

        data = [
            np.array([
                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
        w = v.normalized()
        b = w*(t - R2Point(0., 0.))
        x0 = [w.x, w.y, b]
        print("Initial approximation: ", x0)

        res = minimize(f, x0 = x0)
        print("minimized: w =", res.x)
        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"
        )
        
    def onStart():
        nonlocal lineID, iterLineID, lossFunction, currentPnt, currentW, f
        nonlocal miniBatch, miniBatchSize, miniBatchPortion, stepNumber

        if len(points) == 0:
            return

        print("onStart")
        currentPnt = 0

        if iterLineID != None:
            drawArea.delete(iterLineID)
            iterLineID = None

        data = [
            np.array([
                np.array([points[i].x, points[i].y]),
                1. if mouseButtons[i] == 1 else (-1.)
            ])
            for i in range(len(points))
        ]
        random.shuffle(data)
        print("Data:\n", data)

        miniBatch = 0
        miniBatchSize = miniBatchScale.get()

        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:
                p1 += points[i]
                n1 += 1
            else:
                p0 += points[i]
                n0 += 1
        if n0 > 0:
            p0 *= 1./n0
        if n1 > 0:
            p1 *= 1./n1
        v = p1 - p0
        t = p0 + v*0.5
        w = v.normalized()
        b = w*(t - R2Point(0., 0.))
        x0 = np.array([w.x, w.y, b])
        print("Initial approximation: ", x0)
        currentW = x0.copy()
        stepNumber = 0

        a = R2Vector(currentW[0], currentW[1])
        b = currentW[-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.
        iterLineID = drawArea.create_line(
            map(p0), map(p1), width=2, fill="DarkGreen"
        )
        
    def onNext():
        nonlocal iterLineID, lossFunction, currentPnt, currentW, f
        nonlocal miniBatch, miniBatchSize, miniBatchPortion, stepNumber

        if currentW is None:
            onStart()

        stepNumber += 1

        c = scaleC.get()
        # print("C =", c)
        n = len(points)

        if iterLineID != None:
            drawArea.delete(iterLineID)
            iterLineID = None

        # print("currentW =", currentW)
        
        alpha = stepScale.get()
        eta = etaScale.get()

        # print("miniBatch=", miniBatch, "size=", miniBatchSize)
        # print("alpha=", alpha, "eta=", eta)

        xBatch = []
        i = 0
        while i < miniBatchSize:
            j = (miniBatch + i)%n
            xBatch.append(
                f.data[j]
            )
            i += 1
        miniBatch = (miniBatch + miniBatchSize)%n

        # print("xBatch =", xBatch)
        
        g = f.gradient(currentW, xBatch)
        
        print("gradient =", g)
        
        stepSize = alpha/(1. + eta*stepNumber)
        print("stepSize=", stepSize)

        currentW -= g*stepSize
            
        print("currentW =", currentW)

        # Draw a current line
        a = R2Vector(currentW[0], currentW[1])
        b = currentW[-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.
        iterLineID = drawArea.create_line(
            map(p0), map(p1), width=2, fill="DarkGreen"
        )
        return np.linalg.norm(g)

    def onGo100():
        onGo(100)

    def onGo1000():
        onGo(1000)

    def onGo(maxSteps = 100):
        steps = 0
        while steps < maxSteps:
            gradientNorm = onNext()
            if gradientNorm <= EPS:
                break
            steps += 1
            if steps%10 == 0:
                root.update()

    def clearPicture():
        nonlocal lineID, iterLineID
        for i in objectIDs:
            drawArea.delete(i)
        objectIDs.clear()
        if lineID != None:
            drawArea.delete(lineID)
            lineID = None
        if iterLineID != None:
            drawArea.delete(iterLineID)
            iterLineID = None
            
    def onClear():
        clearPicture()
        points.clear()
        mouseButtons.clear()
        
    def onConfigure(e):
        drawArea.delete("all")
        drawGrid()
        drawPoints()        

    drawButton.configure(command = onDraw)
    startButton.configure(command = onStart)
    nextButton.configure(command = onNext)
    go100Button.configure(command = onGo100)
    go1000Button.configure(command = onGo1000)
    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 = hingeLoss, C = 10., h = 0.001):
        self.data = deepcopy(data)
        self.lossFunc = lossFunc
        self.C = C
        self.h = h

    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

    def miniBatchValue(self, w, xBatch):
        '''Compute the loss function for the given mini-batch'''
        # print("Function call, xBatch =", xBatch)
        ww = w[:-1]
        wb = w[-1]
        # print("ww =", ww, "wb =", wb)
        loss = 0.
        for d in xBatch:
            # print("in miniBatch: d =", d)
            loss += self.lossFunc(d[1] * (np.dot(ww, d[0]) - wb))
        if len(xBatch) > 0:
            loss /= len(xBatch)
        wNorm = np.dot(ww, ww)
        err = wNorm/2. + self.C*loss
        # print("wNorm =", wNorm, "loss =", loss, "err =", err)
        return err

    def gradient(self, w, xBatch):
        '''Compute the gradient of loss function for the given mini-batch'''
        n = len(w)      # Dimension
        g = np.array([0.]*n)
        for i in range(n):
            # Compute approximately the value of 
            # partial derivative d/d w_i
            w_i = w[i]

            w[i] = w_i - self.h
            v0 = self.miniBatchValue(w, xBatch)

            w[i] = w_i + self.h
            v1 = self.miniBatchValue(w, xBatch)

            w[i] = w_i

            g[i] = (v1 - v0)/(2.*self.h)
        return g

if __name__ == "__main__":
    main()