import gemmi
import numpy as np
import matplotlib.pyplot as plt
# -----------------------------
# User inputs
# -----------------------------
apo_pdb   = "../data/hewl_apo.pdb"
occ100_pdb = "../data/hewl_1.0-I.pdb"
occb_pdb  = "../data/hewl_0.5-I.pdb"

b = 0.5
a = 1.0 - b
dmin = 2.0   # resolution cutoff (Å)

# -----------------------------
# Compute structure factors
# -----------------------------
def compute_sf(pdb_file, dmin):
    # Read structure
    st = gemmi.read_structure(pdb_file)
    st.setup_entities()
    st.assign_label_seq_id()
    st.remove_alternative_conformations()

    model = st[0]
    cell = st.cell

    # Space group
    sg = st.find_spacegroup()
    if sg is None:
        if st.spacegroup_hm:
            sg = gemmi.SpaceGroup(st.spacegroup_hm)
        else:
            raise RuntimeError("No space group found")

    # Generate HKLs
    hkl = gemmi.make_miller_array(
        cell=cell,
        spacegroup=sg,
        dmin=dmin,
        unique=True
    )

    # Structure factor calculator
    calc = gemmi.StructureFactorCalculatorX(cell)

    # Compute SFs ONE reflection at a time (this is required)
    F = np.array([
        calc.calculate_sf_from_model(model, [int(h), int(k), int(l)])
        for h, k, l in hkl
    ], dtype=np.complex128)

    return hkl, F


hkl, Fa = compute_sf(apo_pdb, dmin)
_,   F100 = compute_sf(occ100_pdb, dmin)
_,   Fb_model = compute_sf(occb_pdb, dmin)

# -----------------------------
# Sanity: same HKLs?
# -----------------------------
assert len(Fa) == len(F100) == len(Fb_model)

# -----------------------------
# Test forward relation:
# Fab = a Fa + b F100
# -----------------------------
Fab = a * Fa + b * F100
print(Fab)

# Compare to model b-occupancy SFs
delta_forward = Fab - Fb_model

print("FORWARD TEST: Fab = aFa + bF100")
print("Mean |ΔF|:", np.mean(np.abs(delta_forward)))
print("Mean phase Δ (deg):",
      np.mean(np.abs(np.angle(Fab / Fb_model, deg=True))))

# -----------------------------
# Test inverse relation:
# F100 = (Fab − Fa) / b
# -----------------------------
F100_recon = (Fb_model - Fa) / b + Fa

delta_inverse = F100_recon - F100

print("\nINVERSE TEST: F100 = (Fab − Fa)/b + Fa")
print("Mean |ΔF|:", np.mean(np.abs(delta_inverse)))
print("Mean phase Δ (deg):",
      np.mean(np.abs(np.angle(F100_recon / F100, deg=True))))

# -----------------------------
# Optional: inspect worst reflections
# -----------------------------
idx = np.argsort(np.abs(delta_inverse))[-5:]

print("\nWorst reflections (inverse):")
for i in idx:
    print(
        f"|F100|={abs(F100[i]):.2f}  "
        f"|ΔF|={abs(delta_inverse[i]):.2f}  "
        f"Δφ={np.angle(F100_recon[i]/F100[i],deg=True):.1f}°"
    )