import numpy as np
import matplotlib.pyplot as plt

def speaker_frequency_response(fs, Qts, Vas, Re, Bl, Sd, Le):
    """
    Simulate the frequency response of a speaker based on its Thiele-Small parameters.

    Parameters:
        fs: Resonance frequency of the driver (Hz)
        Qts: Total quality factor of the driver (dimensionless)
        Vas: Equivalent compliance volume of the driver (liters)
        Re: DC resistance of the voice coil (ohms)
        Bl: Motor force factor (T·m)
        Sd: Effective cone area (m^2)
        Le: Voice coil inductance (H)

    Returns:
        f: Frequency array (Hz)
        SPL: Frequency response (dB)
    """
    # Constants
    rho0 = 1.18  # Air density (kg/m^3)
    c = 343      # Speed of sound (m/s)

    # Convert Vas to cubic meters
    Vas_m3 = Vas / 1000

    # Frequency range for the simulation
    f = np.logspace(1, 4, 1000)  # 10 Hz to 10 kHz

    # Angular frequency
    omega = 2 * np.pi * f

    # Mechanical compliance
    Cms = Vas_m3 / (rho0 * c**2 * Sd**2)

    # Mechanical resistance
    Rms = (Bl**2) / (Re * Qts * omega[0])

    # Moving mass (approximation)
    Mms = 1 / ((2 * np.pi * fs)**2 * Cms)

    # Impedance of the driver
    Zm = Rms + 1j * omega * Mms + 1 / (1j * omega * Cms)
    Ze = Re + 1j * omega * Le

    # Voltage-to-velocity transfer function
    Hv = Bl / (Zm + Ze)

    # Calculate sound pressure level (SPL)
    SPL = 20 * np.log10(np.abs(Hv) * Sd * omega**2 / (4 * np.pi * rho0 * c))
    return f, SPL

# Define the function to calculate the impedance response
def calculate_impedance_response(frequencies, Fs, Qts, Re, Bl, Vas, Sd):
    omega = 2 * np.pi * frequencies
    omega_s = 2 * np.pi * Fs

    # Mechanical and acoustical compliance
    Cms = Vas / (1.18 * 343**2 * Sd**2)  # Cms from Vas and Sd
    Mms = 1 / ((2 * np.pi * Fs)**2 * Cms)  # Moving mass
    Rms = Bl**2 / (Re * Cms)  # Mechanical resistance

    # Electrical impedance
    Z_mechanical = Rms + 1j * omega * Mms - 1j / (omega * Cms)
    Z_electrical = Re + Bl**2 / Z_mechanical

    return np.abs(Z_electrical)

def impedanceTS(Rvc,Lvc,Cmes,Lces,Res,f):

    omega=2*np.pi*f
    # Equivalent impedance with an electric LCR alanog model of a speaker
    XR=(Res*Lces/Cmes)/(Lces/Cmes+1j*Res*Lces*omega*(1-1/(omega**2*Lces*Cmes)))
    Xtot=Rvc+2*np.pi*f*Lvc+XR
    
    return f, np.abs(Xtot)

# Define Thiele-Small parameters
fs = 129            # Resonance frequency (Hz)
Qts = 1.43          # Total quality factor (dimensionless)
Vas = 10.5          # Equivalent compliance volume (liters)
Re = 3              # Voice coil DC resistance (ohms)
Bl = 3.72           # Motor force factor (T·m)
Sd = 213.8e-4       # Effective cone area (m^2)
Le = 0.21e-3        # Voice coil inductance (H)
Qms = 8.67          # Mechanical Quality Factor (unitless)
Qes = 1.71          # Electrical Quality factor (unitless)
Res =  16           # Lossless electrical resistance (ohms)
Lces = Re/(2*np.pi*fs*Qes)
Cmes = Qes/(2*np.pi*fs*Re)
print(Qms*Qes/(Qms+Qes))
print(Qts)
frequencies = np.linspace(1,10000, 5000)

# Compute the impedance response
f,Xtot = impedanceTS (Re,Le,Cmes,Lces,Res,frequencies)

# Plot the impedance response
plt.figure(figsize=(10, 6))
plt.semilogx(frequencies, Xtot)
plt.title("Speaker Impedance Response")
plt.xlabel("Frequency (Hz)")
plt.ylabel("Impedance (Ohms)")
plt.grid(which="both", linestyle="--", linewidth=0.5)
plt.show()

# Simulate frequency response
frequencies, spl = speaker_frequency_response(fs, Qts, Vas, Re, Bl, Sd, Le)

# Plot the frequency response
plt.figure(figsize=(10, 6))
plt.semilogx(frequencies, spl, label="Frequency Response")
plt.title("Speaker Frequency Response")
plt.xlabel("Frequency (Hz)")
plt.ylabel("SPL (dB)")
plt.grid(which="both", linestyle="--", linewidth=0.5)
plt.legend()
plt.show()