import numpy as np
import matplotlib.pyplot as plt

# -----------------------------
# Generate realistic fake SFs
# -----------------------------
np.random.seed(2)

n_ref = 2000

# Apo SFs
Fa_amp = np.random.uniform(50, 200, n_ref)
Fa_phi = np.random.uniform(-np.pi, np.pi, n_ref)
Fa = Fa_amp * np.exp(1j * Fa_phi)

# Difference signal (correlated phases!)
dF_amp = np.random.uniform(5, 25, n_ref)
dF_phi = Fa_phi + np.random.normal(0, 0.3, n_ref)
dF = dF_amp * np.exp(1j * dF_phi)

# Mixed SFs
Fab = Fa + dF

# -----------------------------
# Occupancies to test
# -----------------------------
b_values = [1.0, 0.5, 0.2, 0.1, 0.05]

# -----------------------------
# Compute resultants
# -----------------------------
R_Fa = np.sum(Fa)
R_Fab = np.sum(Fab)
R_dF = np.sum(Fab - Fa)

R_Fb = []
R_dF_scaled = []

for b in b_values:
    Fb = Fa + (Fab - Fa) / b
    R_Fb.append(np.sum(Fb))
    R_dF_scaled.append(R_dF / b)

# -----------------------------
# Plot resultants
# -----------------------------
plt.figure(figsize=(6, 6))

# Reference resultants
plt.arrow(0, 0, R_Fa.real, R_Fa.imag,
          head_width=200, length_includes_head=True,
          label="R(Fa)", color="black")

plt.arrow(0, 0, R_Fab.real, R_Fab.imag,
          head_width=200, length_includes_head=True,
          label="R(Fab)", color="blue")

# Extrapolated vs difference resultants
for b, Rb, Rd in zip(b_values, R_Fb, R_dF_scaled):
    plt.arrow(0, 0, Rb.real, Rb.imag,
              head_width=200, length_includes_head=True,
              alpha=0.7, label=f"R(Fb), b={b}")
    plt.arrow(0, 0, Rd.real, Rd.imag,
              head_width=200, length_includes_head=True,
              linestyle="dashed", alpha=0.7)

plt.axhline(0, color="grey", lw=0.5)
plt.axvline(0, color="grey", lw=0.5)
plt.gca().set_aspect("equal")

plt.xlabel("Real")
plt.ylabel("Imag")
plt.title("Resultant vectors: extrapolation vs difference")
plt.legend(fontsize=8)
plt.tight_layout()
plt.show()