#
#
# 
import numpy as np
import matplotlib.pyplot as plt

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Input
# Gravitational constant in arbitrary units
G = 1.0

# Masses
M = 1.0
m = 1.0

# Time parameters
dt = 0.01  # Time step
num_steps = 1000  # Number of steps in the simulation

# Gravitational force vector
def gravitational_force(r):
    # Distance from origin
    distance = np.sqrt(r[0]**2 + r[1]**2)
    # Force magnitude F = GMm/r^2 
    force_magnitude = G*M*m / distance**2
    # Force is directed to origin, -rvec/|rvec|
    force_direction = -r / distance
    # Force vector
    return force_magnitude * force_direction


# Initialize position and velocity (e.g., planet orbiting the sun)
r_init = np.array([1.0, 0.0]) # [x,y]
v_init = np.array([0.0, 1.0]) # [vx,vy]

r = r_init.copy()
v = v_init.copy()

# Store positions for plotting
trajectory = [r.copy()]


# Lists to store input-output pairs
inputs = []
outputs = []

# Simulation loop
for step in range(num_steps):
    # Compute the gravitational force at the current position
    F = gravitational_force(r)

    # Store initial point and velocity 
    inputs.append(np.concatenate([r, v]))
    
    # Update velocity (v = v + F * dt)
    v += F * dt
    
    # Update position (r = r + v * dt)
    r += v * dt
    
    # Save the current position
    trajectory.append(r.copy())

    # Store next point [x,y]
    outputs.append( [r[0],r[1]])

     

# Convert to numpy arrays
inputs = np.array(inputs)
outputs = np.array(outputs)
  

# Convert trajectory to numpy array for easy plotting
trajectory = np.array(trajectory)

# Plot the trajectory
plt.plot(trajectory[:, 0], trajectory[:, 1])
plt.xlabel('x')
plt.ylabel('y')
plt.title('Particle Trajectory in Gravitational Field')
plt.grid(True)
plt.axis('equal')  # Keep aspect ratio
plt.draw()
plt.pause(1e-3)



# Build the model
model = Sequential()
model.add(Input(shape=(4,))) # Input is [x, y, vx, vy]
model.add(Dense(64, activation='relu'))
model.add(Dense(64, activation='relu'))
model.add(Dense(2))  # Output is [x_next, y_next]

# Compile the model
model.compile(optimizer='adam', loss='mse')


# Train the model
model.fit(inputs, outputs, epochs=50, batch_size=32)
# Test the model
r_test = r_init.copy()
v_test = v_init.copy()
trajectory_pred = []

print('Calculating predicted trajectory')
for step in range(500):
    input_state = np.concatenate([r_test, v_test])
    
    # Predict the next position
    # model.predict wants argument (None, 4), input_state is (4,)
    r_next_pred = model.predict(input_state[np.newaxis,:], verbose=0)

    # Update velocity 
    v_test += gravitational_force(r_test) * dt

    # Update position
    r_test = r_next_pred[0]
    
    # Save the predicted position
    trajectory_pred.append(r_test)
    

# Convert to numpy array for plotting
trajectory_pred = np.array(trajectory_pred)

print('Plotting true vs. predicted trajectory')

plt.figure()
plt.plot(trajectory[:, 0], trajectory[:, 1], label='True Trajectory')

# Plot the predicted trajectory
plt.plot(trajectory_pred[:, 0], trajectory_pred[:, 1], label='Predicted Trajectory')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.grid(True)
plt.axis('equal')
plt.show()