import gemmi
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable

# -----------------------------
# User input
# -----------------------------
#mtz_file = "../data/on_100k_refine_141.mtz"
#F_label   = "F-model"
#PHI_label = "PHIF-model"
mtz_file = "../data/apo_100k_refine_5.mtz"
F_label   = "F-obs-filtered"
PHI_label = "PHIF-model"
n_max = 10000
n_bins = 10

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

# Unit cell
cell = mtz.cell

# Structure factor data
F   = np.array(mtz.column_with_label(F_label))
PHI = np.array(mtz.column_with_label(PHI_label))

# Gemmi-native HKL
# HKL (must be integers!)
H = np.array(mtz.column_with_label("H"), dtype=int)
K = np.array(mtz.column_with_label("K"), dtype=int)
L = np.array(mtz.column_with_label("L"), dtype=int)

# Remove invalid reflections
mask = np.isfinite(F) & np.isfinite(PHI)

F   = F[mask]
PHI = PHI[mask]
H   = H[mask]
K   = K[mask]
L   = L[mask]

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

# -----------------------------
# Compute resolution (d-spacing)
# -----------------------------
# d = 1 / |h*| where |h*| comes from reciprocal cell
d = np.array([cell.calculate_d([h, k, l]) for h, k, l in zip(H, K, L)])

# -----------------------------
# Sort reflections: low → high resolution
# (large d → small d)
# -----------------------------
idx = np.argsort(-d)

C = C[idx][:n_max]
d = d[idx][:n_max]

# -----------------------------
# Assign resolution bins
# -----------------------------
bins = np.linspace(d.min(), d.max(), n_bins + 1)
colors_array = plt.cm.coolwarm(np.linspace(0, 1, n_bins))  # colormap

bin_indices = np.digitize(d, bins) - 1
bin_indices = np.clip(bin_indices, 0, n_bins - 1)   # <- fix
# -----------------------------
# Head-to-tail Argand walk
# -----------------------------
plt.figure(figsize=(7, 7))

x, y = 0.0, 0.0   # first vector starts at origin

for i, z in enumerate(C):
    dx, dy = z.real, z.imag

    # draw arrow FROM (x_start, y_start) TO (x_start+dx, y_start+dy)
    plt.arrow(
        x, y,
        dx, dy,
        color=colors_array[bin_indices[i]],
        length_includes_head=True,
        head_width=0.03 * abs(z),
        alpha=0.6
    )

    # update start point for next vector
    x += dx
    y += dy

# -----------------------------
# Final tip on top
# -----------------------------
plt.scatter(x, y, c="red", s=80, label="Final tip")

# -----------------------------
# Resultant vector (sum of all SFs)
# -----------------------------
sum_C = np.sum(C)
plt.arrow(
    0, 0,
    sum_C.real, sum_C.imag,
    length_includes_head=True,
    head_width=0.07 * abs(sum_C),
    color="green",
    alpha=0.9,
    label="Resultant ΣF"
)


# Start / end markers
plt.scatter(0, 0, c="black", s=40, label="Start")

plt.xlabel("Re(F)")
plt.ylabel("Im(F)")
plt.title("Head-to-tail Argand walk (sorted by resolution)")
plt.axis("equal")
plt.grid(True)
plt.legend()
plt.show()


# Assuming C is your filtered, complex structure factor array
sum_C = np.sum(C)             # sum all complex vectors

# Magnitude
magnitude = np.abs(sum_C)

# Phase in degrees (-180 to +180)
phase_deg = np.angle(sum_C, deg=True)

print(f"Resultant vector magnitude: {magnitude:.3f}")
print(f"Resultant vector phase: {phase_deg:.2f}°")