import matplotlib.pyplot as plt
import numpy as np
from numpy import random
import scipy.signal as sgl

yi = random.normal(size=2**20)
yr = np.roll(yi, -(2**10))

figy, axy = plt.subplots()
axy.plot(np.arange(2**10, 2**11), yi[2**10 : 2**11])
axy.plot(np.arange(2**10, 2**11), yr[2**10 : 2**11])

figf, axf = plt.subplots()
axf.plot(*sgl.welch(yi))
axf.plot(*sgl.welch(yr))

eta = lambda r: yi + r * yr
u = lambda r: -yi + r * yr


def puv(eta, u):
    f, phi_eta = sgl.welch(eta)
    phi_u = sgl.welch(u)[1]
    phi_eta_u = np.abs(sgl.csd(eta, u)[1].real)

    return f, np.sqrt(
        (phi_eta + phi_u - 2 * phi_eta_u) / (phi_eta + phi_u + 2 * phi_eta_u)
    )


figr, axr = plt.subplots()

for r in np.arange(0, 1.1, 0.1):
    axr.plot(*puv(eta(r), u(r)), c="k")
    Rn = puv(
        eta(r) + 0.4 * random.normal(size=2**20),
        u(r) + 0.4 * random.normal(size=2**20),
    )
    axr.plot(
        *Rn,
        c="#ff6600",
    )
    axr.annotate(
        f"{r=:.1f}", (Rn[0][0], Rn[1][0]), bbox={"boxstyle": "square", "facecolor": "w"}
    )
axr.grid()
axr.autoscale(True, "x", tight=True)
axr.set(ylim=(0, 1), ylabel="R", xlabel="f")
axr.legend(("No noise", "40% noise"), loc="lower left")

figr.savefig("out_r_test.pdf")
plt.show()