from tkinter import *
from R2Graph import *
import math
import random
import numpy as np

EPS = 1e-5
MAX_STEPS = 1000
redrawSteps = 10

def func(p, x):
    return np.polyval(p, x)

def main():
    points = []
    mouseButtons = []
    objectIDs = []
    funcDrawn = False
    plotID = (-1)
    iterPlotID = (-1)
    pol = np.array([])
    iterPol = np.array([])
    stepNumber = 0

    scaleX = 40.; scaleY = 40.
    root = Tk()
    root.title("Least Squares")
    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)

    degree = 1
    degreeLabel = Label(panel, text="Degree:")
    scale = Scale(panel, from_=1, to=10, orient=HORIZONTAL)
    scale.set(degree)
    degreeLabel.pack(side=LEFT, padx=4, pady=4)
    scale.pack(side=LEFT, padx=4, pady=4)

    drawPreciseSolution = IntVar(value=0)
    drawPrecise = False
    def onDrawPrecise():
        nonlocal drawPrecise
        draw_precise_solution = drawPreciseSolution.get()
        drawPrecise = (draw_precise_solution != 0)
        onDraw()
    drawPreciseSolutionButton = Checkbutton(
        panel, text="Draw Precise Solution", variable=drawPreciseSolution,
        command = onDrawPrecise
    )
    drawPreciseSolutionButton.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.5)

    stepNumber = 0

    miniBatch = 0
    miniBatchSize = 4
    miniBatchPortion = 0.2

    miniBatchLabel = Label(panel2, text="Mini\nbatch:", 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 = Button(panel2, text="Start")
    nextButton = Button(panel2, text="Step")
    goButton = Button(panel2, text="100 steps")
    goOptimalButton = Button(panel2, text="1000 steps")

    startButton.pack(side=LEFT, padx=4, pady=4)
    nextButton.pack(side=LEFT, padx=4, pady=4)
    goButton.pack(side=LEFT, padx=4, pady=4)
    goOptimalButton.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)

    def lossFunction():
        nonlocal miniBatch, miniBatchSize
        n = len(points)
        s = 0.
        for i in range(miniBatchSize):
            j = (miniBatch + i)%n
            p = points[j]
            v = np.polyval(iterPol, p.x)
            s += (v - p.y)**2
        s /= miniBatchSize
        return s

    def lossFunctionGradient():
        nonlocal iterPol, miniBatch, miniBatchSize
        n = len(points)
        d = len(iterPol) - 1    # Degree of polynomial
        s = 0.
        g = np.array([0.]*(d + 1))
        for b in range(miniBatchSize):
            i = (miniBatch + b)%n

            # Loss function:
            #     L(w) = 1/m sum_i (p(w, x_i) - y_i)^2
            # Partial derivatives:
            #     dL(w)/dw_j = 2/m sum_i (p(w, x_i) - y_i)*dp(w, x_i)/dw_j =
            #         = 2/m sum_i (p(w, x_i) - y_i)*x_i^(d-j)

            x = points[i].x; y = points[i].y
            p = np.polyval(iterPol, x)
            xPow = 1.
            for j in range(d, -1, -1):
                g[j] += (p - y)*xPow
                xPow *= x
        g *= 2./miniBatchSize

        '''
        # Check whether this is correct
        g0 = np.array([0.]*(d+1))
        h = 0.001
        savedPol = iterPol.copy()
        for j in range(d+1):
            iterPol[j] += h
            v1 = lossFunction()
            iterPol[j] -= 2.*h
            v0 = lossFunction()
            g0[j] = (v1 - v0)/(2.*h)
        iterPol = savedPol

        print(" g:", g)
        print("g0:", g0)
        '''

        return g

    def onStart():
        nonlocal iterPlotID, iterPol
        nonlocal miniBatch, miniBatchSize
        if len(points) == 0:
            return

        print("onStart")

        if iterPlotID >= 0:
            drawArea.delete(iterPlotID)
            iterPlotID = (-1)

        random.shuffle(points)
        print("Points:\n", points)

        stepNumber = 0
        miniBatch = 0
        miniBatchSize = miniBatchScale.get()

        # Compute an initial approximation
        degree = scale.get()

        '''
        if (
            len(iterPol) != degree + 1 or
            (len(iterPol) > 0 and math.isnan(iterPol[0])) or
            (len(iterPol) > 0 and abs(iterPol[0]) >= 1e10)
        ):
            iterPol = np.array([0.]*(degree + 1))
        '''

        # Let it be a constant that equals to mean value of all y's
        s = 0.
        iterPol = np.array([0.]*(degree + 1))
        if len(points) > 0:
            for p in points:
                s += p[1]
            s /= len(points)
            iterPol[-1] = s

        # print("iterPol =", iterPol)
        iterPlotID = plotFunction(iterPol, color="green")

    def onNext():
        nonlocal iterPlotID, iterPol, stepNumber
        nonlocal miniBatch, miniBatchSize
        print("onNext")
        if iterPlotID < 0:
            onStart()
        n = len(points)

        v = lossFunction()
        print("loss function:", v)
        g = lossFunctionGradient()
        print("gradient:", g)

        alpha = stepScale.get()
        eta = etaScale.get()
        stepNumber += 1
        # stepSize = alpha/(1. + eta*(stepNumber**0.75))
        stepSize = alpha/(1. + eta*stepNumber)
        # print("stepSize=", stepSize)

        iterPol -= g*stepSize
        print("iterPol =", iterPol)
        if iterPlotID >= 0:
            drawArea.delete(iterPlotID)
        iterPlotID = plotFunction(iterPol, color="green")
        miniBatch = (miniBatch + miniBatchSize)%n
        return np.linalg.norm(g)

    def onGo():
        doIterations(100)

    def onOptimal():
        doIterations(MAX_STEPS)

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

    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 plotFunction(pol, color="blue"):
        if len(pol) == 0:
            return (-1)
        t = R2Point(xMin(), func(pol, xMin()))
        dx = 0.05
        path = []
        while t.x <= xMax():
            path.append(map(t))
            t.x += dx
            t.y = func(pol, t.x)
        lineID = drawArea.create_line(path, fill=color, width=2)
        return lineID
        
    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.2, 0.)
        vy = R2Vector(0., 0.2)
        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=2
        )
        objectIDs.append(lineID)    
        lineID = drawArea.create_line(
            map(t - vy), map(t + vy), fill=color, width=2
        )    
        objectIDs.append(lineID)
        
    def drawPoints():
        for i in range(len(points)):
            drawPoint(points[i], mouseButtons[i])
        
    def onDraw():
        nonlocal funcDrawn, pol, degree, plotID, drawPrecise
        # clearPicture()
        if funcDrawn:
            drawArea.delete(plotID)
            plotID = (-1)

        m = len(points)
        if m == 0:
            return

        if drawPrecise and len(points) > 0:
            degree = scale.get()
            y = np.array([points[i].y for i in range(m)])
            print("y =", y)
            a = np.vander([points[i].x for i in range(m)], degree + 1)
            print("a =", a)
            pinv_a = np.linalg.pinv(a)
            print("pinv_a =", pinv_a)
            pol = pinv_a.dot(y)
            print("pol =", pol)
            plotID = plotFunction(pol, color="blue")
            funcDrawn = True
        
    def clearPicture():
        for i in objectIDs:
            drawArea.delete(i)
        objectIDs.clear()
        if plotID >= 0:
            drawArea.delete(plotID)
        if iterPlotID >= 0:
            drawArea.delete(iterPlotID)

    def onClear():
        nonlocal funcDrawn
        clearPicture()
        points.clear()
        mouseButtons.clear()
        funcDrawn = False
        stepNumber = 0
        miniBatch = 0
        
    def onConfigure(e):
        drawArea.delete("all")
        drawGrid()
        if funcDrawn:
            plotFunction()
        drawPoints()        

    drawButton.configure(command = onDraw)
    clearButton.configure(command = onClear)
    scale.configure(command = lambda e: onDraw())

    startButton.configure(command = onStart)
    nextButton.configure(command = onNext)
    goButton.configure(command = onGo)
    goOptimalButton.configure(command = onOptimal)

    drawArea.bind("<ButtonRelease-1>", onMouseRelease)
    drawArea.bind("<ButtonRelease-2>", onMouseRelease)
    drawArea.bind("<ButtonRelease-3>", onMouseRelease)
    drawArea.bind("<Configure>", onConfigure)
    
    drawGrid()
    # plotFunction()
    
    root.mainloop()
    
if __name__ == "__main__":
    main()