import itertools
import scgenerator as sc
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["font.family"] = "serif"


def main():
    t = np.linspace(-100, 100, 512) * 1e-15
    w = sc.wspace(t) + sc.units.nm_rads(1056)

    fig, (tops, bots) = plt.subplots(2, 4, constrained_layout=True, figsize=(11, 5))

    spec_kw = dict(cmap="Blues_r", marker=".")
    time_kw = dict(cmap="Greens_r", marker=".")

    for (const, phase), top, bot in zip(
        itertools.product((True, False), (True, False)), tops, bots
    ):
        sn = sc.pulse.shot_noise(w, const, phase)
        sn_t = np.fft.ifft(sn)
        top.scatter(sn.real, sn.imag, c=np.abs(w), **spec_kw)
        bot.scatter(sn_t.real, sn_t.imag, c=t[::-1], **time_kw)
        top.set_aspect(1, adjustable="datalim")
        bot.set_aspect(1, adjustable="datalim")
        top.axis("off")
        bot.axis("off")
        top.set_title(
            f"variance~$\\omega{'_0' if const else ''}$\n{'phase only' if phase else 'Re + Im'}"
        )

    sc.plotting.corner_annotation("↑Im\n→Re", tops[0])
    sc.plotting.corner_annotation(r"$\tilde{A}(\omega)$", tops[0], "bl")
    sc.plotting.corner_annotation("$A(t)$", bots[0], "bl")
    fig.suptitle("Different shot noise implementations, shading according to $|\\omega|$, $t$")
    plt.show()


if __name__ == "__main__":
    main()