import numpy as np
import matplotlib.pyplot as plt

# Problem 2.1
# a)
def verlet(force, x0, p0, m, dt, N):
    assert(x0.shape == p0.shape)

    x = np.zeros((N, *x0.shape))
    p = np.zeros((N, *p0.shape))

    x[0] = x0
    p[0] = p0

    for i in range(1, N):
        p[i] = p[i-1] + 1/2 * force(x[i-1]) * dt
        x[i] = x[i-1] + p[i] / m * dt
        p[i] = p[i] + 1/2 * force(x[i]) * dt

    return x, p

r1_0 = np.array([1/3, 0, 0])
r2_0 = -r1_0
x0 = np.array([r1_0, r2_0]).flatten()

p1_0 = np.array([-6/7, 3/7, -2/7])
p2_0 = -p1_0
p0 = np.array([p1_0, p2_0]).flatten()

m = np.array([1, 1, 1, 1, 1, 1]) # ??
a = 1 # ??
epsilon = 1 # ??

k = 1 # ??
dt = 1e-2 # 1/omega_0
t_max = 10 # 1/omeag_0
t = np.arange(0, t_max, dt)

def F_H(x):
    r1 = x[:3]
    r2 = x[3:]

    return np.array([
        -k * (r1 - r2),
        k * (r1 - r2),
    ]).flatten()

x, p = verlet(F_H, x0, p0, m, dt, int(t_max / dt))
r1 = x[:,:3]
r2 = x[:,3:]
p1 = p[:,:3]
p2 = p[:,3:]

plt.plot(r1[:,0], r1[:,1], label="Particle 1", lw=4)
plt.plot(r2[:,0], r2[:,1], label="Particle 2")

plt.xlabel(r"$x / a$")
plt.ylabel(r"$y / a$")
plt.title("Trajectories")
plt.legend()
plt.savefig("problem2.1a-trajectories.png")
plt.show()

p_tot = p1 + p2
plt.plot(t, np.linalg.norm(p_tot - p_tot[0], axis=1))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$p / \sqrt{m \epsilon}$")
plt.title("Momentum Deviations")
plt.savefig("problem2.1a-dp.png")
plt.show()

J = np.cross(r1, p1, axis=1) + np.cross(r2, p2, axis=1)
plt.plot(t, np.linalg.norm(J - J[0], axis=1))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
plt.title("Angular Momentum Deviations")
plt.savefig("problem2.1a-dJ.png")
plt.show()

plt.plot(t, 1 - np.inner(J, J[0]) / (np.linalg.norm(J, axis=1) * np.linalg.norm(J[0])))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
plt.title("Momentum Deviations function")
plt.savefig("problem2.1a-J?.png")
plt.show()

E = 1/2 * np.linalg.norm(p1, axis=1)**2 + 1/2 * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2
plt.plot(t, E - E[0])
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$E / \epsilon$")
plt.title("Energy Deviations")
plt.savefig("problem2.1a-dE.png")
plt.show()

print(f"Maximum energy deviation: {np.max(np.abs((E - E[0]) / E[0]))}")

# b)
def FENE(x):
    r1 = x[:3]
    r2 = x[3:]

    f1 = -epsilon * (r1 - r2) / (a**2 - np.linalg.norm(r1 - r2)**2)
    f2 = -f1

    return np.array([
        f1,
        f2,
    ]).flatten()

x_F, p_F = verlet(FENE, x0, p0, m, dt, int(t_max / dt))
r1_F = x_F[:,:3]
r2_F = x_F[:,3:]
p1_F = p_F[:,:3]
p2_F = p_F[:,3:]

plt.plot(r1_F[:,0], r1_F[:,1], label="Particle 1")
plt.plot(r2_F[:,0], r2_F[:,1], label="Particle 2")
plt.xlabel(r"$x / a$")
plt.ylabel(r"$y / a$")
plt.title("Trajectories")
plt.legend()
plt.savefig("problem2.1b-trajectories.png")
plt.show()

p_totF = p1_F + p2_F
plt.plot(t, np.linalg.norm(p_totF - p_totF[0], axis=1))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$p / \sqrt{m \epsilon}$")
plt.title("Momentum Deviations")
plt.savefig("problem2.1b-dp.png")
plt.show()

J_F = np.cross(r1_F, p1_F, axis=1) + np.cross(r2_F, p2_F, axis=1)
plt.plot(t, np.linalg.norm(J_F - J_F[0], axis=1))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
plt.title("Angular Momentum Deviations")
plt.savefig("problem2.1b-dJ.png")
plt.show()

plt.plot(t, 1 - np.inner(J, J[0]) / (np.linalg.norm(J, axis=1) * np.linalg.norm(J[0])))
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
plt.title("Momentum Deviations function")
plt.savefig("problem2.1a-J?.png")
plt.show()

E_F = 1/2 * np.linalg.norm(p1_F, axis=1)**2 + 1/2 * np.linalg.norm(p2_F, axis=1)**2 + 1/2 * k * np.linalg.norm(r1_F - r2_F, axis=1)**2
plt.plot(t, E_F - E_F[0])
plt.xlabel(r"$t / \omega_0^{-1}$")
plt.ylabel(r"$E / \epsilon$")
plt.title("Energy Deviations")
plt.savefig("problem2.1b-dE.png")
plt.show()

print(f"Maximum energy deviation: {np.max(np.abs((E_F - E_F[0]) / E_F[0]))}")

# c)
lambda_c = 1/2 - (2 * np.sqrt(326) + 36)**(1/3) / 12 + 1 / (6 * (2 * np.sqrt(326) + 36)**(1/3))

def verlet_BABAB(force, x0, p0, m, dt, N, lambda_):
    assert(x0.shape == p0.shape)

    x = np.zeros((N, *x0.shape))
    p = np.zeros((N, *p0.shape))

    x[0] = x0
    p[0] = p0

    for i in range(1, N):
        p[i] = p[i-1] + lambda_ * force(x[i-1]) * dt
        x[i] = x[i-1] + 1/2 * p[i] / m * dt
        p[i] = p[i] + (1 - 2 * lambda_) * force(x[i]) * dt
        x[i] = x[i] + 1/2 * p[i] / m * dt
        p[i] = p[i] + lambda_ * force(x[i]) * dt

    return x, p

for f in (F_H, FENE,):
    for n in range(1, 5):
        for j in range(1, 6):
            dt = 4**(-j)
            t = np.arange(0, t_max, dt)

            x, p = verlet_BABAB(f, x0, p0, m, dt, len(t), lambda_c)
            r1 = x[:,:3]
            r2 = x[:,3:]
            p1 = p[:,:3]
            p2 = p[:,3:]

            E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2

            plt.plot(t, (E - E[0]) / dt**n, label="$k =" + str(j) + "$")

        plt.xlabel(r"$t / \omega_0^{-1}$")
        plt.ylabel(r"$\frac{E}{t^" + str(n) + r"} / \epsilon \omega_0^" + str(n) + r"$")
        plt.title("$n =" + str(n) + "$")
        plt.legend(loc="lower left")
        plt.savefig(f"problem2.1c-dE-{f.__name__}-n={n}.png")
        plt.show()

for lambda_ in (1/3, 0, 1/2):
    dt = 4**(-2)
    t = np.arange(0, t_max, dt)

    x, p = verlet_BABAB(f, x0, p0, m, dt, len(t), lambda_)
    r1 = x[:,:3]
    r2 = x[:,3:]
    p1 = p[:,:3]
    p2 = p[:,3:]

    E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2

    print(f"Maximum energy deviation for {lambda_}: {np.max(np.abs(E))}")

# d)
theta = 0.08398315262876693
lambda_d = 0.6822365335719091
rho = 0.2539785108410595
mu = -0.03230286765269967

def verlet_d(force, x0, p0, m, dt, N):
    assert(x0.shape == p0.shape)

    x = np.zeros((N, *x0.shape))
    p = np.zeros((N, *p0.shape))

    x[0] = x0
    p[0] = p0

    for i in range(1, N):
        p[i] = p[i-1] + theta * force(x[i-1]) * dt
        x[i] = x[i-1] + rho * p[i] / m * dt
        p[i] = p[i] + lambda_d * force(x[i]) * dt
        x[i] = x[i] + mu * p[i] / m * dt
        p[i] = p[i] + (1 - 2 * (lambda_d + theta)) / 2 * force(x[i]) * dt
        x[i] = x[i] + (1 - 2 * (mu + rho)) * p[i] / m * dt
        p[i] = p[i] + (1 - 2 * (lambda_d + theta)) / 2 * force(x[i]) * dt
        x[i] = x[i] + mu * p[i] / m * dt
        p[i] = p[i] + lambda_d * force(x[i]) * dt
        x[i] = x[i] + rho * p[i] / m * dt
        p[i] = p[i] + theta * force(x[i]) * dt

    return x, p

for f in (F_H, FENE):
    for n in range(1, 5):
        for j in range(1, 4):
            dt = 5 * 2**(-j)
            t = np.arange(0, t_max, dt)

            x, p = verlet_d(f, x0, p0, m, dt, len(t))
            r1 = x[:,:3]
            r2 = x[:,3:]
            p1 = p[:,:3]
            p2 = p[:,3:]

            E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2

            plt.plot(t, (E - E[0]) / dt**n, label="$k =" + str(j) + "$")

        plt.xlabel(r"$t / \omega_0^{-1}$")
        plt.ylabel(r"$\frac{E}{t^" + str(n) + r"} / \epsilon \omega_0^" + str(n) + r"$")
        plt.title("$n =" + str(n) + "$")
        plt.legend(loc="lower left")
        plt.savefig(f"problem2.1d-dE-{f.__name__}-n={n}.png")
        plt.show()