import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Parameters for the bead motion and detector
velocity = 0.11  # Velocity of the bead (m/s)
sensor_width = 2e-4  # Sensor width (2 mm)
t_pass = 2e-3  # Time for bead to pass the sensor (2 ms)
t = np.linspace(0, t_pass, 1000)  # Time array for the bead passing
bead_position = velocity * t  # Position of the bead as a function of time

# Capacitance baseline (C_real) - assume a value for illustration
epsilon_0 = 8.85e-12  # Vacuum permittivity (F/m)
r = 1e-3  # Radius of the silica bead (1 mm)
epsilon_r = 4  # Relative permittivity of silica
C_real = 4 * np.pi * epsilon_0 * epsilon_r * r  # Static capacitance in Farads

# Capacitance change as bead passes (Gaussian profile)
cap_center = sensor_width / 2  # Sensor center at half the width
cap_variation = np.exp(-((bead_position - cap_center) ** 2) / (2 * (sensor_width / 4) ** 2))
C_time_varying = C_real * (1 + 0.1 * cap_variation)  # Scaled around the base capacitance

# Input AC signal parameters
I = 1e-6  # 1 microampere constant current
signal_freq = 1e6  # Signal frequency (1 MHz)
omega = 2 * np.pi * signal_freq  # Angular frequency for 1 MHz
V_in = np.sin(omega * t)  # Input AC voltage signal (sinusoidal)

# Voltage response as the bead passes through the detector
V_time_varying = I * V_in / (1j * omega * C_time_varying)  # Voltage response depends on time-varying C

# Get the magnitude of the voltage response
V_magnitude = np.abs(V_time_varying)  # Magnitude combines real and imaginary parts

# Set up the figure and axis
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(6, 10))

# Bead position plot
ax1.set_xlim(0, sensor_width)
ax1.set_ylim(-0.1, 1.1)
bead_line, = ax1.plot([], [], 'bo', label="Bead")
ax1.set_title("Bead Passing Through the Detector")
ax1.set_xlabel("Position (m)")
ax1.set_ylabel("Bead")
ax1.legend()

# Capacitance plot
ax2.set_xlim(0, t_pass * 1e3)
ax2.set_ylim(min(C_time_varying), max(C_time_varying) * 1.1)
cap_line, = ax2.plot([], [], 'r-', label="Capacitance")
ax2.set_title("Capacitance Change Over Time")
ax2.set_xlabel("Time (ms)")
ax2.set_ylabel("Capacitance (F)")
ax2.legend()

# Voltage response plot (magnitude)
ax3.set_xlim(0, t_pass * 1e3)
ax3.set_ylim(0, np.max(V_magnitude) * 1.1)  # Plot the magnitude with scaling
voltage_line, = ax3.plot([], [], 'g-', label="Voltage Response (Magnitude)")
ax3.set_title("Voltage Response Over Time (Magnitude)")
ax3.set_xlabel("Time (ms)")
ax3.set_ylabel("Voltage (V)")
ax3.legend()

# Initialize function for the animation
def init():
    bead_line.set_data([], [])
    cap_line.set_data([], [])
    voltage_line.set_data([], [])
    return bead_line, cap_line, voltage_line

# Update function for the animation
def update(frame):
    # Update the bead position (needs to be a sequence, even for a single point)
    bead_line.set_data([bead_position[frame]], [0.5])

    # Update the capacitance plot
    cap_line.set_data(t[:frame] * 1e3, C_time_varying[:frame])

    # Update the voltage response plot (magnitude)
    voltage_line.set_data(t[:frame] * 1e3, V_magnitude[:frame])

    return bead_line, cap_line, voltage_line

# Create the animation
anim = FuncAnimation(fig, update, frames=len(t), init_func=init, blit=True, interval=20)

# Show the animation
plt.tight_layout()
plt.show()