#!/usr/bin/env python3
"""
Fit a scattering curve. Two modes:

(a) --ref REF.dat   : fit  I_fit = A*S(q) + B  to a reference curve (experiment,
                      or another simulation for dry-vs-wet cross-checks). Linear
                      least squares in (A,B). A = f1^2 * scale.

(b) --model sphere|cylinder : fit an analytical form factor to the simulated I(q)
                      to EXTRACT SIZE (no experiment needed):
                          I(q) = A * P_model(q; R[,L]) + B
                      sphere   -> micelle radius R
                      cylinder -> rod radius R and length L

Input curves: whitespace files with columns  q  I(q)  [err]  (comments #,@,;).
Use sq.py output (q S I) by pointing at the I column (default col 2, 0-indexed=2).

Examples:
  # dry vs wet shape comparison (fit wet to dry-shaped reference):
  python fit_iq.py wet_sq.dat --ref dry_sq.dat --log
  # extract spherical-micelle radius from a simulated I(q):
  python fit_iq.py sq.dat --model sphere --log
  # extract rod radius+length:
  python fit_iq.py sq.dat --model cylinder --log --R0 2.0 --L0 15
"""
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_curve(path, icol=2):
    q, y = [], []
    with open(path) as fh:
        for ln in fh:
            if not ln.strip() or ln[0] in "#@;":
                continue
            c = ln.split()
            q.append(float(c[0])); y.append(float(c[icol]))
    return np.asarray(q), np.asarray(y)


def P_sphere(q, R):
    x = q * R
    # avoid 0/0 at x->0 (limit is 1)
    f = np.where(x < 1e-6, 1.0, 3.0 * (np.sin(x) - x * np.cos(x)) / x**3)
    return f**2


def P_cylinder(q, R, L, nalpha=90):
    from scipy.special import j1
    a = np.linspace(1e-6, np.pi / 2, nalpha)          # orientation angle
    ca, sa = np.cos(a), np.sin(a)
    P = np.empty_like(q)
    for k, qq in enumerate(q):
        x = qq * R * sa
        bessel = np.where(x < 1e-6, 1.0, 2.0 * j1(x) / x)
        z = qq * L * ca / 2.0
        sinc = np.where(np.abs(z) < 1e-6, 1.0, np.sin(z) / z)
        F2 = (bessel * sinc) ** 2
        P[k] = trapezoid(F2 * sa, a)                   # orientational average
    return P


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Fit CG scattering I(q): to a reference, or to a form-factor model")
    p.add_argument("data", help="simulated curve (q I [err]); e.g. sq.py output")
    p.add_argument("--icol", type=int, default=2, help="0-indexed I column in data (sq.py: 2)")
    p.add_argument("--ref", help="reference curve to fit A*S+B against (experiment or another sim)")
    p.add_argument("--refcol", type=int, default=2, help="0-indexed I column in --ref")
    p.add_argument("--model", choices=["sphere", "cylinder"], help="analytical form factor to fit for size")
    p.add_argument("--log", action="store_true", help="fit in log-space (recommended for SANS dynamic range)")
    p.add_argument("--qmin", type=float, default=None); p.add_argument("--qmax", type=float, default=None)
    p.add_argument("--R0", type=float, default=2.0, help="initial radius guess (nm)")
    p.add_argument("--L0", type=float, default=10.0, help="initial length guess (nm, cylinder)")
    p.add_argument("-o", "--out", default="fit_iq.out")
    a = p.parse_args()

    q, I = read_curve(a.data, a.icol)
    m = np.ones_like(q, bool)
    if a.qmin: m &= q >= a.qmin
    if a.qmax: m &= q <= a.qmax
    q, I = q[m], I[m]

    def resid_space(model, data):
        if a.log:
            eps = 1e-30
            return np.log(np.clip(model, eps, None)) - np.log(np.clip(data, eps, None))
        return model - data

    if a.ref:
        # I_fit = A*Sref + B, linear -> solve via lstsq (in chosen space it's still
        # cleanest to do linear lstsq on raw, then report; log would make it nonlinear)
        qr, Sref = read_curve(a.ref, a.refcol)
        Sr = np.interp(q, qr, Sref)
        M = np.vstack([Sr, np.ones_like(Sr)]).T
        (A, B), *_ = np.linalg.lstsq(M, I, rcond=None)
        fit = A * Sr + B
        chi2 = np.sum(resid_space(fit, I) ** 2)
        print(f"reference fit:  A={A:.5g}  B={B:.5g}  (A = f1^2*scale)   chi2={chi2:.4g}")
        header = f"# ref fit A={A:.5g} B={B:.5g}"
    elif a.model:
        from scipy.optimize import curve_fit
        if a.model == "sphere":
            def f(qq, A, B, R): return A * P_sphere(qq, R) + B
            p0 = [I.max(), min(I.min(), 0.0), a.R0]
            # R must stay positive: P(q) is even in R, so an unbounded fit can
            # return a meaningless negative radius.
            bounds = ([0.0, -np.inf, 1e-3], [np.inf, np.inf, np.inf])
            popt, _ = curve_fit(lambda qq, A, B, R: (np.log(np.clip(f(qq, A, B, R), 1e-30, None)) if a.log else f(qq, A, B, R)),
                                q, (np.log(np.clip(I, 1e-30, None)) if a.log else I), p0=p0, bounds=bounds, maxfev=20000)
            A, B, R = popt
            print(f"sphere fit:  R={R:.3f} nm  (A={A:.4g} B={B:.4g})")
            print(f"  => micelle radius ~{R:.2f} nm; core+shell, compare to rdp.py Rg")
            fit = f(q, *popt); header = f"# sphere R={R:.3f} nm"
        else:
            def f(qq, A, B, R, L): return A * P_cylinder(qq, R, L) + B
            p0 = [I.max(), min(I.min(), 0.0), a.R0, a.L0]
            # R, L must stay positive (P(q) is even in both).
            bounds = ([0.0, -np.inf, 1e-3, 1e-3], [np.inf, np.inf, np.inf, np.inf])
            popt, _ = curve_fit(lambda qq, A, B, R, L: (np.log(np.clip(f(qq, A, B, R, L), 1e-30, None)) if a.log else f(qq, A, B, R, L)),
                                q, (np.log(np.clip(I, 1e-30, None)) if a.log else I), p0=p0, bounds=bounds, maxfev=40000)
            A, B, R, L = popt
            print(f"cylinder fit:  R={R:.3f} nm  L={L:.3f} nm  (A={A:.4g} B={B:.4g})")
            fit = f(q, *popt); header = f"# cylinder R={R:.3f} L={L:.3f} nm"
    else:
        raise SystemExit("Give --ref REF.dat or --model sphere|cylinder")

    with open(a.out, "w") as fh:
        fh.write(header + f"  (log-fit={a.log})\n# q  I_data  I_fit\n")
        for qi, di, fi in zip(q, I, fit):
            fh.write(f"{qi:12.5f}  {di:12.6e}  {fi:12.6e}\n")
    print(f"wrote {a.out}")
