import numpy as np
from matplotlib import pyplot as plt

from loudspeaker_tmatrix.utils import to_db


def plot_loudspeaker_response(
    response_array: np.ndarray,
    freq_array: np.ndarray,
    title: str,
    magnitude_in_db: bool,
    magnitude_units: str,
    shift_phase: bool,
):

    if magnitude_in_db:
        magnitude = to_db(response_array)
    else:
        magnitude = np.abs(response_array)

    if shift_phase:
        phase = np.angle(-response_array, deg=True)
    else:
        phase = np.angle(response_array, deg=True)

    fig, ax1 = plt.subplots(figsize=(12, 4))
    fig.suptitle(title)

    ax1.semilogx(freq_array, magnitude, label="Magnitude", color="C0")
    ax1.set_xlabel("Frequency [Hz]")
    ax1.set_ylabel(f"Magnitude [{magnitude_units}]", color="C0")
    ax1.tick_params(axis="y", labelcolor="C0")

    ax2 = ax1.twinx()
    ax2.semilogx(freq_array, phase, color="r", label="Phase")
    x_ticks = np.sort(np.array([16, 31, 63, 125, 250, 500, 1000, 2000, 4000]))
    ax2.set_xticks(ticks=x_ticks, labels=x_ticks.tolist(), rotation=45)
    ax2.set_ylabel("Phase [º]", color="r")
    ax2.tick_params(axis="y", labelcolor="r")
    y_label1 = [r"$-180º$", r"$-90º$", r"$0º$", r"$90º$", r"$180º$"]
    ax2.set_yticks(np.array([-180, -90, 0, 90, 180]), y_label1)

    ax1.set_xlim(10, 4000)
    ax2.set_ylim(-180, 180)
    ax1.grid(axis="x")
    ax2.grid(axis="y")
    return fig