#!/usr/bin/env python3
"""
CG mesoscopic scattering I(q) from a single total RDF.
Anogiannakis eq 2:  I_coh(q)/N = f1^2 [ 1 + rho0 * INT_0^inf 4*pi*r^2 (g11(r)-1) sin(qr)/(qr) dr ]
                               = f1^2 * S(q)
All CG beads treated as one species (single adjustable form factor f1).

Workflow:
  1) Compute the TOTAL g(r) over the solute beads (all treated as one group):
       gmx rdf -f X-run.xtc -s X-run.tpr -ref <solute> -sel <solute> \
               -o rdf.xvg -bin 0.02 -rmax <~L/2>
     dry: <solute> = all beads.   wet: <solute> = non-water (exclude W) or SANS is washed out.
  2) python sq.py --rdf rdf.xvg --rho <bead_number_density_nm^-3> [options]
     (or give --nbeads and --box to have rho0 computed)

Output columns: q(nm^-1)  S(q)  I(q)=f^2*S(q)

Notes:
  - g(r) here is the FULL pair correlation (INCLUDING intramolecular bonded pairs) —
    those contribute the molecular/aggregate form factor. Do NOT exclude intramolecular
    pairs in gmx rdf (no -excl needed beyond default).
  - The integral is truncated at r=rmax (=last g(r) point). Lorch damping (default on)
    suppresses the truncation ripples; compare --no-lorch to see their size.
  - f is a single scale factor; fit it (and optionally a flat background) to experiment.
    S(q) itself (f=1) carries all the q-shape.
"""
import argparse
import numpy as np

# NumPy 2.0 renamed trapz -> trapezoid (trapz removed). Work on both.
trapezoid = getattr(np, "trapezoid", None) or np.trapz


def read_xvg(path):
    r, g = [], []
    with open(path) as fh:
        for ln in fh:
            if ln[0] in ('#', '@', ';') or not ln.strip():
                continue
            c = ln.split()
            r.append(float(c[0])); g.append(float(c[1]))
    return np.asarray(r), np.asarray(g)


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="CG scattering I(q)=f^2 S(q) from total g(r) (Anogiannakis eq 2)")
    p.add_argument("--rdf", required=True, help="gmx rdf .xvg: columns r(nm) g(r)")
    p.add_argument("--rho", type=float, default=None, help="bead number density rho0 (nm^-3)")
    p.add_argument("--nbeads", type=int, default=None, help="# beads in the RDF selection (with --box, computes rho0)")
    p.add_argument("--box", type=float, default=None, help="cubic box edge (nm) (with --nbeads, computes rho0)")
    p.add_argument("--f", type=float, default=1.0, help="single form factor f1 (scale); default 1")
    p.add_argument("--qmin", type=float, default=0.05, help="min q (nm^-1)")
    p.add_argument("--qmax", type=float, default=15.0, help="max q (nm^-1)")
    p.add_argument("--nq", type=int, default=400, help="# q points")
    p.add_argument("--no-lorch", action="store_true", help="disable Lorch truncation-damping window")
    p.add_argument("-o", "--out", default="sq.dat")
    a = p.parse_args()

    r, g = read_xvg(a.rdf)
    if len(r) < 3:
        raise SystemExit("RDF file has too few points")
    dr = np.mean(np.diff(r))

    if a.rho is not None:
        rho0 = a.rho
    elif a.nbeads and a.box:
        rho0 = a.nbeads / a.box**3
    else:
        raise SystemExit("Provide --rho, or both --nbeads and --box, for the bead number density.")

    rmax = r[-1]
    h = g - 1.0                       # total correlation h(r)=g(r)-1
    if not a.no_lorch:               # Lorch window: sinc(pi r / rmax)
        x = np.pi * r / rmax
        lorch = np.ones_like(r)
        nz = x != 0
        lorch[nz] = np.sin(x[nz]) / x[nz]
        h = h * lorch

    q = np.linspace(a.qmin, a.qmax, a.nq)
    # S(q) = 1 + (4 pi rho0 / q) INT r h(r) sin(qr) dr   (trapezoid over the r grid)
    integrand_r = r * h                                  # shape (Nr,)
    # vectorized: for each q, integrate r*h*sin(qr)
    qr = np.outer(q, r)                                  # (Nq, Nr)
    integ = trapezoid(integrand_r[None, :] * np.sin(qr), r, axis=1)
    S = 1.0 + (4.0 * np.pi * rho0 / q) * integ
    I = a.f**2 * S

    with open(a.out, "w") as fh:
        fh.write(f"# I(q)=f^2 S(q) from {a.rdf}  (Anogiannakis eq 2)\n")
        fh.write(f"# rho0={rho0:.5f} nm^-3  rmax={rmax:.3f} nm  dr={dr:.4f} nm  "
                 f"f={a.f}  lorch={'off' if a.no_lorch else 'on'}\n")
        fh.write("# q(nm^-1)      S(q)          I(q)\n")
        for qi, Si, Ii in zip(q, S, I):
            fh.write(f"{qi:12.5f}  {Si:12.6e}  {Ii:12.6e}\n")
    print(f"wrote {a.out}: {a.nq} q-points, q=[{a.qmin},{a.qmax}] nm^-1, rho0={rho0:.5f} nm^-3")
    print(f"  S(0+) extrapolated ~ {S[0]:.3f}   (compare --no-lorch to gauge truncation ripple)")
