# install Pint if necessary

try:
    import pint
except ImportError:
    !pip install pint

# download modsim.py if necessary

from os.path import exists

filename = 'modsim.py'
if not exists(filename):
    from urllib.request import urlretrieve
    url = 'https://raw.githubusercontent.com/AllenDowney/ModSim/main/'
    local, _ = urlretrieve(url+filename, filename)
    print('Downloaded ' + local)

# import functions from modsim

from modsim import *

"""
Explanation: Chapter 20
Modeling and Simulation in Python
Copyright 2021 Allen Downey
License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
End of explanation
"""




init = State(y=381, v=0)

"""
Explanation: So far the differential equations we've worked with have been first
order, which means they involve only first derivatives. In this
chapter, we turn our attention to second order ODEs, which can involve
both first and second derivatives.
We'll revisit the falling penny example from
Chapter xxx, and use run_solve_ivp to find the position and velocity of the penny as it falls, with and without air resistance.
Newton's second law of motion
First order ODEs can be written 
$$\frac{dy}{dx} = G(x, y)$$ 
where $G$ is some function of $x$ and $y$ (see http://modsimpy.com/ode). Second order ODEs can be written 
$$\frac{d^2y}{dx^2} = H(x, y, \frac{dy}{dt})$$
where $H$ is a function of $x$, $y$, and $dy/dx$.
In this chapter, we will work with one of the most famous and useful
second order ODEs, Newton's second law of motion: 
$$F = m a$$ 
where $F$ is a force or the total of a set of forces, $m$ is the mass of a moving object, and $a$ is its acceleration.
Newton's law might not look like a differential equation, until we
realize that acceleration, $a$, is the second derivative of position,
$y$, with respect to time, $t$. With the substitution
$$a = \frac{d^2y}{dt^2}$$ 
Newton's law can be written
$$\frac{d^2y}{dt^2} = F / m$$ 
And that's definitely a second order ODE.
In general, $F$ can be a function of time, position, and velocity.
Of course, this "law" is really a model in the sense that it is a
simplification of the real world. Although it is often approximately
true:


It only applies if $m$ is constant. If mass depends on time,
    position, or velocity, we have to use a more general form of
    Newton's law (see http://modsimpy.com/varmass).


It is not a good model for very small things, which are better
    described by another model, quantum mechanics.


And it is not a good model for things moving very fast, which are
    better described by yet another model, relativistic mechanics.


However, for medium-sized things with constant mass, moving at
medium-sized speeds, Newton's model is extremely useful. If we can
quantify the forces that act on such an object, we can predict how it
will move.
Dropping pennies
As a first example, let's get back to the penny falling from the Empire State Building, which we considered in
Chapter xxx. We will implement two models of this system: first without air resistance, then with.
Given that the Empire State Building is 381 m high, and assuming that
the penny is dropped from a standstill, the initial conditions are:
End of explanation
"""


g = 9.8

"""
Explanation: where y is height above the sidewalk and v is velocity. 
The units m and s are from the units object provided by Pint:
The only system parameter is the acceleration of gravity:
End of explanation
"""


t_end = 10
dt = 0.1

"""
Explanation: In addition, we'll specify the duration of the simulation and the step
size:
End of explanation
"""




system = System(init=init, g=g, t_end=t_end, dt=dt)

"""
Explanation: With these parameters, the number of time steps is 100, which is good
enough for many problems. Once we have a solution, we will increase the
number of steps and see what effect it has on the results.
We need a System object to store the parameters:
End of explanation
"""


def slope_func(t, state, system):
    y, v = state

    dydt = v
    dvdt = -system.g
    
    return dydt, dvdt

"""
Explanation: Now we need a slope function, and here's where things get tricky. As we have seen, run_solve_ivp can solve systems of first order ODEs, but Newton's law is a second order ODE. However, if we recognize that


Velocity, $v$, is the derivative of position, $dy/dt$, and


Acceleration, $a$, is the derivative of velocity, $dv/dt$,


we can rewrite Newton's law as a system of first order ODEs:
$$\frac{dy}{dt} = v$$ 
$$\frac{dv}{dt} = a$$ 
And we can translate those
equations into a slope function:
End of explanation
"""


dydt, dvdt = slope_func(0, system.init, system)
print(dydt)
print(dvdt)

"""
Explanation: The first parameter, state, contains the position and velocity of the
penny. The last parameter, system, contains the system parameter g,
which is the magnitude of acceleration due to gravity.
The second parameter, t, is time. It is not used in this slope
function because none of the factors of the model are time dependent. I include it anyway because this function will be called by run_solve_ivp, which always provides the same arguments,
whether they are needed or not.
The rest of the function is a straightforward translation of the
differential equations, with the substitution $a = -g$, which indicates that acceleration due to gravity is in the direction of decreasing $y$. slope_func returns a sequence containing the two derivatives.
Before calling run_solve_ivp, it is a good idea to test the slope
function with the initial conditions:
End of explanation
"""




results, details = run_solve_ivp(system, slope_func)
details

results.head()

"""
Explanation: The result is 0 m/s for velocity and 9.8 m/s$^2$ for acceleration. Now we call run_solve_ivp like this:
End of explanation
"""




results.y.plot()

decorate(xlabel='Time (s)',
         ylabel='Position (m)')

"""
Explanation: results is a TimeFrame with two columns: y contains the height of
the penny; v contains its velocity.
We can plot the results like this:
End of explanation
"""


t_end = results.index[-1]
results.y[t_end]

"""
Explanation: Since acceleration is constant, velocity increases linearly and position decreases quadratically; as a result, the height curve is a parabola.
The last value of results.y is negative, which means we ran the simulation too long.
End of explanation
"""




t_crossings = crossings(results.y, 0)
t_crossings

"""
Explanation: One way to solve this problem is to use the results to
estimate the time when the penny hits the sidewalk.
The ModSim library provides crossings, which takes a TimeSeries and a value, and returns a sequence of times when the series passes through the value. We can find the time when the height of the penny is 0 like this:
End of explanation
"""


def event_func(t, state, system):
    y, v = state
    return y

"""
Explanation: The result is an array with a single value, 8.818 s. Now, we could run
the simulation again with t_end = 8.818, but there's a better way.
Events
As an option, run_solve_ivp can take an event function, which
detects an "event", like the penny hitting the sidewalk, and ends the
simulation.
Event functions take the same parameters as slope functions, state,
t, and system. They should return a value that passes through 0
when the event occurs. Here's an event function that detects the penny
hitting the sidewalk:
End of explanation
"""


results, details = run_solve_ivp(system, slope_func,
                                 events=event_func)
details

"""
Explanation: The return value is the height of the penny, y, which passes through
0 when the penny hits the sidewalk.
We pass the event function to run_solve_ivp like this:
End of explanation
"""


t_end = results.index[-1]
t_end

y, v = results.iloc[-1]
print(y)
print(v)

"""
Explanation: Then we can get the flight time and final velocity like this:
End of explanation
"""


# Solution

r_0 = 150e9   # 150 million km in m
v_0 = 0
init = State(r=r_0,
             v=v_0)

# Solution

radius_earth = 6.37e6  # meters
radius_sun = 696e6     # meters
r_final = radius_sun + radius_earth
r_final

r_0 / r_final

t_end = 1e7  # seconds

system = System(init=init,
                G=6.674e-11,   # N m^2 / kg^2
                m1=1.989e30,   # kg
                m2=5.972e24,   # kg
                r_final=radius_sun + radius_earth,
                t_end=t_end)

# Solution

def universal_gravitation(state, system):
    """Computes gravitational force.
    
    state: State object with distance r
    system: System object with m1, m2, and G
    """
    r, v = state
    G, m1, m2 = system.G, system.m1, system.m2
    
    force = G * m1 * m2 / r**2
    return force

# Solution

universal_gravitation(init, system)

# Solution

def slope_func(t, state, system):
    """Compute derivatives of the state.
    
    state: position, velocity
    t: time
    system: System object containing `m2`
    
    returns: derivatives of y and v
    """
    y, v = state
    m2 = system.m2    

    force = universal_gravitation(state, system)
    dydt = v
    dvdt = -force / m2
    
    return dydt, dvdt

# Solution

slope_func(0, system.init, system)

# Solution

def event_func(t, state, system):
    r, v = state
    return r - system.r_final

# Solution

event_func(0, init, system)

# Solution

results, details = run_solve_ivp(system, slope_func, 
                                 events=event_func)
details

# Solution

t_event = results.index[-1]
t_event

# Solution



seconds = t_event * units.second
days = seconds.to(units.day)

# Solution

results.index /= 60 * 60 * 24

# Solution

results.r /= 1e9

# Solution

results.r.plot(label='r')

decorate(xlabel='Time (day)',
         ylabel='Distance from sun (million km)')

"""
Explanation: If there were no air resistance, the penny would hit the sidewalk (or someone's head) at more than 300 km/h.
So it's a good thing there is air resistance.
Summary
But air resistance...
Exercises
Exercise: Here's a question from the web site Ask an Astronomer:
"If the Earth suddenly stopped orbiting the Sun, I know eventually it would be pulled in by the Sun's gravity and hit it. How long would it take the Earth to hit the Sun? I imagine it would go slowly at first and then pick up speed."
Use run_solve_ivp to answer this question.
Here are some suggestions about how to proceed:


Look up the Law of Universal Gravitation and any constants you need.  I suggest you work entirely in SI units: meters, kilograms, and Newtons.


When the distance between the Earth and the Sun gets small, this system behaves badly, so you should use an event function to stop when the surface of Earth reaches the surface of the Sun.


Express your answer in days, and plot the results as millions of kilometers versus days.


If you read the reply by Dave Rothstein, you will see other ways to solve the problem, and a good discussion of the modeling decisions behind them.
You might also be interested to know that it's actually not that easy to get to the Sun.
End of explanation
"""


%psource crossings

"""
Explanation: Under the hood
solve_ivp
Here is the source code for crossings so you can see what's happening under the hood:
End of explanation
"""


