import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from skimage import exposure, filters
import cv2
from scipy.spatial import cKDTree
from scipy.optimize import curve_fit
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# STEP 1: Load image and calibrate pixel size
# ============================================================
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260401_082752/results/analysis__easy__LLTO_ImageAnalysis_20260401_083025_001/image_0000/temp_image_0.npy'
image_raw = np.load(path).astype(np.float64)
print(f"Image shape: {image_raw.shape}, dtype: {image_raw.dtype}")
print(f"Intensity range: [{image_raw.min():.1f}, {image_raw.max():.1f}]")

# Pixel size calibration from known FOV
image_size_px = image_raw.shape[0]  # 1024
FOV_nm = 20.88  # known FOV in nm
FOV_angstrom = FOV_nm * 10  # 208.8 Angstrom
pixel_size = FOV_angstrom / image_size_px  # ~0.2039 Angstrom/px
print(f"Pixel size: {pixel_size:.4f} Angstrom/px")

# ============================================================
# STEP 2: Local contrast normalization (CLAHE)
# ============================================================
# Normalize to 0-1 for processing
img_norm = (image_raw - image_raw.min()) / (image_raw.max() - image_raw.min())

# Convert to uint8 for CLAHE
img_uint8 = (img_norm * 255).astype(np.uint8)

# Apply CLAHE with tile size ~64 px
clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(16, 16))  # 1024/16 = 64 px tiles
img_clahe_uint8 = clahe.apply(img_uint8)
img_clahe = img_clahe_uint8.astype(np.float64) / 255.0

print(f"CLAHE applied: tile=64px, clipLimit=2.5")
print(f"CLAHE image range: [{img_clahe.min():.3f}, {img_clahe.max():.3f}]")

# ============================================================
# STEP 3: FFT analysis
# ============================================================
# Compute FFT of CLAHE-normalized image
fft_image = np.fft.fft2(img_clahe)
fft_shift = np.fft.fftshift(fft_image)
fft_mag = np.abs(fft_shift)
fft_log = np.log1p(fft_mag)

center = np.array([image_size_px // 2, image_size_px // 2])

# Find peaks in FFT (excluding DC)
fft_log_search = fft_log.copy()
# Mask out DC region
y_grid, x_grid = np.mgrid[0:image_size_px, 0:image_size_px]
dist_from_center = np.sqrt((y_grid - center[0])**2 + (x_grid - center[1])**2)
fft_log_search[dist_from_center < 10] = 0
fft_log_search[dist_from_center > 400] = 0  # ignore very high freq

# Find peaks via local maximum filter
from scipy.ndimage import maximum_filter, label
fft_smooth = ndimage.gaussian_filter(fft_log_search, sigma=1.5)
local_max = maximum_filter(fft_smooth, size=15)
peak_mask = (fft_smooth == local_max) & (fft_smooth > np.percentile(fft_smooth[fft_smooth > 0], 97))
peak_labels, num_peaks = label(peak_mask)

fft_peaks = []
for i in range(1, min(num_peaks + 1, 50)):
    ys, xs = np.where(peak_labels == i)
    cy, cx = ys.mean(), xs.mean()
    dist = np.sqrt((cy - center[0])**2 + (cx - center[1])**2)
    intensity = fft_log[int(round(cy)), int(round(cx))]
    fft_peaks.append({'y': cy, 'x': cx, 'dist_px': dist, 'intensity': intensity})

fft_peaks.sort(key=lambda p: p['dist_px'])

# Convert FFT peak distances to real-space periods
print("\nFFT peaks (sorted by distance from center):")
superstructure_peak = None
fundamental_peaks = []
for p in fft_peaks[:20]:
    if p['dist_px'] > 5:
        real_space_period = (image_size_px / p['dist_px']) * pixel_size
        p['real_space_A'] = real_space_period
        print(f"  dist={p['dist_px']:.1f} px -> period={real_space_period:.2f} A, intensity={p['intensity']:.2f}")
        # Identify superstructure peak (~8 A)
        if 7.0 < real_space_period < 9.5 and superstructure_peak is None:
            superstructure_peak = p
        # Identify fundamental peaks (~3.87 A)
        if 3.4 < real_space_period < 4.3:
            fundamental_peaks.append(p)

if superstructure_peak:
    ss_direction = np.array([superstructure_peak['x'] - center[1], superstructure_peak['y'] - center[0]])
    ss_direction_norm = ss_direction / np.linalg.norm(ss_direction)
    print(f"\nSuperstructure peak: period={superstructure_peak['real_space_A']:.2f} A, direction=({ss_direction_norm[0]:.3f}, {ss_direction_norm[1]:.3f})")
else:
    ss_direction_norm = np.array([1.0, 0.0])  # default
    print("\nNo superstructure peak found in 7-9.5 A range")

print(f"Fundamental peaks found: {len(fundamental_peaks)}")

# ============================================================
# STEP 4: Bandpass filtering
# ============================================================
# Bandpass: retain 1.5-8 A -> 7-40 px
low_sigma = 1.2   # high-freq cutoff (Gaussian sigma for low-pass, removes noise < ~7 px)
high_sigma = 20.0  # low-freq cutoff (Gaussian sigma for high-pass, removes > ~40 px)

# Bandpass = low-pass - very-low-pass (difference of Gaussians)
img_lowpass = ndimage.gaussian_filter(img_clahe, sigma=low_sigma)
img_very_lowpass = ndimage.gaussian_filter(img_clahe, sigma=high_sigma)
img_bandpass = img_lowpass - img_very_lowpass

# Normalize
img_bandpass = (img_bandpass - img_bandpass.min()) / (img_bandpass.max() - img_bandpass.min())
print(f"\nBandpass filter applied: low_sigma={low_sigma}, high_sigma={high_sigma}")

# ============================================================
# STEP 5: Optional gentle smoothing
# ============================================================
img_filtered = ndimage.gaussian_filter(img_bandpass, sigma=0.8)
print("Gentle Gaussian smoothing applied (sigma=0.8)")

# ============================================================
# STEP 6: Atomic column detection
# ============================================================
# Import detect_atoms etc from scilink if available, otherwise implement
try:
    from scilink.tools.atom_finding import detect_atoms, find_zone_axes, find_missing_atoms, subtract_atoms
    HAS_SCILINK_ATOMS = True
    print("Using scilink atom finding tools")
except ImportError:
    HAS_SCILINK_ATOMS = False
    print("scilink atom tools not available, using custom implementation")

def custom_detect_atoms(image, separation, threshold_rel=0.02, refine=True, percent_to_nn=0.4):
    """Custom atomic column detection using peak_local_max + Gaussian refinement."""
    from skimage.feature import peak_local_max
    
    # Detect peaks
    min_distance = max(int(separation * 0.7), 3)
    coordinates = peak_local_max(image, min_distance=min_distance, 
                                  threshold_rel=threshold_rel,
                                  exclude_border=5)
    
    positions = []  # (x, y) where x=col, y=row
    sigma_x_list = []
    sigma_y_list = []
    amplitude_list = []
    rotation_list = []
    
    half_win = max(int(separation * percent_to_nn), 4)
    
    for row, col in coordinates:
        # Extract local window
        r0 = max(0, row - half_win)
        r1 = min(image.shape[0], row + half_win + 1)
        c0 = max(0, col - half_win)
        c1 = min(image.shape[1], col + half_win + 1)
        
        patch = image[r0:r1, c0:c1]
        
        if refine and patch.size > 9:
            try:
                # 2D Gaussian fit
                yy, xx = np.mgrid[0:patch.shape[0], 0:patch.shape[1]]
                yy = yy.astype(float)
                xx = xx.astype(float)
                
                def gaussian2d(coords, amp, x0, y0, sx, sy, bg):
                    x, y = coords
                    return (amp * np.exp(-0.5 * ((x - x0)**2 / sx**2 + (y - y0)**2 / sy**2)) + bg).ravel()
                
                p0 = [patch.max() - patch.min(), 
                      col - c0, row - r0,
                      2.5, 2.5, patch.min()]
                bounds_low = [0, -2, -2, 1.0, 1.0, -np.inf]
                bounds_high = [np.inf, patch.shape[1] + 2, patch.shape[0] + 2, 8.0, 8.0, np.inf]
                
                popt, _ = curve_fit(gaussian2d, (xx, yy), patch.ravel(), p0=p0,
                                   bounds=(bounds_low, bounds_high), maxfev=500)
                
                amp_fit, x0_fit, y0_fit, sx_fit, sy_fit, bg_fit = popt
                
                # Convert back to image coordinates
                refined_col = c0 + x0_fit
                refined_row = r0 + y0_fit
                
                # Validate sigma bounds 1.5-7 px
                if 1.0 <= sx_fit <= 7.5 and 1.0 <= sy_fit <= 7.5:
                    positions.append([refined_col, refined_row])
                    sigma_x_list.append(sx_fit)
                    sigma_y_list.append(sy_fit)
                    amplitude_list.append(amp_fit)
                    rotation_list.append(0.0)
                else:
                    # Use unrefined position with default sigma
                    positions.append([float(col), float(row)])
                    sigma_x_list.append(2.5)
                    sigma_y_list.append(2.5)
                    amplitude_list.append(float(patch.max() - patch.min()))
                    rotation_list.append(0.0)
            except:
                positions.append([float(col), float(row)])
                sigma_x_list.append(2.5)
                sigma_y_list.append(2.5)
                amplitude_list.append(float(patch.max() - patch.min()))
                rotation_list.append(0.0)
        else:
            positions.append([float(col), float(row)])
            sigma_x_list.append(2.5)
            sigma_y_list.append(2.5)
            amplitude_list.append(float(patch.max() - patch.min()))
            rotation_list.append(0.0)
    
    return {
        'positions': np.array(positions),
        'sigma_x': np.array(sigma_x_list),
        'sigma_y': np.array(sigma_y_list),
        'amplitude': np.array(amplitude_list),
        'rotation': np.array(rotation_list)
    }

def custom_find_zone_axes(positions, n_neighbors=9, distance_tolerance=None):
    """Find lattice translation vectors from atomic positions."""
    from sklearn.cluster import DBSCAN
    
    tree = cKDTree(positions)
    # Get displacement vectors to n nearest neighbors
    dists, indices = tree.query(positions, k=n_neighbors + 1)
    
    all_displacements = []
    for i in range(len(positions)):
        for j in range(1, n_neighbors + 1):
            nn_idx = indices[i, j]
            dx = positions[nn_idx, 0] - positions[i, 0]
            dy = positions[nn_idx, 1] - positions[i, 1]
            all_displacements.append([dx, dy])
    
    all_displacements = np.array(all_displacements)
    
    # Cluster displacement vectors
    eps = 3.0 if distance_tolerance is None else distance_tolerance
    clustering = DBSCAN(eps=eps, min_samples=max(10, len(positions) // 20)).fit(all_displacements)
    
    labels = clustering.labels_
    unique_labels = set(labels) - {-1}
    
    cluster_centers = []
    for lab in unique_labels:
        mask = labels == lab
        center = all_displacements[mask].mean(axis=0)
        count = mask.sum()
        cluster_centers.append({'center': center, 'count': count, 'dist': np.linalg.norm(center)})
    
    # Sort by distance
    cluster_centers.sort(key=lambda c: c['dist'])
    
    # Filter: keep vectors that are not inverses of each other
    zone_axes = []
    for cc in cluster_centers:
        if cc['dist'] < 3.0:  # too close to origin
            continue
        # Check if inverse already exists
        is_inverse = False
        for za in zone_axes:
            diff = cc['center'] + za  # inverse check
            if np.linalg.norm(diff) < eps:
                is_inverse = True
                break
        if not is_inverse:
            zone_axes.append(cc['center'])
        if len(zone_axes) >= 4:
            break
    
    return zone_axes

# Expected nearest-neighbor distance in pixels
expected_nn_px = 3.87 / pixel_size  # ~19 px for pseudocubic lattice parameter
# But in perovskite [001], A-B distance is a*sqrt(2)/2 ~ 2.74 A ~ 13.4 px
# A-A or B-B distance is a ~ 3.87 A ~ 19 px
# All columns (A+B) nearest neighbor is ~2.74 A ~ 13.4 px

separation_px = int(expected_nn_px * 0.5)  # ~9-10 px (slightly below A-B distance)
print(f"\nExpected lattice parameter: {3.87/pixel_size:.1f} px ({3.87:.2f} A)")
print(f"Expected A-B distance: {2.74/pixel_size:.1f} px ({2.74:.2f} A)")
print(f"Using separation: {separation_px} px")

# Run detection
if HAS_SCILINK_ATOMS:
    result1 = detect_atoms(img_filtered, separation=separation_px, threshold_rel=0.02, 
                           refine=True, percent_to_nn=0.4)
else:
    result1 = custom_detect_atoms(img_filtered, separation=separation_px, 
                                   threshold_rel=0.02, refine=True, percent_to_nn=0.4)

positions_all = result1['positions']  # (N, 2) as (x, y) where x=col, y=row
amplitudes_all = result1['amplitude']
sigma_x_all = result1['sigma_x']
sigma_y_all = result1['sigma_y']

print(f"\nInitial detection: {len(positions_all)} columns")

# Filter by amplitude (reject low-intensity false positives)
if len(positions_all) > 0:
    median_amp = np.median(amplitudes_all)
    amp_threshold = 0.15 * median_amp
    valid_mask = amplitudes_all > amp_threshold
    
    # Also reject sigma outliers
    valid_mask &= (sigma_x_all >= 1.0) & (sigma_x_all <= 7.5)
    valid_mask &= (sigma_y_all >= 1.0) & (sigma_y_all <= 7.5)
    
    positions_all = positions_all[valid_mask]
    amplitudes_all = amplitudes_all[valid_mask]
    sigma_x_all = sigma_x_all[valid_mask]
    sigma_y_all = sigma_y_all[valid_mask]
    
    print(f"After amplitude/sigma filtering: {len(positions_all)} columns")

# Get intensities from original image (HAADF) at each position
original_intensities = np.zeros(len(positions_all))
for i, (x, y) in enumerate(positions_all):
    r, c = int(round(y)), int(round(x))
    # Average in small window
    r0 = max(0, r - 2)
    r1 = min(image_raw.shape[0], r + 3)
    c0 = max(0, c - 2)
    c1 = min(image_raw.shape[1], c + 3)
    original_intensities[i] = image_raw[r0:r1, c0:c1].mean()

# ============================================================
# STEP 7: Detection coverage validation
# ============================================================
total_cols = len(positions_all)
print(f"\nTotal detected columns: {total_cols}")

# Check quadrant distribution
rows_all = positions_all[:, 1]
cols_all = positions_all[:, 0]
mid_r = image_size_px / 2
mid_c = image_size_px / 2

q1 = np.sum((rows_all < mid_r) & (cols_all < mid_c))
q2 = np.sum((rows_all < mid_r) & (cols_all >= mid_c))
q3 = np.sum((rows_all >= mid_r) & (cols_all < mid_c))
q4 = np.sum((rows_all >= mid_r) & (cols_all >= mid_c))
print(f"Quadrant distribution: TL={q1}, TR={q2}, BL={q3}, BR={q4}")

top60 = np.sum(rows_all < 0.6 * image_size_px)
bottom40 = np.sum(rows_all >= 0.6 * image_size_px)
print(f"Top 60%: {top60} columns, Bottom 40%: {bottom40} columns")
print(f"Bottom/Top density ratio: {(bottom40/0.4)/(top60/0.6+1e-10):.2f}")

# ============================================================
# STEP 8: Lattice vector determination
# ============================================================
if HAS_SCILINK_ATOMS:
    zone_axes = find_zone_axes(positions_all, n_neighbors=9)
else:
    zone_axes = custom_find_zone_axes(positions_all, n_neighbors=9)

print(f"\nZone axes found: {len(zone_axes)}")
lattice_vectors = []
for i, za in enumerate(zone_axes[:4]):
    mag_px = np.linalg.norm(za)
    mag_A = mag_px * pixel_size
    print(f"  Vector {i}: ({za[0]:.2f}, {za[1]:.2f}) px, magnitude = {mag_px:.2f} px = {mag_A:.2f} A")
    lattice_vectors.append({'dx': float(za[0]), 'dy': float(za[1]), 
                           'mag_px': float(mag_px), 'mag_A': float(mag_A)})

# Identify the two shortest vectors that form the unit cell
# For perovskite with all columns detected, shortest vectors should be A-B distance (~2.74 A)
# The lattice parameter vectors should be ~3.87 A
# We need vectors close to 3.87 A for the pseudocubic lattice

# Sort by magnitude
lattice_vectors.sort(key=lambda v: v['mag_A'])

# Find vectors closest to expected lattice parameter (~3.87 A)
# First check if shortest vectors are ~2.74 A (A-B vectors) or ~3.87 A (lattice vectors)
vec_a = None
vec_b = None

# Try to find two ~3.87 A vectors at ~90 degrees
candidates_387 = [v for v in lattice_vectors if 3.2 < v['mag_A'] < 4.5]
candidates_274 = [v for v in lattice_vectors if 2.2 < v['mag_A'] < 3.2]

print(f"\nCandidates near 3.87 A: {len(candidates_387)}")
print(f"Candidates near 2.74 A: {len(candidates_274)}")

if len(candidates_387) >= 2:
    # Use the 3.87 A vectors directly
    vec_a = candidates_387[0]
    # Find second vector at ~90 degrees
    for v in candidates_387[1:]:
        angle = np.degrees(np.arccos(np.clip(
            (vec_a['dx']*v['dx'] + vec_a['dy']*v['dy']) / (vec_a['mag_px'] * v['mag_px']), -1, 1)))
        if 70 < angle < 110:
            vec_b = v
            break
    if vec_b is None and len(candidates_387) >= 2:
        vec_b = candidates_387[1]
elif len(candidates_274) >= 2:
    # We have A-B vectors; lattice vectors would be combinations
    # a = v1 + v2, b = v1 - v2 (approximately)
    v1 = np.array([candidates_274[0]['dx'], candidates_274[0]['dy']])
    v2 = np.array([candidates_274[1]['dx'], candidates_274[1]['dy']])
    
    # Try different combinations
    combos = [v1 + v2, v1 - v2, -v1 + v2]
    for combo in combos:
        mag = np.linalg.norm(combo) * pixel_size
        if 3.2 < mag < 4.5:
            if vec_a is None:
                vec_a = {'dx': float(combo[0]), 'dy': float(combo[1]),
                         'mag_px': float(np.linalg.norm(combo)), 'mag_A': float(mag)}
            elif vec_b is None:
                vec_b = {'dx': float(combo[0]), 'dy': float(combo[1]),
                         'mag_px': float(np.linalg.norm(combo)), 'mag_A': float(mag)}

if vec_a is None or vec_b is None:
    # Fallback: use two shortest vectors
    if len(lattice_vectors) >= 2:
        vec_a = lattice_vectors[0]
        vec_b = lattice_vectors[1]
    else:
        vec_a = {'dx': 19.0, 'dy': 0.0, 'mag_px': 19.0, 'mag_A': 3.87}
        vec_b = {'dx': 0.0, 'dy': 19.0, 'mag_px': 19.0, 'mag_A': 3.87}

# Compute angle between lattice vectors
cos_angle = (vec_a['dx']*vec_b['dx'] + vec_a['dy']*vec_b['dy']) / (vec_a['mag_px'] * vec_b['mag_px'])
lattice_angle = np.degrees(np.arccos(np.clip(cos_angle, -1, 1)))

print(f"\nLattice vector a: ({vec_a['dx']:.2f}, {vec_a['dy']:.2f}) px, |a| = {vec_a['mag_A']:.2f} A")
print(f"Lattice vector b: ({vec_b['dx']:.2f}, {vec_b['dy']:.2f}) px, |b| = {vec_b['mag_A']:.2f} A")
print(f"Angle between a and b: {lattice_angle:.1f} degrees")

# ============================================================
# STEP 9: Sublattice separation via lattice-vector-based positional classification
# ============================================================
# Use the detected lattice vectors to define unit cell
# Transform all positions into fractional coordinates

# Build transformation matrix from lattice vectors
A_mat = np.array([[vec_a['dx'], vec_b['dx']],
                   [vec_a['dy'], vec_b['dy']]])

try:
    A_inv = np.linalg.inv(A_mat)
except np.linalg.LinAlgError:
    # If singular, use approximate values
    A_mat = np.array([[19.0, 0.0], [0.0, 19.0]])
    A_inv = np.linalg.inv(A_mat)

# Choose a reference position (first detected atom)
ref_pos = positions_all[0]

# Compute fractional coordinates for all positions
rel_positions = positions_all - ref_pos  # relative to reference
frac_coords = np.zeros((len(positions_all), 2))
for i in range(len(positions_all)):
    frac = A_inv @ rel_positions[i]
    # Map to [0, 1)
    frac_coords[i] = frac - np.floor(frac)

# Classify based on fractional coordinates
# A-site at (0, 0) corners -> frac near (0,0) or (1,1) etc
# B-site at (0.5, 0.5) body center

# Distance to (0,0) in fractional space (with periodic boundary)
def frac_dist_to_corner(frac):
    """Distance from fractional coord to nearest corner (0,0), (1,0), (0,1), (1,1)"""
    fx, fy = frac
    # Wrap to [-0.5, 0.5]
    fx_w = fx - round(fx)
    fy_w = fy - round(fy)
    return np.sqrt(fx_w**2 + fy_w**2)

def frac_dist_to_center(frac):
    """Distance from fractional coord to (0.5, 0.5)"""
    fx, fy = frac
    return np.sqrt((fx - 0.5)**2 + (fy - 0.5)**2)

sublattice_labels = np.zeros(len(positions_all), dtype=int)  # 0 = corner, 1 = body center

for i in range(len(positions_all)):
    d_corner = frac_dist_to_corner(frac_coords[i])
    d_center = frac_dist_to_center(frac_coords[i])
    if d_corner < d_center:
        sublattice_labels[i] = 0  # corner sublattice
    else:
        sublattice_labels[i] = 1  # body-center sublattice

mask_sub0 = sublattice_labels == 0
mask_sub1 = sublattice_labels == 1

count_sub0 = np.sum(mask_sub0)
count_sub1 = np.sum(mask_sub1)
ratio_01 = count_sub0 / max(count_sub1, 1)

print(f"\nSublattice 0 (corner): {count_sub0} columns")
print(f"Sublattice 1 (body-center): {count_sub1} columns")
print(f"Ratio: {ratio_01:.3f}")

# Determine which is A-site (brighter in HAADF = higher Z = La/Li) vs B-site (dimmer = Ti)
mean_intensity_sub0 = np.mean(original_intensities[mask_sub0])
mean_intensity_sub1 = np.mean(original_intensities[mask_sub1])

print(f"\nMean HAADF intensity - Sub0: {mean_intensity_sub0:.1f}, Sub1: {mean_intensity_sub1:.1f}")

if mean_intensity_sub0 > mean_intensity_sub1:
    A_mask = mask_sub0  # A-site = brighter = La/Li
    B_mask = mask_sub1  # B-site = dimmer = Ti
    print("Sublattice 0 -> A-site (La/Li, brighter)")
    print("Sublattice 1 -> B-site (Ti, dimmer)")
else:
    A_mask = mask_sub1
    B_mask = mask_sub0
    print("Sublattice 1 -> A-site (La/Li, brighter)")
    print("Sublattice 0 -> B-site (Ti, dimmer)")

A_positions = positions_all[A_mask]
B_positions = positions_all[B_mask]
A_intensities = original_intensities[A_mask]
B_intensities = original_intensities[B_mask]

A_count = len(A_positions)
B_count = len(B_positions)
AB_ratio = A_count / max(B_count, 1)

print(f"\nA-site (La/Li) columns: {A_count}")
print(f"B-site (Ti) columns: {B_count}")
print(f"A/B ratio: {AB_ratio:.3f}")
print(f"A-site mean intensity: {np.mean(A_intensities):.1f}")
print(f"B-site mean intensity: {np.mean(B_intensities):.1f}")

# ============================================================
# STEP 10: Nearest-neighbor distance validation
# ============================================================
def compute_nn_distances(pos):
    if len(pos) < 2:
        return np.array([0.0])
    tree = cKDTree(pos)
    dists, _ = tree.query(pos, k=2)
    return dists[:, 1]  # nearest neighbor distance

# Within A sublattice
if len(A_positions) > 1:
    nn_A_px = compute_nn_distances(A_positions)
    nn_A_A = nn_A_px * pixel_size
    nn_A_mean = np.mean(nn_A_A)
    nn_A_std = np.std(nn_A_A)
    nn_A_cv = nn_A_std / nn_A_mean * 100
    print(f"\nA-sublattice NN distance: {nn_A_mean:.2f} +/- {nn_A_std:.2f} A (CV={nn_A_cv:.1f}%)")
else:
    nn_A_mean, nn_A_cv = 0, 0

# Within B sublattice
if len(B_positions) > 1:
    nn_B_px = compute_nn_distances(B_positions)
    nn_B_A = nn_B_px * pixel_size
    nn_B_mean = np.mean(nn_B_A)
    nn_B_std = np.std(nn_B_A)
    nn_B_cv = nn_B_std / nn_B_mean * 100
    print(f"B-sublattice NN distance: {nn_B_mean:.2f} +/- {nn_B_std:.2f} A (CV={nn_B_cv:.1f}%)")
else:
    nn_B_mean, nn_B_cv = 0, 0

# Between sublattices (A to nearest B)
if len(A_positions) > 0 and len(B_positions) > 0:
    tree_B = cKDTree(B_positions)
    nn_AB_px, _ = tree_B.query(A_positions, k=1)
    nn_AB_A = nn_AB_px * pixel_size
    nn_AB_mean = np.mean(nn_AB_A)
    nn_AB_std = np.std(nn_AB_A)
    print(f"A-B inter-sublattice NN distance: {nn_AB_mean:.2f} +/- {nn_AB_std:.2f} A")
else:
    nn_AB_mean = 0

# ============================================================
# STEP 11: La/Li A-site ordering analysis
# ============================================================
ordering_period_A = None
ordering_strength = None
la_li_direction = None

if len(A_positions) > 10:
    # Project A-site positions onto the superstructure ordering direction
    # The ordering direction is perpendicular to the FFT superstructure peak direction in image space
    # (FFT peak direction in reciprocal space = real-space direction)
    
    ordering_dir = ss_direction_norm  # direction in image coordinates (x, y)
    
    # Project A-site positions onto this direction
    projections = A_positions @ ordering_dir
    
    # Sort by projection
    sort_idx = np.argsort(projections)
    sorted_proj = projections[sort_idx]
    sorted_int = A_intensities[sort_idx]
    
    # Bin the projections to create 1D intensity profile
    # Use bins of ~0.5 lattice parameter width
    bin_width_px = vec_a['mag_px'] * 0.3  # ~6 px bins
    proj_min = sorted_proj.min()
    proj_max = sorted_proj.max()
    n_bins = int((proj_max - proj_min) / bin_width_px)
    
    if n_bins > 10:
        bin_edges = np.linspace(proj_min, proj_max, n_bins + 1)
        bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
        bin_means = np.zeros(n_bins)
        bin_counts = np.zeros(n_bins)
        
        for i in range(len(sorted_proj)):
            bin_idx = int((sorted_proj[i] - proj_min) / bin_width_px)
            bin_idx = min(bin_idx, n_bins - 1)
            bin_means[bin_idx] += sorted_int[i]
            bin_counts[bin_idx] += 1
        
        valid_bins = bin_counts > 0
        bin_means[valid_bins] /= bin_counts[valid_bins]
        
        # Use only valid bins for analysis
        valid_centers = bin_centers[valid_bins] * pixel_size  # convert to Angstrom
        valid_means = bin_means[valid_bins]
        
        if len(valid_means) > 10:
            # FFT of the 1D intensity profile to find periodicity
            profile_mean = valid_means - np.mean(valid_means)
            
            # Zero-pad for better FFT resolution
            n_pad = max(1024, len(profile_mean) * 4)
            profile_padded = np.zeros(n_pad)
            profile_padded[:len(profile_mean)] = profile_mean
            
            fft_1d = np.abs(np.fft.rfft(profile_padded))
            freqs = np.fft.rfftfreq(n_pad, d=(valid_centers[1] - valid_centers[0]))
            
            # Look for peak in period range 6-10 A (corresponding to ~2a ordering)
            periods = np.zeros_like(freqs)
            periods[1:] = 1.0 / freqs[1:]
            
            # Find peak in the range 6-10 A
            mask_range = (periods > 6.0) & (periods < 10.0)
            if np.any(mask_range):
                peak_idx = np.argmax(fft_1d[mask_range])
                ordering_period_A = periods[mask_range][peak_idx]
                ordering_peak_power = fft_1d[mask_range][peak_idx]
                print(f"\nA-site ordering period from 1D profile FFT: {ordering_period_A:.2f} A")
            
            # Also try direct row-by-row intensity analysis
            # Group A-site columns by rows (along ordering direction)
            row_spacing = vec_a['mag_px']  # one lattice parameter in px
            # Assign each A-site column to a row based on projection
            row_assignments = np.round((projections - projections.min()) / row_spacing).astype(int)
            unique_rows = np.unique(row_assignments)
            
            row_mean_intensities = []
            row_positions = []
            for r in unique_rows:
                mask_r = row_assignments == r
                if np.sum(mask_r) >= 3:  # need enough columns per row
                    row_mean_intensities.append(np.mean(A_intensities[mask_r]))
                    row_positions.append(r * row_spacing * pixel_size)  # in Angstrom
            
            row_mean_intensities = np.array(row_mean_intensities)
            row_positions = np.array(row_positions)
            
            if len(row_mean_intensities) > 4:
                # Check for alternating bright/dim pattern
                diffs = np.diff(row_mean_intensities)
                sign_changes = np.sum(np.diff(np.sign(diffs)) != 0)
                
                # Compute intensity contrast
                # Separate into odd and even rows
                even_int = row_mean_intensities[0::2]
                odd_int = row_mean_intensities[1::2]
                
                if len(even_int) > 0 and len(odd_int) > 0:
                    mean_even = np.mean(even_int)
                    mean_odd = np.mean(odd_int)
                    brighter = max(mean_even, mean_odd)
                    dimmer = min(mean_even, mean_odd)
                    ordering_strength = (brighter - dimmer) / (brighter + dimmer) * 2
                    
                    # Ordering period = 2 * row spacing
                    if ordering_period_A is None:
                        if len(row_positions) > 1:
                            mean_row_spacing = np.mean(np.diff(row_positions))
                            ordering_period_A = 2 * mean_row_spacing
                    
                    print(f"La-rich row mean intensity: {brighter:.1f}")
                    print(f"Li-rich row mean intensity: {dimmer:.1f}")
                    print(f"Ordering strength (contrast ratio): {ordering_strength:.3f}")
                    if ordering_period_A:
                        print(f"Ordering period: {ordering_period_A:.2f} A")
                    la_li_direction = f"({ss_direction_norm[0]:.3f}, {ss_direction_norm[1]:.3f})"
                    print(f"Ordering direction: {la_li_direction}")

# ============================================================
# STEP 12: Save outputs
# ============================================================
# Save all positions
np.save('atomic_positions_all.npy', positions_all)
np.save('A_site_positions.npy', A_positions)
np.save('B_site_positions.npy', B_positions)

# Save sublattice labels
label_array = np.zeros(len(positions_all), dtype=int)
label_array[A_mask] = 1  # A-site
label_array[B_mask] = 2  # B-site
np.save('sublattice_labels.npy', label_array)

# Save intensities
np.save('column_intensities.npy', original_intensities)

# ============================================================
# STEP 13: Visualization
# ============================================================
fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)

# 1. Original image
ax = axes[0, 0]
ax.imshow(image_raw, cmap='gray')
ax.set_title('Original HAADF Image', fontsize=10)
ax.axis('off')

# 2. CLAHE-normalized image
ax = axes[0, 1]
ax.imshow(img_clahe, cmap='gray')
ax.set_title('CLAHE Normalized', fontsize=10)
ax.axis('off')

# 3. Detected columns overlay
ax = axes[0, 2]
ax.imshow(image_raw, cmap='gray')
if len(A_positions) > 0:
    ax.scatter(A_positions[:, 0], A_positions[:, 1], c='red', s=3, alpha=0.6, label=f'A-site (La/Li) N={A_count}')
if len(B_positions) > 0:
    ax.scatter(B_positions[:, 0], B_positions[:, 1], c='cyan', s=3, alpha=0.6, label=f'B-site (Ti) N={B_count}')
ax.legend(fontsize=7, loc='upper right')
ax.set_title(f'Sublattice Assignment (A/B={AB_ratio:.2f})', fontsize=10)
ax.axis('off')

# 4. FFT
ax = axes[1, 0]
fft_display = fft_log.copy()
fft_display = (fft_display - fft_display.min()) / (fft_display.max() - fft_display.min())
# Crop center region for better visibility
crop = 150
c0r = center[0] - crop
c1r = center[0] + crop
c0c = center[1] - crop
c1c = center[1] + crop
ax.imshow(fft_display[c0r:c1r, c0c:c1c], cmap='inferno')
if superstructure_peak:
    sp_y = superstructure_peak['y'] - c0r
    sp_x = superstructure_peak['x'] - c0c
    ax.plot(sp_x, sp_y, 'wo', markersize=8, fillstyle='none', linewidth=1.5)
    ax.annotate(f"{superstructure_peak['real_space_A']:.1f} A", (sp_x+5, sp_y), color='white', fontsize=8)
ax.set_title('FFT (center crop)', fontsize=10)
ax.axis('off')

# 5. Nearest-neighbor distance histogram
ax = axes[1, 1]
if len(A_positions) > 1:
    ax.hist(nn_A_A, bins=50, alpha=0.6, color='red', label=f'A-A: {nn_A_mean:.2f}±{nn_A_std:.2f} Å')
if len(B_positions) > 1:
    ax.hist(nn_B_A, bins=50, alpha=0.6, color='cyan', label=f'B-B: {nn_B_mean:.2f}±{nn_B_std:.2f} Å')
if len(A_positions) > 0 and len(B_positions) > 0:
    ax.hist(nn_AB_A, bins=50, alpha=0.6, color='green', label=f'A-B: {nn_AB_mean:.2f} Å')
ax.axvline(3.87, color='k', linestyle='--', label='Expected a=3.87 Å')
ax.axvline(2.74, color='gray', linestyle='--', label='Expected A-B=2.74 Å')
ax.legend(fontsize=7)
ax.set_xlabel('Distance (Å)', fontsize=9)
ax.set_ylabel('Count', fontsize=9)
ax.set_title('NN Distances', fontsize=10)

# 6. A-site intensity map for ordering analysis
ax = axes[1, 2]
if len(A_positions) > 0:
    scatter = ax.scatter(A_positions[:, 0], A_positions[:, 1], c=A_intensities, 
                         cmap='hot', s=5, alpha=0.8)
    plt.colorbar(scatter, ax=ax, shrink=0.7, label='HAADF Intensity')
    ax.set_xlim(0, image_size_px)
    ax.set_ylim(image_size_px, 0)
    if ordering_period_A:
        ax.set_title(f'A-site Intensity Map\nOrdering period: {ordering_period_A:.1f} Å', fontsize=10)
    else:
        ax.set_title('A-site Intensity Map', fontsize=10)
ax.set_aspect('equal')

plt.tight_layout()
plt.savefig('analysis_visualization.png', dpi=100, bbox_inches='tight')
plt.close()
print("\nVisualization saved.")

# ============================================================
# STEP 14: Compile and print results
# ============================================================
results = {
    "analysis_type": "LLTO perovskite HAADF-STEM atomic column analysis with sublattice separation and A-site ordering",
    "extracted_features": {
        "pixel_size_A_per_px": round(pixel_size, 4),
        "FOV_nm": FOV_nm,
        "total_detected_columns": int(total_cols),
        "A_site_count": int(A_count),
        "B_site_count": int(B_count),
        "A_B_ratio": round(float(AB_ratio), 3),
        "lattice_vector_a_A": round(vec_a['mag_A'], 2),
        "lattice_vector_b_A": round(vec_b['mag_A'], 2),
        "lattice_vector_a_components_px": [round(vec_a['dx'], 2), round(vec_a['dy'], 2)],
        "lattice_vector_b_components_px": [round(vec_b['dx'], 2), round(vec_b['dy'], 2)],
        "lattice_angle_degrees": round(float(lattice_angle), 1),
        "A_sublattice_NN_distance_A": round(float(nn_A_mean), 2),
        "A_sublattice_NN_CV_percent": round(float(nn_A_cv), 1),
        "B_sublattice_NN_distance_A": round(float(nn_B_mean), 2),
        "B_sublattice_NN_CV_percent": round(float(nn_B_cv), 1),
        "AB_inter_sublattice_NN_distance_A": round(float(nn_AB_mean), 2),
        "A_site_mean_HAADF_intensity": round(float(np.mean(A_intensities)), 1),
        "B_site_mean_HAADF_intensity": round(float(np.mean(B_intensities)), 1),
        "superstructure_peak_period_A": round(float(superstructure_peak['real_space_A']), 2) if superstructure_peak else None,
        "ordering_period_realspace_A": round(float(ordering_period_A), 2) if ordering_period_A else None,
        "ordering_strength_contrast_ratio": round(float(ordering_strength), 3) if ordering_strength else None,
        "ordering_direction": la_li_direction,
        "quadrant_distribution": {"TL": int(q1), "TR": int(q2), "BL": int(q3), "BR": int(q4)}
    },
    "quality_metrics": {
        "detection_coverage_top60_pct": round(float(top60 / max(total_cols, 1) * 100), 1),
        "detection_coverage_bottom40_pct": round(float(bottom40 / max(total_cols, 1) * 100), 1),
        "A_sublattice_NN_CV_percent": round(float(nn_A_cv), 1),
        "B_sublattice_NN_CV_percent": round(float(nn_B_cv), 1),
        "A_B_ratio_deviation_from_1": round(abs(float(AB_ratio) - 1.0), 3)
    },
    "summary": f"Detected {total_cols} atomic columns in LLTO HAADF-STEM image. Sublattice separation via lattice-vector-based positional classification yielded {A_count} A-site (La/Li) and {B_count} B-site (Ti) columns (A/B={AB_ratio:.2f}). Lattice parameters: a={vec_a['mag_A']:.2f} A, b={vec_b['mag_A']:.2f} A, angle={lattice_angle:.1f} deg. A-site ordering period: {ordering_period_A:.1f} A." if ordering_period_A else f"Detected {total_cols} atomic columns. A/B={AB_ratio:.2f}. a={vec_a['mag_A']:.2f} A, b={vec_b['mag_A']:.2f} A.",
    "saved_arrays": {
        "atomic_positions_all.npy": {
            "description": f"All {total_cols} atomic column positions as (x,y) where x=col, y=row in pixel coordinates",
            "shape": list(positions_all.shape),
            "dtype": str(positions_all.dtype)
        },
        "A_site_positions.npy": {
            "description": f"{A_count} A-site (La/Li) column positions as (x,y)",
            "shape": list(A_positions.shape),
            "dtype": str(A_positions.dtype)
        },
        "B_site_positions.npy": {
            "description": f"{B_count} B-site (Ti) column positions as (x,y)",
            "shape": list(B_positions.shape),
            "dtype": str(B_positions.dtype)
        },
        "sublattice_labels.npy": {
            "description": "Sublattice label for each detected column: 1=A-site (La/Li), 2=B-site (Ti)",
            "shape": list(label_array.shape),
            "dtype": str(label_array.dtype)
        },
        "column_intensities.npy": {
            "description": "HAADF intensity at each detected column position from original image",
            "shape": list(original_intensities.shape),
            "dtype": str(original_intensities.dtype)
        }
    }
}

print(f"IMAGE_ANALYSIS_RESULTS_JSON:{json.dumps(results)}")
