import numpy as np
import matplotlib.pyplot as plt

# Given parameters
k = 1.0   # Spring constant (kg/s^2)
m = 1.0   # Mass (kg)
v0 = 1.0  # Initial velocity (m/s)
t_max = 100.0  # Simulation time (s)
dt = 0.001  # Time step
N_steps = int(t_max / dt)  # Number of time steps
omega = np.sqrt(k / m)  # Angular frequency

# Initialize arrays
t_values = np.linspace(0, t_max, N_steps)
x_exact = (v0 / omega) * np.sin(omega * t_values)  # Analytical solution

# Arrays for numerical solutions
x_euler = np.zeros(N_steps)
v_euler = np.zeros(N_steps)
x_verlet = np.zeros(N_steps)
v_verlet = np.zeros(N_steps)

# Initial conditions
x_euler[0], v_euler[0] = 0.0, v0
x_verlet[0], v_verlet[0] = 0.0, v0
x_verlet[1] = x_verlet[0] + v_verlet[0] * dt - 0.5 * (k/m) * x_verlet[0] * dt**2  # One step for Verlet

# Euler Integration
for i in range(N_steps - 1):
    a = - (k / m) * x_euler[i]  # Acceleration
    x_euler[i + 1] = x_euler[i] + v_euler[i] * dt
    v_euler[i + 1] = v_euler[i] + a * dt

# Verlet Integration
for i in range(1, N_steps - 1):
    a = - (k / m) * x_verlet[i]  # Acceleration
    x_verlet[i + 1] = 2 * x_verlet[i] - x_verlet[i - 1] + a * dt**2  # Verlet update

# Compute total energy for each method
def compute_energy(x, v):
    K = 0.5 * m * v**2  # Kinetic Energy
    U = 0.5 * k * x**2  # Potential Energy
    return K + U  # Total Energy

E_euler = compute_energy(x_euler, v_euler)
E_verlet = compute_energy(x_verlet[:-1], (x_verlet[1:] - x_verlet[:-1]) / dt)
E_exact = compute_energy(x_exact, v0 * np.cos(omega * t_values))  # Using v(t) from exact solution

# Plot results
plt.figure(figsize=(12, 5))
plt.plot(t_values, x_exact, label="Analytical Solution", linestyle="dashed", color="black")
plt.plot(t_values, x_euler, label="Euler Method", linestyle="dotted", color="red")
plt.plot(t_values, x_verlet, label="Verlet Method", linestyle="solid", color="blue")
plt.xlabel("Time (s)")
plt.ylabel("Displacement x(t)")
plt.title("Simple Harmonic Motion: Euler vs. Verlet vs. Analytical")
plt.legend()
plt.grid()
plt.show()

# Energy Plot
plt.figure(figsize=(12, 5))
plt.plot(t_values, E_exact, label="Exact Energy", linestyle="dashed", color="black")
plt.plot(t_values, E_euler, label="Euler Energy", linestyle="dotted", color="red")
plt.plot(t_values[:-1], E_verlet, label="Verlet Energy", linestyle="solid", color="blue")
plt.xlabel("Time (s)")
plt.ylabel("Total Energy")
plt.title("Energy Conservation: Euler vs. Verlet vs. Exact")
plt.legend()
plt.grid()
plt.show()