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


params = Params(y_attach = 80,   # m,
                 v_init = 0,     # m / s,
                 g = 9.8,        # m/s**2,
                 M = 75,         # kg,
                 m_cord = 75,    # kg
                 area = 1,       # m**2,
                 rho = 1.2,      # kg/m**3,
                 v_term = 60,    # m / s,
                 L = 25,         # m,
                 k = 40,         # N / m
               )

"""
Explanation: In the previous case study, we simulated a bungee jump with a model that took into account gravity, air resistance, and the spring force of the bungee cord, but we ignored the weight of the cord.
It is tempting to say that the weight of the cord doesn't matter because it falls along with the jumper.  But that intuition is incorrect, as explained by Heck, Uylings, and Kędzierska.  As the cord falls, it transfers energy to the jumper.   They derive a differential equation that relates the acceleration of the jumper to position and velocity:
$a = g + \frac{\mu v^2/2}{\mu(L+y) + 2L}$ 
where $a$ is the net acceleration of the jumper, $g$ is acceleration due to gravity, $v$ is the velocity of the jumper, $y$ is the position of the jumper relative to the starting point (usually negative), $L$ is the length of the cord, and $\mu$ is the mass ratio of the cord and jumper.
If you don't believe this model is correct, this video might convince you.
Following the previous case study, we'll model the jump with the following assumptions:


Initially the bungee cord hangs from a crane with the attachment point 80 m above a cup of tea.


Until the cord is fully extended, it applies a force to the jumper as explained above.


After the cord is fully extended, it obeys Hooke's Law; that is, it applies a force to the jumper proportional to the extension of the cord beyond its resting length.


The jumper is subject to drag force proportional to the square of their velocity, in the opposite of their direction of motion.


First I'll create a Param object to contain the quantities we'll need:


Let's assume that the jumper's mass is 75 kg and the cord's mass is also 75 kg, so mu=1.


The jumpers's frontal area is 1 square meter, and terminal velocity is 60 m/s.  I'll use these values to back out the coefficient of drag.


The length of the bungee cord is L = 25 m.


The spring constant of the cord is k = 40 N / m when the cord is stretched, and 0 when it's compressed.


I adopt the coordinate system and most of the variable names from Heck, Uylings, and Kędzierska.
End of explanation
"""


def make_system(params):
    """Makes a System object for the given params.
    
    params: Params object
    
    returns: System object
    """
    M, m_cord = params.M, params.m_cord
    g, rho, area =  params.g, params.rho, params.area
    v_init, v_term = params.v_init, params.v_term
    
    # back out the coefficient of drag
    C_d = 2 * M * g / (rho * area * v_term**2)
    
    mu = m_cord / M
    init = State(y=params.y_attach, v=v_init)
    t_end = 8

    return System(params, C_d=C_d, mu=mu,
                  init=init, t_end=t_end)

"""
Explanation: Now here's a version of make_system that takes a Params object as a parameter.
make_system uses the given value of v_term to compute the drag coefficient C_d.
It also computes mu and the initial State object.
End of explanation
"""


system1 = make_system(params)

"""
Explanation: Let's make a System
End of explanation
"""


def drag_force(v, system):
    """Computes drag force in the opposite direction of `v`.
    
    v: velocity
    
    returns: drag force in N
    """
    rho, C_d, area = system.rho, system.C_d, system.area

    f_drag = -np.sign(v) * rho * v**2 * C_d * area / 2
    return f_drag

"""
Explanation: drag_force computes drag as a function of velocity:
End of explanation
"""


drag_force(20, system1)

"""
Explanation: Here's drag force at 20 m/s.
End of explanation
"""


def cord_acc(y, v, system):
    """Computes the force of the bungee cord on the jumper:
    
    y: height of the jumper
    v: velocity of the jumpter
    
    returns: acceleration in m/s
    """
    L, mu = system.L, system.mu
    
    a_cord = -v**2 / 2 / (2*L/mu + (L+y))
    return a_cord

"""
Explanation: The following function computes the acceleration of the jumper due to tension in the cord.
$a_{cord} = \frac{\mu v^2/2}{\mu(L+y) + 2L}$
End of explanation
"""


y = -20
v = -20
cord_acc(y, v, system1)

"""
Explanation: Here's acceleration due to tension in the cord if we're going 20 m/s after falling 20 m.
End of explanation
"""


def slope_func1(t, state, system):
    """Compute derivatives of the state.
    
    state: position, velocity
    t: time
    system: System object containing g, rho,
            C_d, area, and mass
    
    returns: derivatives of y and v
    """
    y, v = state
    M, g = system.M, system.g
    
    a_drag = drag_force(v, system) / M
    a_cord = cord_acc(y, v, system)
    dvdt = -g + a_cord + a_drag
    
    return v, dvdt

"""
Explanation: Now here's the slope function:
End of explanation
"""


slope_func1(0, system1.init, system1)

"""
Explanation: As always, let's test the slope function with the initial params.
End of explanation
"""


def event_func1(t, state, system):
    """Run until y=-L.
    
    state: position, velocity
    t: time
    system: System object containing g, rho,
            C_d, area, and mass
    
    returns: difference between y and y_attach-L
    """
    y, v = state   
    return y - (system.y_attach - system.L)

"""
Explanation: We'll need an event function to stop the simulation when we get to the end of the cord.
End of explanation
"""


event_func1(0, system1.init, system1)

"""
Explanation: We can test it with the initial conditions.
End of explanation
"""


results1, details1 = run_solve_ivp(system1, slope_func1, 
                                  events=event_func1)
details1.message

"""
Explanation: And then run the simulation.
End of explanation
"""


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

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


min(results1.y)

"""
Explanation: We can use min to find the lowest point:
End of explanation
"""


def spring_force(y, system):
    """Computes the force of the bungee cord on the jumper:
    
    y: height of the jumper
    
    Uses these variables from system:
    y_attach: height of the attachment point
    L: resting length of the cord
    k: spring constant of the cord
    
    returns: force in N
    """
    L, k = system.L, system.k
    
    distance_fallen = system.y_attach - y
    extension = distance_fallen - L
    f_spring = k * extension
    return f_spring

"""
Explanation: As expected, Phase 1 ends when the jumper reaches an altitude of 55 m.
Phase 2
Once the jumper has falled more than the length of the cord, acceleration due to energy transfer from the cord stops abruptly.  As the cord stretches, it starts to exert a spring force.  So let's simulate this second phase.
spring_force computes the force of the cord on the jumper:
End of explanation
"""


spring_force(55, system1)

spring_force(56, system1)

"""
Explanation: The spring force is 0 until the cord is fully extended.  When it is extended 1 m, the spring force is 40 N.
End of explanation
"""


def slope_func2(t, state, system):
    """Compute derivatives of the state.
    
    state: position, velocity
    t: time
    system: System object containing g, rho,
            C_d, area, and mass
    
    returns: derivatives of y and v
    """
    y, v = state
    M, g = system.M, system.g
    
    a_drag = drag_force(v, system) / M
    a_spring = spring_force(y, system) / M
    dvdt = -g + a_drag + a_spring
    
    return v, dvdt

"""
Explanation: The slope function for Phase 2 includes the spring force, and drops the acceleration due to the cord.
End of explanation
"""


t_final = results1.index[-1]
t_final

state_final = results1.iloc[-1]
state_final

"""
Explanation: The initial state for Phase 2 is the final state from Phase 1.
End of explanation
"""


system2 = System(system1, t_0=t_final, init=state_final)

"""
Explanation: And that gives me the starting conditions for Phase 2.
End of explanation
"""


results2, details2 = run_solve_ivp(system2, slope_func2)
details2.message

t_final = results2.index[-1]
t_final

"""
Explanation: Here's how we run Phase 2, setting the direction of the event function so it doesn't stop the simulation immediately.
End of explanation
"""


plot_position(results1, label='Phase 1')
plot_position(results2, label='Phase 2')

"""
Explanation: We can plot the results on the same axes.
End of explanation
"""


min(results2.y)

"""
Explanation: And get the lowest position from Phase 2.
End of explanation
"""


def run_two_phases(params):
    system1 = make_system(params)
    results1, details1 = run_solve_ivp(system1, slope_func1, 
                                       events=event_func1)
    t_final = results1.index[-1]
    state_final = results1.iloc[-1]
    
    system2 = system1.set(t_0=t_final, init=state_final)
    results2, details2 = run_solve_ivp(system2, slope_func2)
    return results1.append(results2)

"""
Explanation: To see how big the effect of the cord is, I'll collect the previous code in a function.
End of explanation
"""


results = run_two_phases(params)

plot_position(results)

params_no_cord = params.set(m_cord=1)
results_no_cord = run_two_phases(params_no_cord);

plot_position(results, label='m_cord = 75 kg')
plot_position(results_no_cord, label='m_cord = 1 kg')

min(results_no_cord.y)

diff = min(results.y) - min(results_no_cord.y)
diff

"""
Explanation: Now we can run both phases and get the results in a single TimeFrame.
End of explanation
"""


