# 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 21
Modeling and Simulation in Python
Copyright 2021 Allen Downey
License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
End of explanation
"""


from modsim import Params

params = Params(
    mass = 0.0025,      # kg
    diameter = 0.019,   # m
    rho = 1.2,     # kg/m**3
    g = 9.8,       # m/s**2
    v_init = 0,    # m / s
    v_term = 18,   # m / s
    height = 381,  # m
    t_end = 30,    # s
)

"""
Explanation: In the previous chapter we simulated a penny falling in a vacuum, that
is, without air resistance. But the computational framework we used is
very general; it is easy to add additional forces, including drag.
In this chapter, I present a model of drag force and add it to the
simulation.
Drag force
As an object moves through a fluid, like air, the object applies force
to the air and, in accordance with Newton's third law of motion, the air applies an equal and opposite force to the object (see
http://modsimpy.com/newton).
The direction of this drag force is opposite the direction of
travel, and its magnitude is given by the drag equation (see
http://modsimpy.com/drageq): 
$$F_d = \frac{1}{2}~\rho~v^2~C_d~A$$
where


$F_d$ is force due to drag, in newtons (N).


$\rho$ is the density of the fluid in kg/m^3^.


$v$ is the magnitude of velocity in m/s.


$A$ is the reference area of the object, in m^2^. In this
    context, the reference area is the projected frontal area, that is, the visible area of the object as seen from a point on its line of
    travel (and far away).


$C_d$ is the drag coefficient, a dimensionless quantity that
    depends on the shape of the object (including length but not frontal area), its surface properties, and how it interacts with the fluid.


For objects moving at moderate speeds through air, typical drag
coefficients are between 0.1 and 1.0, with blunt objects at the high end of the range and streamlined objects at the low end (see
http://modsimpy.com/dragco).
For simple geometric objects we can sometimes guess the drag coefficient with reasonable accuracy; for more complex objects we usually have to take measurements and estimate $C_d$ from data.
Of course, the drag equation is itself a model, based on the assumption that $C_d$ does not depend on the other terms in the equation: density, velocity, and area. For objects moving in air at moderate speeds (below 45 mph or 20 m/s), this model might be good enough, but we should remember to revisit this assumption.
For the falling penny, we can use measurements to estimate $C_d$. In
particular, we can measure terminal velocity, $v_{term}$, which is
the speed where drag force equals force due to gravity:
$$\frac{1}{2}~\rho~v_{term}^2~C_d~A = m g$$ 
where $m$ is the mass of the object and $g$ is acceleration due to gravity. Solving this equation for
$C_d$ yields: 
$$C_d = \frac{2~m g}{\rho~v_{term}^2~A}$$ 
According to Mythbusters, the terminal velocity of a penny is between 35 and 65 mph (see http://modsimpy.com/mythbust). Using the low end of their range, 40 mph or about 18 m/s, the estimated value of $C_d$ is 0.44, which is close to the drag coefficient of a smooth sphere.
Now we are ready to add air resistance to the model.
The Params Object
As the number of system parameters increases, and as we need to do more work to compute them, we will find it useful to define a Params object to contain the quantities we need to make a System object. Params objects are similar to System objects, and we initialize them the same way.
Here's the Params object for the falling penny:
End of explanation
"""


from numpy import pi
from modsim import State
from modsim import System

def make_system(params):
    init = State(y=params.height, v=params.v_init)

    area = pi * (params.diameter/2)**2

    C_d = (2 * params.mass * params.g / 
           (params.rho * area * params.v_term**2))

    return System(init=init,
                  area=area,
                  C_d=C_d,
                  mass=params.mass,
                  rho=params.rho,
                  g=params.g,
                  t_end=params.t_end
    )

"""
Explanation: The mass and diameter are from http://modsimpy.com/penny. The density
of air depends on temperature, barometric pressure (which depends on
altitude), humidity, and composition (http://modsimpy.com/density). I
chose a value that might be typical in Boston, Massachusetts at 20 °C.
Here's a version of make_system that takes the Params object and computes the inital state, init, the area, and the coefficient of drag.
Then it returns a System object with the quantities we'll need for the simulation.
End of explanation
"""


system = make_system(params)

"""
Explanation: And here's how we call it.
End of explanation
"""


system.C_d

"""
Explanation: Based on the mass and diameter of the penny, the density of air, and acceleration due to gravity, and the observed terminal velocity, we estimate that the coefficient of drag is about 0.44.
End of explanation
"""


params2 = params.set(v_term=20)

"""
Explanation: It might not be obvious why it is useful to create a Params object just to create a System object.
In fact, if we only run one simulation, it might not be useful.  But it helps when we want to change or sweep the parameters.
For example, suppose we learn that the terminal velocity of a penny is actually closer to 20 m/s.
We can make a Params object with the new value, and a corresponding System object, like this:
End of explanation
"""


system2 = make_system(params2)
system2.C_d

"""
Explanation: The result from set is a new Params object that is identical to the original except for the given value of v_term. 
If we pass params2 to make_system, we see that it computes a different value of C_d.
End of explanation
"""


def slope_func(t, state, system):
    y, v = state
    rho, C_d, area = system.rho, system.C_d, system.area
    mass, g = system.mass, system.g
    
    f_drag = rho * v**2 * C_d * area / 2
    a_drag = f_drag / mass
    
    dydt = v
    dvdt = -g + a_drag
    
    return dydt, dvdt

"""
Explanation: If the terminal velocity of the penny is 20 m/s, rather than 18 m/s, that implies that the coefficient of drag is 0.36, rather than 0.44.
And that makes sense, since lower drag implies faster terminal velocity.
Using Params objects to make System objects helps make sure that relationships like this are consistent.  And since we are always making new objects, rather than modifying existing objects, we are less likely to make a mistake.
Simulation
Now let's get to the simulation.  Here's a version of the slope function that includes drag:
End of explanation
"""


slope_func(0, system.init, system)

"""
Explanation: f_drag is force due to drag, based on the drag equation. a_drag is
acceleration due to drag, based on Newton's second law.
To compute total acceleration, we add accelerations due to gravity and
drag. g is negated because it is in the direction of decreasing y,
and a_drag is positive because it is in the direction of increasing
y. In the next chapter we will use Vector objects to keep track of
the direction of forces and add them up in a less error-prone way.
End of explanation
"""


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

event_func(0, system.init, system)

"""
Explanation: To stop the simulation when the penny hits the sidewalk, we'll use the
event function from Section xxx:
End of explanation
"""


from modsim import run_solve_ivp

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

"""
Explanation: Now we can run the simulation like this:
End of explanation
"""


results.tail()

"""
Explanation: Here are the last few time steps:
End of explanation
"""


t_sidewalk = results.index[-1]
t_sidewalk

"""
Explanation: The final height is close to 0, as expected.
Interestingly, the final velocity is not exactly terminal velocity, which suggests that there are some numerical errors.
We can get the flight time from results.
End of explanation
"""


from modsim import decorate

def plot_position(results):
    results.y.plot()
        
    decorate(xlabel='Time (s)',
         ylabel='Position (m)')
    
plot_position(results)

"""
Explanation: Here's the plot of position as a function of time.
End of explanation
"""


def plot_velocity(results):

    results.v.plot(color='C1', label='v')
        
    decorate(xlabel='Time (s)',
             ylabel='Velocity (m/s)')
    
plot_velocity(results)

"""
Explanation: And velocity as a function of time:
End of explanation
"""


# Solution

params = params.set(v_init=-30)
system2 = make_system(params)

# Solution

results2, details2 = run_solve_ivp(system2, slope_func, 
                                 events=event_func)
details2.message

t_sidewalk = results2.index[-1]
t_sidewalk

plot_position(results2)

# Solution

plot_velocity(results2)

"""
Explanation: From an initial velocity of 0, the penny accelerates downward until it reaches terminal velocity; after that, velocity is constant.
Summary
Exercises
Exercise: Run the simulation with a downward initial velocity that exceeds the penny's terminal velocity.
What do you expect to happen?  Plot velocity and position as a function of time, and see if they are consistent with your prediction.
Hint: Use params.set to make a new Params object with a different initial velocity.
End of explanation
"""


# Solution

params_quarter = params.set(
    mass = 0.0057,      # kg
    diameter = 0.024,   # m
    flight_time = 19.1,   # s
)

# Solution

system3 = make_system(params_quarter)

# Solution

# Run the simulation

results3, details3 = run_solve_ivp(system3, slope_func, 
                                   events=event_func)
details3.message

# Solution

# And get the flight time

t_sidewalk = results3.index[-1]
t_sidewalk

# Solution

# The flight time is a little long, 
# so we could increase `v_term` and try again.

# Or we could write an error function

def error_func(guess, params):
    """Final height as a function of C_d.
    
    guess: guess at v_term
    params: Params object
    
    returns: height in m
    """
    print(guess)
    params = params.set(v_term=guess)
    system = make_system(params)
    results, details = run_solve_ivp(system, slope_func, 
                                     events=event_func)
    t_sidewalk = results.index[-1]
    error = t_sidewalk - params.flight_time
    return error

# Solution

# We can test the error function like this
v_guess1 = 18
error_func(v_guess1, params_quarter)

# Solution

v_guess2 = 22
error_func(v_guess2, params_quarter)

# Solution

# Now we can use `root_scalar` to find the value of 
# `v_term` that yields the measured flight time.

from scipy.optimize import root_scalar

res = root_scalar(error_func, params_quarter,
                  bracket=[v_guess1, v_guess2])

# Solution

v_term = res.root
v_term

# Solution

# Plugging in the estimated value, 
# we can use `make_system` to compute `C_d`

system4 = make_system(params_quarter.set(v_term=res.root))
system4.C_d

"""
Explanation: Exercise: Suppose we drop a quarter from the Empire State Building and find that its flight time is 19.1 seconds.  Use this measurement to estimate terminal velocity and coefficient of drag.
You can get the relevant dimensions of a quarter from https://en.wikipedia.org/wiki/Quarter_(United_States_coin).


Create a Params object with new values of mass and diameter. We don't know v_term, so we'll start with the initial guess 18 m/s.


Use make_system to create a System object.  


Call run_solve_ivp to simulate the system.  How does the flight time of the simulation compare to the measurement?


Try a few different values of v_term and see if you can get the simulated flight time close to 19.1 seconds.


Optionally, write an error function and use root_scalar to improve your estimate.


Use your best estimate of v_term to compute C_d.


Note: I fabricated the observed flight time, so don't take the results of this exercise too seriously.
End of explanation
"""


