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


params = Params(
    R1 = 1e6,   # * ohm
    C1 = 1e-9,  # * farad
    A = 5,      # * volt
    f = 1000,   # * Hz
)
params

"""
Explanation: The following circuit diagram (from Wikipedia) shows a low-pass filter built with one resistor and one capacitor.  

A "filter" is a circuit takes a signal, $V_{in}$, as input and produces a signal, $V_{out}$, as output.  In this context, a "signal" is a voltage that changes over time.
A filter is "low-pass" if it allows low-frequency signals to pass from $V_{in}$ to $V_{out}$ unchanged, but it reduces the amplitude of high-frequency signals.
By applying the laws of circuit analysis, we can derive a differential equation that describes the behavior of this system.  By solving the differential equation, we can predict the effect of this circuit on any input signal.
Suppose we are given $V_{in}$ and $V_{out}$ at a particular instant in time.  By Ohm's law, which is a simple model of the behavior of resistors, the instantaneous current through the resistor is:
$ I_R = (V_{in} - V_{out}) / R $
where $R$ is resistance in ohms.
Assuming that no current flows through the output of the circuit, Kirchhoff's current law implies that the current through the capacitor is:
$ I_C = I_R $
According to a simple model of the behavior of capacitors, current through the capacitor causes a change in the voltage across the capacitor:
$ I_C = C \frac{d V_{out}}{dt} $
where $C$ is capacitance in farads (F).
Combining these equations yields a differential equation for $V_{out}$:
$ \frac{d }{dt} V_{out} = \frac{V_{in} - V_{out}}{R C} $
Follow the instructions below to simulate the low-pass filter for input signals like this:
$ V_{in}(t) = A \cos (2 \pi f t) $
where $A$ is the amplitude of the input signal, say 5 V, and $f$ is the frequency of the signal in Hz.
Params and System objects
Here's a Params object to contain the quantities we need.  I've chosen values for R1 and C1 that might be typical for a circuit that works with audio signal.
End of explanation
"""


from numpy import pi

def make_system(params):
    """Makes a System object for the given conditions.
    
    params: Params object
    
    returns: System object
    """
    f, R1, C1 = params.f, params.R1, params.C1
    
    init = State(V_out = 0)
    omega = 2 * pi * f
    tau = R1 * C1
    cutoff = 1 / R1 / C1 / 2 / pi
    t_end = 4 / f
    
    return System(params, 
                  init=init, 
                  t_end=t_end, num=401,
                  omega=omega, tau=tau, 
                  cutoff=cutoff)

"""
Explanation: Now we can pass the Params object make_system which computes some additional parameters and defines init.


omega is the frequency of the input signal in radians/second.


tau is the time constant for this circuit, which is the time it takes to get from an initial startup phase to 


cutoff is the cutoff frequency for this circuit (in Hz), which marks the transition from low frequency signals, which pass through the filter unchanged, to high frequency signals, which are attenuated.


t_end is chosen so we run the simulation for 4 cycles of the input signal.
End of explanation
"""


system = make_system(params)
system

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


# Solution goes here

"""
Explanation: The system variable num controls how many time steps we get from run_solve_ivp.  The default is 101; in this case we increase it to 401 because the methods we'll use to analyze the results require high resolution in time.
Exercise: Write a slope function that takes as an input a State object that contains V_out, and returns the derivative of V_out.
End of explanation
"""


slope_func(0, system.init, system)

"""
Explanation: Test the slope function with the initial conditions.
End of explanation
"""


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

results.tail()

"""
Explanation: And then run the simulation.  I suggest using t_eval=ts to make sure we have enough data points to plot and analyze the results.
End of explanation
"""


def plot_results(results):
    V_out = results.V_out.copy()
    t_end = results.index[-1]
    
    if t_end < 0.1:
        V_out.index *= 1000
        xlabel = 'Time (ms)'
    else:
        V_out = results.V_out
        xlabel = 'Time (s)'
        
    V_out.plot(label='_nolegend')
    decorate(xlabel=xlabel,
             ylabel='$V_{out}$ (volt)')
    
plot_results(results)

"""
Explanation: Here's a function you can use to plot V_out as a function of time.
End of explanation
"""


from matplotlib.pyplot import subplot

fs = [1, 10, 100, 1000, 10000, 100000]

for i, f in enumerate(fs):
    system = make_system(params.set(f=f))
    results, details = run_solve_ivp(system, slope_func)
    subplot(3, 2, i+1)
    plot_results(results)

"""
Explanation: If things have gone according to plan, the amplitude of the output signal should be about 0.8 V.
Also, you might notice that it takes a few cycles for the signal to get to the full amplitude. 
Sweeping frequency
Here's what V_out looks like for a range of frequencies:
End of explanation
"""


system = make_system(params.set(f=1000))
results, details = run_solve_ivp(system, slope_func)
V_out = results.V_out
plot_results(results)

"""
Explanation: At low frequencies, notice that there is an initial "transient" before the output gets to a steady-state sinusoidal output.  The duration of this transient is a small multiple of the time constant, tau, which is 1 ms.
Estimating the output ratio
Let's compare the amplitudes of the input and output signals.  Below the cutoff frequency, we expect them to be about the same.  Above the cutoff, we expect the amplitude of the output signal to be smaller.
We'll start with a signal at the cutoff frequency, f=1000 Hz.
End of explanation
"""


def compute_vin(results, system):
    """Computes V_in as a TimeSeries.
    
    results: TimeFrame with simulation results
    system: System object with A and omega
    
    returns: TimeSeries
    """
    A, omega = system.A, system.omega
    
    ts = results.index
    V_in = A * np.cos(omega * ts)
    return TimeSeries(V_in, results.index, name='V_in')

"""
Explanation: The following function computes V_in as a TimeSeries:
End of explanation
"""


V_in = compute_vin(results, system)

V_out.plot()
V_in.plot()

decorate(xlabel='Time (s)',
         ylabel='V (volt)')

"""
Explanation: Here's what the input and output look like.  Notice that the output is not just smaller; it is also "out of phase"; that is, the peaks of the output are shifted to the right, relative to the peaks of the input.
End of explanation
"""


def estimate_A(series):
    """Estimate amplitude.
    
    series: TimeSeries
    
    returns: amplitude in volts
    """
    return (series.max() - series.min()) / 2

"""
Explanation: The following function estimates the amplitude of a signal by computing half the distance between the min and max.
End of explanation
"""


A_in = estimate_A(V_in)
A_in

"""
Explanation: The amplitude of V_in should be near 5 (but not exact because we evaluated it at a finite number of points).
End of explanation
"""


A_out = estimate_A(V_out)
A_out

"""
Explanation: The amplitude of V_out should be lower.
End of explanation
"""


ratio = A_out / A_in
ratio

"""
Explanation: And here's the ratio between them.
End of explanation
"""


# Solution goes here

"""
Explanation: Exercise: Encapsulate the code we have so far in a function that takes two TimeSeries objects and returns the ratio between their amplitudes.
End of explanation
"""


estimate_ratio(V_out, V_in)

"""
Explanation: And test your function.
End of explanation
"""


corr = np.correlate(V_out, V_in, mode='same')
corr_series = make_series(V_in.index, corr)
corr_series.plot(color='C4')
decorate(xlabel='Lag (s)',
         ylabel='Correlation')

"""
Explanation: Estimating phase offset
The delay between the peak of the input and the peak of the output is call a "phase shift" or "phase offset", usually measured in fractions of a cycle, degrees, or radians.
To estimate the phase offset between two signals, we can use cross-correlation.  Here's what the cross-correlation looks like between V_out and V_in:
End of explanation
"""


peak_time = corr_series.idxmax()
peak_time

"""
Explanation: The location of the peak in the cross correlation is the estimated shift between the two signals, in seconds.
End of explanation
"""


period = 1 / system.f
period

peak_time / period

"""
Explanation: We can express the phase offset as a multiple of the period of the input signal:
End of explanation
"""


frac, whole = np.modf(peak_time / period)
frac

"""
Explanation: We don't care about whole period offsets, only the fractional part, which we can get using modf:
End of explanation
"""


frac * 360

"""
Explanation: Finally, we can convert from a fraction of a cycle to degrees:
End of explanation
"""


# Solution goes here

"""
Explanation: Exercise: Encapsulate this code in a function that takes two TimeSeries objects and a System object, and returns the phase offset in degrees.
Note: by convention, if the output is shifted to the right, the phase offset is negative.
End of explanation
"""


estimate_offset(V_out, V_in, system)

"""
Explanation: Test your function.
End of explanation
"""


# Solution goes here

"""
Explanation: Sweeping frequency again
Exercise: Write a function that takes as parameters an array of input frequencies and a Params object.
For each input frequency it should run a simulation and use the results to estimate the output ratio (dimensionless) and phase offset (in degrees).
It should return two SweepSeries objects, one for the ratios and one for the offsets.
End of explanation
"""


fs = 10 ** linspace(0, 4, 9)

ratios, offsets = sweep_frequency(fs, params)

"""
Explanation: Run your function with these frequencies.
End of explanation
"""


ratios.plot(color='C2', label='output ratio')
decorate(xlabel='Frequency (Hz)',
         ylabel='$V_{out} / V_{in}$')

"""
Explanation: We can plot output ratios like this:
End of explanation
"""


def plot_ratios(ratios, system):
    """Plot output ratios.
    """
    # axvline can't handle a Quantity with units
    cutoff = magnitude(system.cutoff)
    plt.axvline(cutoff, color='gray', alpha=0.4)
    
    ratios.plot(color='C2', label='output ratio')
    decorate(xlabel='Frequency (Hz)',
             ylabel='$V_{out} / V_{in}$',
             xscale='log', yscale='log')

plot_ratios(ratios, system)

"""
Explanation: But it is useful and conventional to plot ratios on a log-log scale.  The vertical gray line shows the cutoff frequency.
End of explanation
"""


def plot_offsets(offsets, system):
    """Plot phase offsets.
    """
    # axvline can't handle a Quantity with units
    cutoff = magnitude(system.cutoff)
    plt.axvline(cutoff, color='gray', alpha=0.4)
    
    offsets.plot(color='C9', label='phase offset')
    decorate(xlabel='Frequency (Hz)',
             ylabel='Phase offset (degree)',
             xscale='log')

plot_offsets(offsets, system)

"""
Explanation: This plot shows the cutoff behavior more clearly.  Below the cutoff, the output ratio is close to 1.  Above the cutoff, it drops off linearly, on a log scale, which indicates that output ratios for high frequencies are practically 0.
Here's the plot for phase offset, on a log-x scale:
End of explanation
"""


# Solution goes here

"""
Explanation: For low frequencies, the phase offset is near 0.  For high frequencies, it approaches 90 degrees.
Analysis
By analysis we can show that the output ratio for this signal is
$A = \frac{1}{\sqrt{1 + (R C \omega)^2}}$ 
where $\omega = 2 \pi f$, and the phase offset is
$ \phi = \arctan (- R C \omega)$ 
Exercise: Write functions that take an array of input frequencies and returns $A(f)$ and $\phi(f)$ as SweepSeries objects.  Plot these objects and compare them with the results from the previous section.
End of explanation
"""


A = output_ratios(fs, system)

# Solution goes here

"""
Explanation: Test your function:
End of explanation
"""


phi = phase_offsets(fs, system)

"""
Explanation: Test your function:
End of explanation
"""


A.plot(style=':', color='gray', label='analysis')
plot_ratios(ratios, system)

phi.plot(style=':', color='gray', label='analysis')
plot_offsets(offsets, system)

"""
Explanation: Plot the theoretical results along with the simulation results and see if they agree.
End of explanation
"""


