# %% imports
import numpy as np
from utils import AnimatedPlot, db10, db20, dir_assets, wrap_phase


# %%
class PcolorPlot(AnimatedPlot):
    def __init__(self, elements: int, angle_deg: float = 0, spacing_lambda: float = 0.5, **kwargs):
        super().__init__(**kwargs)

        self.ax.axes.xaxis.set_visible(False)
        self.ax.axes.yaxis.set_visible(False)
        self.ax.set_aspect("equal")

        self.spacing_lambda = spacing_lambda
        x_range = spacing_lambda * (elements - 1)
        self.element_x = np.linspace(-x_range / 2, x_range / 2, elements)
        self.element_y = np.zeros(elements)
        self.angle = np.deg2rad(angle_deg)
        self.element_phase = (
            np.dot(np.array([self.element_x, self.element_y]).T, [np.sin(self.angle), np.cos(self.angle)]) * 2 * np.pi
        )

        x_max = np.max([x_range * 1.5, 10])
        x_pixels = 200
        self.x = np.linspace(-1, 1, x_pixels) * x_max
        self.y = np.linspace(-0.2, 1, x_pixels) * x_max

    def update(self, t: float):
        self.ax.clear()

        x = self.x
        y = self.y

        xx, yy = np.meshgrid(x, y)

        cdata = []
        for x, y, phase in zip(self.element_x, self.element_y, self.element_phase):
            distance = np.sqrt((xx - x) ** 2 + (yy - y) ** 2) + float(np.finfo(float).tiny)
            voltage = np.exp(1j * (phase + 2 * np.pi * distance - 2 * np.pi * t))  # * (1 / distance**2)
            cdata.append(voltage)

        cc = np.array(cdata).sum(axis=0)

        self.ax.scatter(self.element_x, self.element_y, marker="x", color="k")
        x_max = np.max(self.x)
        self.ax.arrow(
            x=np.sin(self.angle) * 0.05 * x_max,
            y=np.cos(self.angle) * 0.05 * x_max,
            dx=np.sin(self.angle) * 0.9 * x_max,
            dy=np.cos(self.angle) * 0.9 * x_max,
            head_width=0.05 * x_max,
            head_length=0.05 * x_max,
            width=0.01 * x_max,
            fc="k",
            ec="k",
        )

        return xx, yy, cc


class PcolorPhasePlot(PcolorPlot):
    def update(self, t: float):
        xx, yy, cc = super().update(t)

        self.ax.set_title(
            "Nearfield Phase\n"
            f"{len(self.element_x)} Element{'s' if len(self.element_x) > 1 else ''}, "
            f"{self.spacing_lambda}$\\lambda$ Spacing, "
            f"{np.rad2deg(self.angle):.0f}° Steer"
        )
        self.ax.pcolormesh(
            xx,
            yy,
            np.angle(cc, deg=False),
            clim=(-np.pi, np.pi),
            cmap="twilight",
            zorder=0,
        )


# TODO: don't animate this, it is constant
class PcolorMagPlot(PcolorPlot):
    def update(self, t: float):
        xx, yy, cc = super().update(t)

        self.ax.set_title(
            # Note that this ignores 1/r**2 losses
            "Nearfield Power Density\n"
            f"{len(self.element_x)} Element{'s' if len(self.element_x) > 1 else ''}, "
            f"{self.spacing_lambda}$\\lambda$ Spacing, "
            f"{np.rad2deg(self.angle):.0f}° Steer"
        )
        self.ax.pcolormesh(
            xx,
            yy,
            np.abs(cc) ** 2,
            # 20 * np.log10(np.abs(cc / len(self.element_x))),
            clim=(0, len(self.element_x) ** 2),
            # clim=(-30, 0),
            cmap="viridis",
            zorder=0,
        )


class Steer(AnimatedPlot):
    def __init__(self, elements: int, spacing_lambda: float = 0.5, **kwargs):
        super().__init__(**kwargs)

        self.ax.remove()
        self.axs = [
            self.fig.add_subplot(3, 1, 1),
            self.fig.add_subplot(3, 1, 2),
            self.fig.add_subplot(3, 1, 3),
        ]

        self.spacing_lambda = spacing_lambda
        self.elements = elements

    def update(self, t: float):
        for ax in self.axs:
            ax.clear()

        # rewind, don't reset
        if t > 0.5:
            t = 1 - t
        t *= 2

        angle = np.deg2rad((t * 2 - 1) * 90)  # steering angle

        def get_phase(position):
            return wrap_phase(np.sin(angle) * position * 2 * np.pi, deg=False)

        ideal_position = self.spacing_lambda * np.linspace(-(self.elements - 1) / 2, (self.elements - 1) / 2, 1001)
        ideal_phase = get_phase(ideal_position)

        element_position = self.spacing_lambda * np.linspace(
            -(self.elements - 1) / 2, (self.elements - 1) / 2, self.elements
        )
        element_phase = get_phase(element_position)
        element_taper = np.ones(self.elements)
        element_excitation = element_taper * np.exp(1j * element_phase)

        theta = np.linspace(-90, 90, 361)
        fft_points = 128
        fft_period = 90 / self.spacing_lambda
        theta_fft = np.linspace(0, fft_period, fft_points, endpoint=False)
        ff_pattern = np.interp(
            theta,
            theta_fft,
            np.fft.fft(np.concat([element_excitation, np.zeros(fft_points - self.elements)]), norm="backward")
            / self.elements,
            period=fft_period,
        )

        self.fig.suptitle(f"{np.rad2deg(angle):+5.1f}° Steer")

        rolloff = np.cos(np.deg2rad(theta))
        self.axs[0].set_ylabel("Farfield Magnitude [dB]")
        self.axs[0].plot(theta, db20(rolloff), color="gray")
        self.axs[0].plot(theta, db20(ff_pattern * rolloff))
        self.axs[0].set_xlim(-90, 90)
        self.axs[0].set_ylim(-30, 5)
        self.axs[0].set_xlabel("Theta [°]")
        self.axs[0].grid(True)
        self.axs[0].axvline(np.rad2deg(angle))

        self.axs[1].set_ylabel("Excitation Phase")
        self.axs[1].stem(element_position, np.rad2deg(element_phase))
        self.axs[1].plot(ideal_position, np.rad2deg(ideal_phase), linestyle="--")
        self.axs[1].set_ylim(-200, 200)

        self.axs[2].set_ylabel("Excitation Magnitude")
        self.axs[2].stem(element_position, element_taper)
        # self.axs[2].set_xlabel("Element")


# %%
def generate():
    # for elements in [
    #     1,
    #     2,
    #     4,
    #     10,
    # ]:
    #     PcolorPhasePlot(elements=elements, angle_deg=45).save(dir_assets / "beamforming" / f"phase_xz_{elements}el.gif")
    #     PcolorMagPlot(elements=elements, angle_deg=45).save(
    #         dir_assets / "beamforming" / f"magnitude_xz_{elements}el.gif"
    #     )
    Steer(elements=16, spacing_lambda=0.5, frames=200).save(dir_assets / "beamforming" / "steering.gif", framerate=15)


# %%
if __name__ == "__main__":
    generate()