%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint

from initial_velocities import velocities_m, velocities_S
from DE_solver import derivs, equationsolver

"""
Explanation: Additional Question 2 Implementation
My first additional question is basically my base question but with my disrupting galaxy having a mass of 3 times that of the main galaxy
End of explanation
"""


ic_add2 = np.zeros(484)

"""
Explanation: Defining empty initial condition array:
End of explanation
"""


max_time_add2 = 1.5
time_step_add2 = 120
S_add2 = 1e11
M_add2 = S_add2/3

S_y_add2 = 70
S_x_add2 = -.01*S_y_add2**2+25
vxS_add2 = velocities_S(M_add2,S_add2,S_x_add2,S_y_add2)[0]
vyS_add2 = velocities_S(M_add2,S_add2,S_x_add2,S_y_add2)[1]

"""
Explanation: Setting values for S,M, and t:
End of explanation
"""


ic_add2[0] = S_x_add2
ic_add2[1] = S_y_add2
ic_add2[2] = vxS_add2
ic_add2[3] = vyS_add2

"""
Explanation: Setting initial condition array values pertaining to S:
End of explanation
"""


f = open('star_positions.npz','r')
r = np.load('star_positions.npz')
x_y = r['arr_0']
f.close()

"""
Explanation: Loading the positions of my stars:
End of explanation
"""


for i in range(0,120):
    ic_add2[(i+1)*4] = x_y[0][i]
    ic_add2[((i+1)*4)+1] = x_y[1][i] 

for n in range(1,int(len(ic_add2)/4)):
    ic_add2[n*4+2] = velocities_m(M_add2,ic_add2[n*4],ic_add2[n*4+1])[0]
    ic_add2[n*4+3] = velocities_m(M_add2,ic_add2[n*4],ic_add2[n*4+1])[1]

"""
Explanation: Putting these values into my initial condition array, as well calling the initial velocity function on each position:
End of explanation
"""


sol_add2 = equationsolver(ic_add2,max_time_add2,time_step_add2,M_add2,S_add2)

np.savez('additional_2_data.npz',sol_add2,ic_add2)

"""
Explanation: Calling my differential equation solver, and saving the data to disk:
End of explanation
"""


