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 = 0.3 * 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 = yi + yr
u = -yi + 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()

axr.plot(*puv(eta, u), label="Without noise")
axr.plot(
    *puv(
        eta + 0.4 * random.normal(size=2**20), u + 0.4 * random.normal(size=2**20)
    ),
    label="With noise"
)
axr.grid()
axr.autoscale(True, "x", tight=True)
axr.set(ylim=(0, 1), ylabel="R", xlabel="f")
axr.legend(loc="lower left")

axpsd = axr.twinx()
axpsd.plot(*sgl.welch(eta), label=r"PSD ($\eta$)", c="k", alpha=0.2, lw=1)
axpsd.legend(loc="lower right")
axpsd.set(ylabel="PSD")

plt.show()