import numpy as np
import seaborn as sns
from scipy.integrate import odeint
from solve import *

M = 1e1
S = 1e1
Rmin = 25

e = 7 #Eccentricity

r_array = np.array([.2,.3,.4,.5,.6])*Rmin
N_array = np.array([12,18,24,30,36])

steps = 1e3
t = np.linspace(0,.4,steps) #Timescale of 1 billion years

atol=1e-6
rtol=1e-6
gamma = 4.49933e4 #Units ((kpc)^3)/((M_sun^10)(billion_years)^2)

Boolean = np.array([True,False])

"""
Explanation: This is a utility notebook/script that goes through and writes all of the possible combinations of solutions to npz files.
Hyperbolic/Parabolic
Retrograde/Direct
CM Frame/M Frame
Equal Mass Disruptor/Heavy Mass Disruptor
This notebook probably takes 2 minutes or so to run in entirety
End of explanation
"""


for bool1 in Boolean:
    CenterOfMass = bool1
    for bool2 in Boolean:
        Hyperbolic_Approach = bool2
        
        IC = set_IC(Hyperbolic_Approach,M,S)
        for bool3 in Boolean:
            Retrograde_Orbit = bool3
            
            allic = all_ic(N_array,r_array,IC,Retrograde_Orbit)
            
            sln = solution(t,allic,N_array,atol,rtol,gamma,Hyperbolic_Approach,M,S)
            x,dx,y,dy,X,dX,Y,dY = coordinate_solution(sln,steps)
            soln = np.array(frame(x,dx,y,dy,X,dX,Y,dY,CenterOfMass,M,S))
            
            if CenterOfMass == False:
                if Hyperbolic_Approach == True:
                    if Retrograde_Orbit == True:
                        np.savez('Hyp_Ret_M',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                    else:
                        np.savez('Hyp_Dir_M',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                else:
                    if Retrograde_Orbit == True:
                        np.savez('Par_Ret_M',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                    else:
                        np.savez('Par_Dir_M',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
            else:
                if Hyperbolic_Approach == True:
                    if Retrograde_Orbit == True:
                        np.savez('Hyp_Ret_Cen',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                    else:
                        np.savez('Hyp_Dir_Cen',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                else:
                    if Retrograde_Orbit == True:
                        np.savez('Par_Ret_Cen',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                    else:
                        np.savez('Par_Dir_Cen',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])

"""
Explanation: Equal Mass Disruptor Solutions
End of explanation
"""


M = 1e1
S = 3e1

for bool1 in Boolean:
    CenterOfMass = bool1
    for bool2 in Boolean:
        Hyperbolic_Approach = bool2
        
        IC = set_IC(Hyperbolic_Approach,M,S)
        for bool3 in Boolean:
            Retrograde_Orbit = bool3
            
            allic = all_ic(N_array,r_array,IC,Retrograde_Orbit)
            
            sln = solution(t,allic,N_array,atol,rtol,gamma,Hyperbolic_Approach,M,S)
            x,dx,y,dy,X,dX,Y,dY = coordinate_solution(sln,steps)
            soln = np.array(frame(x,dx,y,dy,X,dX,Y,dY,CenterOfMass,M,S))
            
            if CenterOfMass == False:
                if Hyperbolic_Approach == True:
                    if Retrograde_Orbit == True:
                        np.savez('Hyp_Ret_M_H',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                    else:
                        np.savez('Hyp_Dir_M_H',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                else:
                    if Retrograde_Orbit == True:
                        np.savez('Par_Ret_M_H',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
                    else:
                        np.savez('Par_Dir_M_H',x = soln[0], dx = soln[1],y = soln[2],dy = soln[3],X = soln[4],dX = soln[5],Y = soln[6],dY = soln[7])
            else:
                if Hyperbolic_Approach == True:
                    if Retrograde_Orbit == True:
                        np.savez('Hyp_Ret_Cen_H',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                    else:
                        np.savez('Hyp_Dir_Cen_H',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                else:
                    if Retrograde_Orbit == True:
                        np.savez('Par_Ret_Cen_H',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])
                    else:
                        np.savez('Par_Dir_Cen_H',x = soln[0], y = soln[1],X1 = soln[2],Y1 = soln[3],X2 = soln[4],Y2 = soln[5])

"""
Explanation: Writing Heavy Mass Disruptor Solutions
End of explanation
"""


