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


# %%
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.zeros(elements)
        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,
        )


# %%
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"
        )


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