import gemmi
import numpy as np
import matplotlib.pyplot as plt

# -----------------------------
# User input
# -----------------------------
mtz_file = "../data/Fext_tests_refine_5.mtz"
F_label   = "F-model"
PHI_label = "PHIF-model"
n_max = 30000

# -----------------------------
# Load MTZ
# -----------------------------
mtz = gemmi.read_mtz_file(mtz_file)

F   = np.array(mtz.column_with_label(F_label))[:n_max]
PHI = np.array(mtz.column_with_label(PHI_label))[:n_max]

# Convert to complex structure factors
C = F * np.exp(1j * np.deg2rad(PHI))

# -----------------------------
# Head-to-tail vector plot
# -----------------------------
plt.figure(figsize=(7, 7))

x_start, y_start = 0.0, 0.0   # first vector starts at origin

for z in C:
    dx, dy = z.real, z.imag

    # draw arrow FROM (x_start, y_start) TO (x_start+dx, y_start+dy)
    plt.arrow(
        x_start, y_start,
        dx, dy,
        length_includes_head=True,
        head_width=0.03 * abs(z),
        alpha=0.6
    )

    # update start point for next vector
    x_start += dx
    y_start += dy



# Mark start and end explicitly
plt.scatter(0, 0, c="black", s=40, label="Start")
plt.scatter(x_start, y_start, c="red", s=60, label="End")

plt.xlabel("Re(F)")
plt.ylabel("Im(F)")
plt.title("Head-to-tail Argand plot of structure factors")
plt.axis("equal")
plt.grid(True)
plt.legend()
plt.show()