import numpy as np
import matplotlib.pyplot as plt

# -----------------------------
# Fake but realistic SFs
# -----------------------------
np.random.seed(1)

n_ref = 30

# Apo structure factors
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 (e.g. ligand / backbone change)
dF_amp = np.random.uniform(5, 40, n_ref)
dF_phi = Fa_phi + np.random.normal(0, 0.4, n_ref)
dF = dF_amp * np.exp(1j * dF_phi)

# Mixed data
Fab = Fa + dF

# -----------------------------
# Choose occupancies
# -----------------------------
b_values = [1.0, 0.5, 0.2, 0.1]

# -----------------------------
# Plot
# -----------------------------
fig, axes = plt.subplots(2, len(b_values), figsize=(14, 6), sharex=True, sharey=True)

for i, b in enumerate(b_values):
    Fb = Fa + (Fab - Fa) / b          # extrapolated SF
    dF_scaled = (Fab - Fa) / b        # pure difference term

    # ----- Extrapolated vectors -----
    ax = axes[0, i]
    ax.set_title(f"Fb extrapolated (b={b})")
    for z in Fb:
        ax.arrow(0, 0, z.real, z.imag,
                 head_width=2, length_includes_head=True,
                 alpha=0.6)
    ax.axhline(0, color='grey', lw=0.5)
    ax.axvline(0, color='grey', lw=0.5)
    ax.set_aspect('equal')

    # ----- Difference vectors -----
    ax = axes[1, i]
    ax.set_title("ΔF / b")
    for z in dF_scaled:
        ax.arrow(0, 0, z.real, z.imag,
                 head_width=2, length_includes_head=True,
                 alpha=0.6, color='crimson')
    ax.axhline(0, color='grey', lw=0.5)
    ax.axvline(0, color='grey', lw=0.5)
    ax.set_aspect('equal')

axes[0, 0].set_ylabel("Imag(F)")
axes[1, 0].set_ylabel("Imag(F)")
for ax in axes[1]:
    ax.set_xlabel("Real(F)")

plt.tight_layout()
plt.show()