from sympy import *
import sympy as sy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import math

def drawLine(pa, pb):
    lx = np.zeros(2)
    ly = np.zeros(2)

    lx[0] = -100
    lx[1] = 100

    ly[0] = pa*lx[0] + pb
    ly[1] = pa*lx[1] + pb

    return lx, ly

def drawCircle(sx, sy):
    cx = np.zeros(100)
    cy = np.zeros(100)
    step = (2*math.pi)/100
    for i in range(0, 100):
        cx[i] = sx+math.sin(i*step)
        cy[i] = sy+math.cos(i*step)

    return cx, cy


def quadraticcase2(t0, tend, dt, k, fun1, fun2, teta, x0):  # solver for the case of quadratic energy
    # fun is a linear function with time in argument; dt is the change in time
    f1 = sy.sympify(fun1)  # converting string into function with t in argument
    f2 = sy.sympify(fun2)
    t = np.arange(t0, tend, dt)  # time grid\
    n = len(t)
    x = np.zeros((n, 2))  # x-grid
    e = np.zeros((n, 2))
    a = np.zeros(n)
    b = np.zeros(n)

    ki = k #np.linalg.inv(k)

    x[0] = x0
    e[0][0] = (f1.subs("t", t[0]) - teta) * ki[0][0] + (f1.subs("t", t[0]) + teta) * ki[1][0]
    e[0][1] = (f2.subs("t", t[0]) - teta) * ki[0][1] + (f2.subs("t", t[0]) + teta) * ki[1][1]
    print(e[0][0], ",", e[0][1])
    x[0] = e[0]

    for i in range(1, n):
        e[i][0] = (f1.subs("t", t[i]) - teta) * ki[0][0] + (f1.subs("t", t[i]) + teta) * ki[1][0]
        e[i][1] = (f2.subs("t", t[i]) - teta) * ki[0][1] + (f2.subs("t", t[i]) + teta) * ki[1][1]

        dist = np.linalg.norm(e[i]-x[i-1])

        if dist <= 1:
            x[i] = x[i - 1]
        else:
            Vx = e[i][0]-e[i-1][0]
            Vy = e[i][1]-e[i-1][1]


            a[i] = Vy/Vx
            b[i] = x[i-1][1] - a[i]*x[i-1][0]

            A = a[i]+1
            B = 0 - 2*e[i][0] + 2*a[i]*(b[i]-e[i][1])
            C = math.pow(e[i][0], 2) - math.pow((b[i]-e[i][1]), 2) - 1

            delta = B*B - 4*A*C
            sd = math.sqrt(delta)

            p1x = (0 - B - sd) / (2*A)
            p2x = (0 - B + sd) / (2*A)

            p1y = a[i]*p1x + b[i]
            p2y = a[i]*p2x + b[i]

            print(e[i][0], e[i][1], A, B, C, delta, "; p1 ", p1x, ",", p1y, " p2 ", p2x, ",", p2y)

            distp1 = math.sqrt(math.pow(x[i - 1][0] - p1x, 2) + math.pow(x[i - 1][1] - p1y, 2))
            distp2 = math.sqrt(math.pow(x[i - 1][0] - p2x, 2) + math.pow(x[i - 1][1] - p2y, 2))

            if distp1 < distp2:
                x[i][0] = p1x
                x[i][1] = p1y
            else:
                x[i][0] = p2x
                x[i][1] = p2y

    fig = plt.figure()
    ax = plt.axes(xlim=(-4, 4), ylim=(-4, 4))
    line1, = ax.plot([], [])
    line2, = ax.plot([], [], ':')
    line3, = ax.plot([], [], ':')
    line4, = ax.plot([], [])

    xdata = []
    ydata = []
    exdata = []
    eydata = []

    def init():  # only required for blitting to give a clean slate.
        line1.set_data([], [])
        line2.set_data([], [])
        line3.set_data([], [])
        line4.set_data([], [])
        xdata.clear()
        ydata.clear()
        exdata.clear()
        eydata.clear()
        return line1, line2, line3,

    def animate(i):
        xdata.append(x[i][0])
        ydata.append(x[i][1])
        exdata.append(e[i][0])
        eydata.append(e[i][1])
        cx, cy = drawCircle(e[i][0], e[i][1])
        lx, ly = drawLine(a[i], b[i])
        line1.set_data(cx, cy)
        line2.set_data(xdata, ydata)
        line3.set_data(exdata, eydata)
        line4.set_data(lx, ly)
        return line1, line2, line3, line4

    ani = animation.FuncAnimation(fig, animate, init_func=init, interval=200, blit=True, save_count=1, frames=100)
    ani.save("movie.gif")

    plt.show()

    return x, t


quadraticcase2(t0=0.0, tend=1.6, dt=0.1, k=np.array([[1, 0], [0, 1]]), fun1="1.5*sin(t)", fun2="1.5*cos(t)", teta=0, x0=np.array([0.5, 0.0]))