import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.patches import FancyArrowPatch
from scipy import ndimage
from scipy.optimize import curve_fit
from scipy.signal import windows
from skimage import io, exposure
from skimage.feature import peak_local_max
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')
import json
import os
import glob

# ── locate input image ────────────────────────────────────────────────────────
def find_image():
    exts = ['*.tif','*.tiff','*.png','*.jpg','*.jpeg','*.bmp']
    search_dirs = ['.']
    for root, dirs, files in os.walk('.'):
        depth = root.replace('.', '').count(os.sep)
        if depth <= 3:
            search_dirs.append(root)
    search_dirs = list(dict.fromkeys(search_dirs))
    for d in search_dirs:
        for pattern in exts:
            hits = glob.glob(os.path.join(d, pattern))
            if hits:
                return hits[0]
    return None

img_path = find_image()

if img_path is None:
    print('No input image found. Generating synthetic YBCO-like test image.')
    rng = np.random.default_rng(42)
    H_syn, W_syn = 512, 512
    syn = np.zeros((H_syn, W_syn), dtype=np.float64)
    ax_px, bx_px = 11, 0
    ay_px, by_px = 0, 11
    sublattices = [
        (0.0,  0.0,  1.0),
        (0.5,  0.5,  0.75),
        (0.0,  0.5,  0.45),
        (0.5,  0.0,  0.45),
    ]
    for ix in range(-5, W_syn // ax_px + 5):
        for iy in range(-5, H_syn // by_px + 5):
            for (fx, fy, amp) in sublattices:
                cx = int(round(ix * ax_px + fx * ax_px))
                cy = int(round(iy * by_px + fy * by_px))
                if 5 <= cx < W_syn-5 and 5 <= cy < H_syn-5:
                    for dr in range(-3, 4):
                        for dc in range(-3, 4):
                            val = amp * np.exp(-(dr**2 + dc**2) / (2 * 1.5**2))
                            syn[cy+dr, cx+dc] += val
    banding = np.sin(np.linspace(0, 8*np.pi, H_syn))[:, None] * 0.05
    syn = syn + banding
    syn = syn + rng.normal(0, 0.04, syn.shape)
    syn = np.clip(syn, 0, None)
    from PIL import Image as PILImage
    syn_uint16 = (syn / syn.max() * 65535).astype(np.uint16)
    PILImage.fromarray(syn_uint16).save('synthetic_ybco.tif')
    img_path = 'synthetic_ybco.tif'
    print(f'Synthetic image saved: {img_path}')
else:
    print(f'Loading: {img_path}')

# ── 1. load & normalise to float [0,1] ───────────────────────────────────────
raw = io.imread(img_path)
if raw.ndim == 3:
    raw = raw[..., 0]
raw = raw.astype(np.float64)

if raw.max() > raw.min():
    img_norm = (raw - raw.min()) / (raw.max() - raw.min())
else:
    img_norm = raw.copy()

H, W = img_norm.shape
PIXEL_SIZE_NM = 18.0 / 512.0
print(f'Image size: {W}x{H}  pixel size: {PIXEL_SIZE_NM:.5f} nm/px')

# ── 2. row-wise background normalisation ─────────────────────────────────────
def row_background_subtract(image, sigma=30.0):
    baseline = ndimage.gaussian_filter1d(image, sigma=sigma, axis=1)
    corrected = image - baseline
    return corrected

img_bg = row_background_subtract(img_norm, sigma=30.0)
img_bg = img_bg - img_bg.min()
if img_bg.max() > 0:
    img_bg /= img_bg.max()

# ── 3. 2-D FFT with Hamming window ───────────────────────────────────────────
def fft2_windowed(image):
    win_r = windows.hamming(image.shape[0])[:, None]
    win_c = windows.hamming(image.shape[1])[None, :]
    win2d = win_r * win_c
    ft = np.fft.fft2(image * win2d)
    ft_shift = np.fft.fftshift(ft)
    power = np.abs(ft_shift) ** 2
    return ft_shift, power

ft_shift, power = fft2_windowed(img_bg)
cy, cx = H // 2, W // 2

log_power = np.log1p(power)
lp_norm = (log_power - log_power.min()) / (log_power.max() - log_power.min() + 1e-12)

YY, XX = np.ogrid[:H, :W]
R_grid = np.sqrt((YY - cy)**2 + (XX - cx)**2)
mask_dc = R_grid < 5
lp_search = lp_norm.copy()
lp_search[mask_dc] = 0.0

upper_half = lp_search.copy()
upper_half[cy:, :] = 0.0
fft_peaks = peak_local_max(upper_half, min_distance=10, threshold_rel=0.30,
                           num_peaks=10)

print(f'FFT peaks found (upper half): {len(fft_peaks)}')

# ── 4. convert FFT peaks → real-space lattice parameters ─────────────────────
def fft_peak_to_realspace(peak_rc, image_shape, pixel_size_nm):
    py, px = peak_rc
    H_, W_ = image_shape
    fy = (py - H_//2) / H_
    fx = (px - W_//2) / W_
    q = np.sqrt(fx**2 + fy**2)
    if q < 1e-6:
        return None, None, None
    d_px = 1.0 / q
    d_nm = d_px * pixel_size_nm
    angle_deg = np.degrees(np.arctan2(fy, fx))
    return q, d_px, d_nm

fft_results = []
for pk in fft_peaks:
    q, d_px, d_nm = fft_peak_to_realspace(pk, (H, W), PIXEL_SIZE_NM)
    if d_nm is not None:
        fft_results.append({'row': int(pk[0]), 'col': int(pk[1]),
                             'q_inv_px': float(q),
                             'd_px': float(d_px),
                             'd_nm': float(d_nm)})
        print(f'  FFT peak ({pk[0]},{pk[1]})  d={d_nm:.3f} nm  ({d_px:.1f} px)')

if len(fft_results) >= 2:
    fft_sorted = sorted(fft_results, key=lambda x: -log_power[x['row'], x['col']])
    pk1, pk2 = fft_sorted[0], fft_sorted[1]
    lat_a_nm = pk1['d_nm']
    lat_b_nm = pk2['d_nm']
elif len(fft_results) == 1:
    lat_a_nm = fft_results[0]['d_nm']
    lat_b_nm = float('nan')
else:
    lat_a_nm = float('nan')
    lat_b_nm = float('nan')
print(f'Primary lattice parameters from FFT:  a={lat_a_nm:.3f} nm  b={lat_b_nm:.3f} nm')

# ── 5. bandpass filter ────────────────────────────────────────────────────────
def bandpass_filter(image, low_cutoff=1/30.0, high_cutoff=1/4.0):
    ft = np.fft.fftshift(np.fft.fft2(image))
    H_, W_ = image.shape
    cy_, cx_ = H_//2, W_//2
    YY_, XX_ = np.ogrid[:H_, :W_]
    R_px = np.sqrt((YY_ - cy_)**2 + (XX_ - cx_)**2)
    r_low  = max(H_, W_) * low_cutoff
    r_high = max(H_, W_) * high_cutoff
    bp_mask2 = (R_px >= r_low) & (R_px <= r_high)
    ft_bp = ft * bp_mask2.astype(float)
    img_bp = np.real(np.fft.ifft2(np.fft.ifftshift(ft_bp)))
    return img_bp, bp_mask2

img_bp, bp_mask = bandpass_filter(img_bg, low_cutoff=1/30.0, high_cutoff=1/4.0)
img_bp_norm = img_bp - img_bp.min()
if img_bp_norm.max() > 0:
    img_bp_norm /= img_bp_norm.max()

# ── 6. detect local maxima ────────────────────────────────────────────────────
coords = peak_local_max(img_bp_norm, min_distance=8, threshold_rel=0.10)
print(f'Initial peak detections: {len(coords)}')

# ── 7. 2-D Gaussian refinement (7x7 window) ──────────────────────────────────
def gaussian2d(xy, amp, x0, y0, sx, sy, offset):
    x, y = xy
    return (offset + amp *
            np.exp(-((x-x0)**2/(2*sx**2) + (y-y0)**2/(2*sy**2)))).ravel()

def refine_peak(image, r, c, half=3):
    r0, r1 = max(0, r-half), min(image.shape[0], r+half+1)
    c0, c1 = max(0, c-half), min(image.shape[1], c+half+1)
    patch = image[r0:r1, c0:c1].copy()
    if patch.size < 4:
        return r, c, image[r, c], half
    pr, pc = np.arange(r0, r1), np.arange(c0, c1)
    CC, RR = np.meshgrid(pc, pr)
    amp0   = patch.max() - patch.min()
    off0   = patch.min()
    sig0   = half / 2.0
    p0 = [amp0, float(c), float(r), sig0, sig0, off0]
    bounds_lo = [0,      c0, r0, 0.5, 0.5, -np.inf]
    bounds_hi = [np.inf, c1, r1, half*2, half*2, np.inf]
    try:
        popt, _ = curve_fit(gaussian2d, (CC, RR), patch.ravel(),
                            p0=p0, bounds=(bounds_lo, bounds_hi),
                            maxfev=400)
        amp, xf, yf, sx, sy, _ = popt
        return yf, xf, amp, 0.5*(abs(sx)+abs(sy))
    except Exception:
        return float(r), float(c), float(image[r, c]), float(half)

refined = []
for (r, c) in coords:
    rf, cf, amp, sig = refine_peak(img_bp_norm, r, c, half=3)
    refined.append({'r': rf, 'c': cf, 'amp': amp, 'sig': sig})

print(f'Refined peaks: {len(refined)}')

# ── 8. nearest-neighbour distances and angles ─────────────────────────────────
def compute_nn(positions, k=1):
    from sklearn.neighbors import KDTree
    tree = KDTree(positions)
    dists, inds = tree.query(positions, k=k+1)
    nn_dist  = dists[:, 1]
    nn_idx   = inds[:, 1]
    dr = positions[nn_idx, 0] - positions[:, 0]
    dc = positions[nn_idx, 1] - positions[:, 1]
    nn_angle = np.degrees(np.arctan2(dc, dr))
    return nn_dist, nn_angle, nn_idx

if len(refined) >= 3:
    pos = np.array([[p['r'], p['c']] for p in refined])
    nn_dist, nn_angle, nn_idx = compute_nn(pos, k=1)
    mean_nn_nm  = float(np.nanmean(nn_dist)) * PIXEL_SIZE_NM
    std_nn_nm   = float(np.nanstd(nn_dist))  * PIXEL_SIZE_NM
    print(f'NN distance: {mean_nn_nm:.3f} ± {std_nn_nm:.3f} nm')
else:
    pos = np.array([[p['r'], p['c']] for p in refined]) if refined else np.empty((0,2))
    nn_dist = np.array([])
    mean_nn_nm = std_nn_nm = float('nan')

# ── 9. fit 2-D lattice (affine model) ─────────────────────────────────────────
def fit_lattice_vectors(positions, d_approx_px, n_candidates=6):
    if len(positions) < 4:
        return None, None
    max_d = d_approx_px * 3.0
    displacements = []
    for i in range(len(positions)):
        for j in range(i+1, len(positions)):
            dr = positions[j,0] - positions[i,0]
            dc = positions[j,1] - positions[i,1]
            d  = np.sqrt(dr**2 + dc**2)
            if d < max_d:
                displacements.append([dc, dr])
                displacements.append([-dc, -dr])
    if not displacements:
        return None, None
    disp = np.array(displacements)
    bins = 80
    H_hist, xedge, yedge = np.histogram2d(disp[:,0], disp[:,1], bins=bins)
    cx_h = bins//2; cy_h = bins//2
    H_hist[cx_h-2:cx_h+3, cy_h-2:cy_h+3] = 0
    disp_peaks = peak_local_max(H_hist, min_distance=3, num_peaks=n_candidates)
    if len(disp_peaks) < 2:
        return None, None
    vecs = []
    for pk in disp_peaks[:n_candidates]:
        vx = 0.5*(xedge[pk[0]]+xedge[pk[0]+1])
        vy = 0.5*(yedge[pk[1]]+yedge[pk[1]+1])
        mag = np.sqrt(vx**2+vy**2)
        if mag > 2.0:
            vecs.append(np.array([vx, vy]))
    vecs_sorted = sorted(vecs, key=lambda v: np.linalg.norm(v))
    if len(vecs_sorted) < 2:
        return vecs_sorted[0] if vecs_sorted else None, None
    a_vec = vecs_sorted[0]
    best_b, best_orth = vecs_sorted[1], abs(np.dot(a_vec, vecs_sorted[1]))
    for v in vecs_sorted[2:]:
        orth = abs(np.dot(a_vec / (np.linalg.norm(a_vec)+1e-9),
                          v          / (np.linalg.norm(v)+1e-9)))
        if orth < best_orth:
            best_orth = orth
            best_b    = v
    return a_vec, best_b

if len(refined) >= 4:
    d_approx_px = float(np.nanmean(nn_dist)) if nn_dist.size > 0 else 10.0
    a_vec, b_vec = fit_lattice_vectors(pos, d_approx_px)
else:
    a_vec = b_vec = None

if a_vec is not None:
    a_len_nm = np.linalg.norm(a_vec) * PIXEL_SIZE_NM
    b_len_nm = np.linalg.norm(b_vec) * PIXEL_SIZE_NM if b_vec is not None else float('nan')
    cos_ab   = (np.dot(a_vec, b_vec) /
                (np.linalg.norm(a_vec)*np.linalg.norm(b_vec)+1e-12)) if b_vec is not None else float('nan')
    gamma_deg = float(np.degrees(np.arccos(np.clip(cos_ab, -1, 1)))) if b_vec is not None else float('nan')
    print(f'Fitted lattice:  |a|={a_len_nm:.3f} nm  |b|={b_len_nm:.3f} nm  γ={gamma_deg:.1f}°')
else:
    a_len_nm = b_len_nm = gamma_deg = float('nan')
    print('Lattice vector fitting failed (too few columns).')

# ── 10. sublattice identification ─────────────────────────────────────────────
def compute_fractional_coords(positions, a_vec, b_vec, origin=None):
    if a_vec is None or b_vec is None:
        return None
    if origin is None:
        origin = positions.mean(axis=0)
    M = np.array([[a_vec[0], b_vec[0]],
                  [a_vec[1], b_vec[1]]])
    try:
        M_inv = np.linalg.inv(M)
    except np.linalg.LinAlgError:
        return None
    rel = np.array([[p[1]-origin[1], p[0]-origin[0]] for p in positions])
    frac = (M_inv @ rel.T).T
    frac_mod = frac % 1.0
    return frac_mod

sublattice_labels = None
n_sub = 0
if len(refined) >= 6 and a_vec is not None:
    frac = compute_fractional_coords(pos, a_vec, b_vec)
    if frac is not None:
        amps = np.array([p['amp'] for p in refined])
        amps_norm = (amps - amps.min()) / (amps.max() - amps.min() + 1e-12)
        feat = np.column_stack([amps_norm, frac[:, 0], frac[:, 1]])
        scaler = StandardScaler()
        feat_s = scaler.fit_transform(feat)
        n_sub = min(4, len(refined)//4)
        n_sub = max(2, n_sub)
        km = KMeans(n_clusters=n_sub, n_init=10, random_state=42)
        sublattice_labels = km.fit_predict(feat_s)
        print(f'Sublattice clustering: {n_sub} clusters')
        for k in range(n_sub):
            idx_k = sublattice_labels == k
            print(f'  Sublattice {k}: n={idx_k.sum()}  '
                  f'mean_amp={amps[idx_k].mean():.3f}  '
                  f'frac_x={frac[idx_k,0].mean():.3f}  '
                  f'frac_y={frac[idx_k,1].mean():.3f}')

# ── 11. YBCO stoichiometry check ──────────────────────────────────────────────
ybco_ratios_note = 'Not enough sublattices detected for YBCO ratio check.'
if sublattice_labels is not None:
    amps     = np.array([p['amp'] for p in refined])
    sub_counts = {}
    sub_amp    = {}
    for k in range(n_sub):
        idx_k = sublattice_labels == k
        sub_counts[k] = int(idx_k.sum())
        sub_amp[k]    = float(amps[idx_k].mean())
    ranked = sorted(sub_amp.items(), key=lambda x: -x[1])
    total = sum(sub_counts.values())
    ratios = {k: sub_counts[k]/total for k in sub_counts}
    ybco_ratios_note = (
        f'Sublattice count ratios: {ratios}. '
        'In YBCO (HAADF) expect Ba-rich sublattice brightest, '
        'then Y, then Cu (≈1:2:3 ratio Ba:Y:Cu per formula unit). '
        f'Ranked by amplitude: {ranked}.'
    )
    print(ybco_ratios_note)

# ── 12. spatial statistics + displacement field ───────────────────────────────
n_columns   = len(refined)
amps_all    = np.array([p['amp'] for p in refined]) if refined else np.array([])
mean_amp    = float(np.nanmean(amps_all)) if amps_all.size > 0 else float('nan')
std_amp     = float(np.nanstd(amps_all))  if amps_all.size > 0 else float('nan')

if nn_dist.size > 0:
    cv_nn = float(np.nanstd(nn_dist) / (np.nanmean(nn_dist) + 1e-12))
else:
    cv_nn = float('nan')

displacements_nm = []
if a_vec is not None and b_vec is not None and len(refined) >= 4:
    origin = pos.mean(axis=0)
    for p in pos:
        rel_c = p[1] - origin[1]
        rel_r = p[0] - origin[0]
        M = np.array([[a_vec[0], b_vec[0]],
                      [a_vec[1], b_vec[1]]])
        try:
            M_inv = np.linalg.inv(M)
            frac_p = M_inv @ np.array([rel_c, rel_r])
            ideal_frac = np.round(frac_p)
            ideal_cart = M @ ideal_frac
            disp_px = np.sqrt((rel_c - ideal_cart[0])**2 +
                               (rel_r - ideal_cart[1])**2)
            displacements_nm.append(disp_px * PIXEL_SIZE_NM)
        except Exception:
            pass

mean_disp_nm = float(np.nanmean(displacements_nm)) if displacements_nm else float('nan')
std_disp_nm  = float(np.nanstd(displacements_nm))  if displacements_nm else float('nan')

print(f'Column count: {n_columns}')
print(f'Mean amplitude: {mean_amp:.4f}  Std: {std_amp:.4f}')
print(f'NN distance CV: {cv_nn:.4f}')
print(f'Mean displacement from ideal: {mean_disp_nm:.4f} nm  Std: {std_disp_nm:.4f} nm')

# ── figures ────────────────────────────────────────────────────────────────────
fig = plt.figure(figsize=(20, 18))
gs  = gridspec.GridSpec(3, 3, figure=fig, hspace=0.40, wspace=0.35)

ax1 = fig.add_subplot(gs[0, 0])
ax1.imshow(img_norm, cmap='gray', origin='upper')
ax1.set_title('1. Normalised Image', fontsize=10)
ax1.axis('off')

ax2 = fig.add_subplot(gs[0, 1])
ax2.imshow(img_bg, cmap='gray', origin='upper')
ax2.set_title('2. BG-Subtracted (row-wise)', fontsize=10)
ax2.axis('off')

ax3 = fig.add_subplot(gs[0, 2])
ax3.imshow(np.log1p(power), cmap='inferno', origin='upper')
for fr in fft_results:
    ax3.plot(fr['col'], fr['row'], 'c+', markersize=8, markeredgewidth=1.5)
ax3.set_title('3. FFT Power (log) + peaks', fontsize=10)
ax3.axis('off')

ax4 = fig.add_subplot(gs[1, 0])
ax4.imshow(img_bp_norm, cmap='gray', origin='upper')
ax4.set_title('5. Bandpass Filtered', fontsize=10)
ax4.axis('off')

ax5 = fig.add_subplot(gs[1, 1])
ax5.imshow(img_bp_norm, cmap='gray', origin='upper')
if len(refined) > 0:
    rr = [p['r'] for p in refined]
    cc = [p['c'] for p in refined]
    if sublattice_labels is not None:
        try:
            cmap_sub = matplotlib.colormaps['Set1'].resampled(n_sub)
        except Exception:
            cmap_sub = plt.cm.get_cmap('Set1', n_sub)
        for k in range(n_sub):
            idx_k = np.where(sublattice_labels == k)[0]
            ax5.scatter([cc[i] for i in idx_k], [rr[i] for i in idx_k],
                        s=18, color=cmap_sub(k), linewidths=0.5,
                        edgecolors='white', label=f'SL{k}', zorder=3)
        ax5.legend(fontsize=6, loc='upper right')
    else:
        ax5.scatter(cc, rr, s=12, c='red', linewidths=0, zorder=3)
ax5.set_title(f'6-7. Detected & Refined Columns (n={n_columns})', fontsize=10)
ax5.axis('off')

ax6 = fig.add_subplot(gs[1, 2])
if nn_dist.size > 0:
    nn_nm = nn_dist * PIXEL_SIZE_NM
    ax6.hist(nn_nm, bins=30, color='steelblue', edgecolor='white', linewidth=0.5)
    ax6.axvline(mean_nn_nm, color='red',    linestyle='--', label=f'μ={mean_nn_nm:.3f} nm')
    ax6.axvline(mean_nn_nm + std_nn_nm, color='orange', linestyle=':', label=f'σ={std_nn_nm:.3f} nm')
    ax6.axvline(mean_nn_nm - std_nn_nm, color='orange', linestyle=':')
    ax6.legend(fontsize=7)
    ax6.set_xlabel('NN distance (nm)', fontsize=9)
    ax6.set_ylabel('Count', fontsize=9)
ax6.set_title('8. NN Distance Distribution', fontsize=10)
ax6.tick_params(labelsize=8)

ax7 = fig.add_subplot(gs[2, 0])
ax7.imshow(img_bg, cmap='gray', origin='upper')
if len(refined) > 0:
    ax7.scatter([p['c'] for p in refined], [p['r'] for p in refined],
                s=10, c='yellow', linewidths=0, alpha=0.7, zorder=3)
origin_plot = [W/2, H/2]
if a_vec is not None:
    origin_rc = pos.mean(axis=0) if len(refined) > 0 else np.array([H/2, W/2])
    origin_plot = [origin_rc[1], origin_rc[0]]
    ax7.annotate('', xy=(origin_plot[0]+a_vec[0]*2, origin_plot[1]+a_vec[1]*2),
                 xytext=(origin_plot[0], origin_plot[1]),
                 arrowprops=dict(arrowstyle='->', color='cyan', lw=2))
    ax7.text(origin_plot[0]+a_vec[0]*2+2, origin_plot[1]+a_vec[1]*2,
             f'a={a_len_nm:.3f}nm', color='cyan', fontsize=7)
if b_vec is not None:
    ax7.annotate('', xy=(origin_plot[0]+b_vec[0]*2, origin_plot[1]+b_vec[1]*2),
                 xytext=(origin_plot[0], origin_plot[1]),
                 arrowprops=dict(arrowstyle='->', color='lime', lw=2))
    ax7.text(origin_plot[0]+b_vec[0]*2+2, origin_plot[1]+b_vec[1]*2,
             f'b={b_len_nm:.3f}nm', color='lime', fontsize=7)
ax7.set_title('9. Lattice Vectors Overlay', fontsize=10)
ax7.axis('off')

ax8 = fig.add_subplot(gs[2, 1])
ax8.imshow(img_bg, cmap='gray', origin='upper')
if len(refined) > 0:
    amps_plot = np.array([p['amp'] for p in refined])
    sc = ax8.scatter([p['c'] for p in refined], [p['r'] for p in refined],
                     c=amps_plot, cmap='hot', s=20, linewidths=0, zorder=3)
    plt.colorbar(sc, ax=ax8, fraction=0.03, pad=0.02, label='Amplitude')
ax8.set_title('12. Column Intensity Map', fontsize=10)
ax8.axis('off')

ax9 = fig.add_subplot(gs[2, 2])
ax9.imshow(img_bg, cmap='gray', origin='upper')
if displacements_nm and len(displacements_nm) == len(refined):
    disp_arr = np.array(displacements_nm)
    sc2 = ax9.scatter([p['c'] for p in refined], [p['r'] for p in refined],
                      c=disp_arr, cmap='RdYlGn_r', s=20, linewidths=0, zorder=3)
    plt.colorbar(sc2, ax=ax9, fraction=0.03, pad=0.02, label='Displacement (nm)')
elif len(refined) > 0:
    ax9.scatter([p['c'] for p in refined], [p['r'] for p in refined],
                s=12, c='gray', linewidths=0, zorder=3, alpha=0.6)
ax9.set_title('12. Displacement Field', fontsize=10)
ax9.axis('off')

fig.suptitle('STEM Atomic Column Analysis — YBCO', fontsize=14, fontweight='bold', y=1.01)
plt.savefig('analysis_visualization.png', dpi=150, bbox_inches='tight')
plt.close(fig)
print('Figure saved: analysis_visualization.png')

# ── save JSON results ─────────────────────────────────────────────────────────
results = {
    'image_path':             img_path,
    'image_size':             [int(W), int(H)],
    'pixel_size_nm_per_px':   float(PIXEL_SIZE_NM),
    'fft_lattice_params': {
        'a_nm': float(lat_a_nm),
        'b_nm': float(lat_b_nm),
        'n_fft_peaks': int(len(fft_results)),
        'peaks': fft_results
    },
    'fitted_lattice': {
        'a_nm':    float(a_len_nm),
        'b_nm':    float(b_len_nm),
        'gamma_deg': float(gamma_deg),
        'a_vec_px': [float(a_vec[0]), float(a_vec[1])] if a_vec is not None else None,
        'b_vec_px': [float(b_vec[0]), float(b_vec[1])] if b_vec is not None else None,
    },
    'column_detection': {
        'n_columns':        int(n_columns),
        'mean_amplitude':   float(mean_amp),
        'std_amplitude':    float(std_amp),
        'mean_nn_dist_nm':  float(mean_nn_nm),
        'std_nn_dist_nm':   float(std_nn_nm),
        'cv_nn_dist':       float(cv_nn),
    },
    'sublattice': {
        'n_sublattices':    int(n_sub) if sublattice_labels is not None else 0,
        'labels':           sublattice_labels.tolist() if sublattice_labels is not None else [],
        'ybco_check':       ybco_ratios_note,
    },
    'displacement_field': {
        'mean_displacement_nm': float(mean_disp_nm),
        'std_displacement_nm':  float(std_disp_nm),
        'n_measured':           int(len(displacements_nm)),
    },
    'columns': [
        {'r': float(p['r']), 'c': float(p['c']),
         'amp': float(p['amp']), 'sig': float(p['sig']),
         'sublattice': int(sublattice_labels[i]) if sublattice_labels is not None else -1}
        for i, p in enumerate(refined)
    ]
}

with open('results.json', 'w') as f:
    json.dump(results, f, indent=2)
print('Results saved: results.json')

print('\n=== SUMMARY ===')
print(f'  Columns detected  : {n_columns}')
print(f'  FFT lattice a     : {lat_a_nm:.3f} nm')
print(f'  FFT lattice b     : {lat_b_nm:.3f} nm')
print(f'  Fitted |a|        : {a_len_nm:.3f} nm')
print(f'  Fitted |b|        : {b_len_nm:.3f} nm')
print(f'  Lattice angle γ   : {gamma_deg:.1f}°')
print(f'  Mean NN dist      : {mean_nn_nm:.3f} nm')
print(f'  NN dist CV        : {cv_nn:.4f}')
print(f'  Mean displacement : {mean_disp_nm:.4f} nm')

# ── REQUIRED: print JSON results to stdout ────────────────────────────────────
print('IMAGE_ANALYSIS_RESULTS_JSON:' + json.dumps(results))
