import numpy as np
import json
import os
import warnings
warnings.filterwarnings('ignore')

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

from scipy import ndimage, signal
from scipy.optimize import curve_fit
from scipy.interpolate import griddata
from scipy.fft import fft, fftfreq
from skimage.filters import threshold_otsu
from skimage.feature import peak_local_max
from skimage.measure import label, regionprops
import cv2

# ─────────────────────────────────────────────
# 1. Load image
# ─────────────────────────────────────────────
image_path = '/Users/maxim.ziatdinov/Code/SciLink/test_ybco_output/image_0000/temp_image_0.npy'
image_raw = np.load(image_path)
print(f'Image shape: {image_raw.shape}, dtype: {image_raw.dtype}')
print(f'Intensity range: [{image_raw.min()}, {image_raw.max()}]')

if image_raw.ndim == 3:
    img = image_raw[:, :, 0].astype(np.float64)
else:
    img = image_raw.astype(np.float64)

# Normalise to [0, 1]
img_min, img_max = img.min(), img.max()
img_norm = (img - img_min) / (img_max - img_min)

H, W = img_norm.shape
print(f'Working image shape: {H}x{W}')

# ─────────────────────────────────────────────
# 2. Attempt to load Tier 1 JSON; fall back to detection
# ─────────────────────────────────────────────
tier1_dir = '/Users/maxim.ziatdinov/Code/SciLink/test_ybco_output/image_0000'
tier1_json_path = os.path.join(tier1_dir, 'results.json')

columns = None  # list of dicts with keys: x, y, intensity, sublattice
sublattice_labels = None

try:
    with open(tier1_json_path, 'r') as f:
        tier1 = json.load(f)
    print('Loaded Tier 1 JSON')
    # Try to extract column positions from various possible JSON structures
    if 'columns' in tier1:
        columns = tier1['columns']
    elif 'extracted_features' in tier1 and 'columns' in tier1['extracted_features']:
        columns = tier1['extracted_features']['columns']
except Exception as e:
    print(f'No Tier 1 JSON available ({e}); will detect columns from scratch.')

# ─────────────────────────────────────────────
# 3. Column detection (if not loaded from JSON)
# ─────────────────────────────────────────────
if columns is None:
    print('Running column detection pipeline...')
    
    # Background subtraction with large Gaussian
    bg = ndimage.gaussian_filter(img_norm, sigma=20)
    img_bg = img_norm - bg
    img_bg = (img_bg - img_bg.min()) / (img_bg.max() - img_bg.min() + 1e-12)
    
    # Bandpass: subtract broad blur, keep medium-scale features
    img_smooth = ndimage.gaussian_filter(img_norm, sigma=1.5)
    img_coarse = ndimage.gaussian_filter(img_norm, sigma=4.0)
    img_bp = img_smooth - img_coarse
    img_bp = (img_bp - img_bp.min()) / (img_bp.max() - img_bp.min() + 1e-12)
    
    # Peak detection on bandpass image
    # Expected inter-column spacing ~4-8 pixels for YBCO HAADF at typical magnification
    min_distance = 5  # pixels
    threshold_abs = 0.2
    
    coords = peak_local_max(
        img_bp,
        min_distance=min_distance,
        threshold_abs=threshold_abs,
        exclude_border=5
    )
    print(f'Detected {len(coords)} column candidates')
    
    # Refine with 2D Gaussian fit in small window
    def gaussian2d(xy, amp, x0, y0, sx, sy, bg_val):
        x, y = xy
        return amp * np.exp(-((x-x0)**2/(2*sx**2) + (y-y0)**2/(2*sy**2))) + bg_val
    
    refined_cols = []
    win = 4  # half-window
    for (row, col) in coords:
        r0, r1 = max(0, row-win), min(H, row+win+1)
        c0, c1 = max(0, col-win), min(W, col+win+1)
        patch = img_norm[r0:r1, c0:c1]
        pr, pc = np.mgrid[r0:r1, c0:c1]
        try:
            p0 = [patch.max()-patch.min(), float(row), float(col), 1.5, 1.5, patch.min()]
            bounds_lo = [0, row-win, col-win, 0.3, 0.3, 0]
            bounds_hi = [2, row+win, col+win, 5, 5, 1]
            popt, _ = curve_fit(gaussian2d, (pr.ravel(), pc.ravel()),
                                 patch.ravel(), p0=p0,
                                 bounds=(bounds_lo, bounds_hi),
                                 maxfev=400)
            amp, x0, y0, sx, sy, bg_v = popt
            peak_int = amp
            refined_cols.append({'x': float(y0), 'y': float(x0),
                                  'intensity': float(peak_int),
                                  'sublattice': -1})
        except Exception:
            peak_int = float(img_norm[row, col])
            refined_cols.append({'x': float(col), 'y': float(row),
                                  'intensity': float(peak_int),
                                  'sublattice': -1})
    
    columns = refined_cols
    print(f'Refined to {len(columns)} columns')

# Convert to arrays
col_x = np.array([c['x'] for c in columns])  # x = horizontal (col)
col_y = np.array([c['y'] for c in columns])  # y = vertical (row)
col_int = np.array([c['intensity'] for c in columns])
col_sl = np.array([c.get('sublattice', -1) for c in columns], dtype=int)

print(f'Total columns: {len(col_x)}')
print(f'Intensity range of columns: [{col_int.min():.4f}, {col_int.max():.4f}]')

# ─────────────────────────────────────────────
# 4. Sublattice assignment via intensity clustering
# ─────────────────────────────────────────────
# If sublattice labels not from JSON, cluster by intensity
if np.all(col_sl == -1) or len(np.unique(col_sl[col_sl >= 0])) < 2:
    print('Assigning sublattices by intensity clustering...')
    from sklearn.cluster import KMeans
    
    # Use up to 4 sublattices (Ba, Y, Cu, O)
    n_sl = 4
    intensities = col_int.reshape(-1, 1)
    km = KMeans(n_clusters=n_sl, n_init=20, random_state=42)
    labels_km = km.fit_predict(intensities)
    
    # Sort clusters by mean intensity (descending: brightest = heaviest)
    cluster_means = [intensities[labels_km == k].mean() for k in range(n_sl)]
    order = np.argsort(cluster_means)[::-1]  # highest first
    remap = {old: new for new, old in enumerate(order)}
    col_sl = np.array([remap[l] for l in labels_km], dtype=int)
    print(f'Sublattice counts: {[int((col_sl==k).sum()) for k in range(n_sl)]}')
else:
    n_sl = int(col_sl.max()) + 1
    print(f'Using {n_sl} sublattices from JSON')

# ─────────────────────────────────────────────
# 5 (A). Per-sublattice intensity statistics & Z^1.7 comparison
# ─────────────────────────────────────────────
# Corrected Z^1.7 reference values (normalised to Ba=1.000)
Z_ref = {
    'Ba': 1.000,
    'Y':  0.542,
    'Cu': 0.351,
    'O':  0.034
}
Z_order = ['Ba', 'Y', 'Cu', 'O']
Z_vals = np.array([Z_ref[k] for k in Z_order])  # [1.000, 0.542, 0.351, 0.034]

sl_stats = {}
for k in range(n_sl):
    mask_k = col_sl == k
    if mask_k.sum() == 0:
        sl_stats[k] = {'mean': 0.0, 'std': 0.0, 'count': 0, 'element': 'unknown'}
        continue
    ints = col_int[mask_k]
    sl_stats[k] = {
        'mean': float(ints.mean()),
        'std':  float(ints.std()),
        'count': int(mask_k.sum())
    }

# Normalise observed means to sum/max for comparison
if n_sl >= 4:
    obs_means = np.array([sl_stats[k]['mean'] for k in range(4)])
else:
    obs_means = np.array([sl_stats[k]['mean'] for k in range(n_sl)])
    obs_means = np.pad(obs_means, (0, 4 - len(obs_means)))

# Normalise observed means so that max = 1.000 (assume SL0 is Ba)
if obs_means.max() > 0:
    obs_norm = obs_means / obs_means.max()
else:
    obs_norm = obs_means

# Assign chemical identities (SL0=Ba, SL1=Y, SL2=Cu, SL3=O)
chem_names = Z_order[:min(n_sl, 4)]
for k in range(min(n_sl, 4)):
    sl_stats[k]['element'] = chem_names[k]
    sl_stats[k]['obs_norm'] = float(obs_norm[k])
    sl_stats[k]['Z_pred'] = float(Z_vals[k])
    sl_stats[k]['residual'] = float(obs_norm[k] - Z_vals[k])

print('\nPer-sublattice stats:')
for k, v in sl_stats.items():
    print(f'  SL{k} ({v.get("element","?")}): mean={v["mean"]:.4f}, std={v["std"]:.4f}, '
          f'count={v["count"]}, obs_norm={v.get("obs_norm",0):.3f}, '
          f'Z_pred={v.get("Z_pred",0):.3f}, residual={v.get("residual",0):.3f}')

# Intensity ratio matrix (observed vs predicted)
obs_ratio_matrix = [float(obs_norm[k]) for k in range(min(n_sl, 4))]
pred_ratio_matrix = list(Z_vals[:min(n_sl, 4)])
residual_matrix = [float(obs_ratio_matrix[k] - pred_ratio_matrix[k]) for k in range(min(n_sl, 4))]

print(f'Observed ratios:  {obs_ratio_matrix}')
print(f'Predicted ratios: {pred_ratio_matrix}')
print(f'Residuals:        {residual_matrix}')

# ─────────────────────────────────────────────
# 5 (B). 2D spatial intensity map (32×32 grid)
# ─────────────────────────────────────────────
grid_n = 32
cell_size_px = H / grid_n  # pixels per cell
# Physical scale: assume YBCO HAADF ~0.0195 nm/pixel (512 px ~ 10 nm FOV)
# Adjust: typical HAADF YBCO images are ~5-10 nm FOV at 512 px
# cell_size_nm will be calculated from lattice spacing detected
# For now use a nominal scale; will refine from sublattice spacing

# Build 2D intensity grid using local averaging
intensity_grid = np.full((grid_n, grid_n), np.nan)

for gy in range(grid_n):
    for gx in range(grid_n):
        y0_cell = gy * cell_size_px
        y1_cell = (gy + 1) * cell_size_px
        x0_cell = gx * cell_size_px
        x1_cell = (gx + 1) * cell_size_px
        in_cell = ((col_y >= y0_cell) & (col_y < y1_cell) &
                   (col_x >= x0_cell) & (col_x < x1_cell))
        if in_cell.sum() > 0:
            intensity_grid[gy, gx] = col_int[in_cell].mean()

# Fill NaN by interpolation for visualization
grid_mask_valid = ~np.isnan(intensity_grid)
if grid_mask_valid.sum() > 3:
    gy_pts, gx_pts = np.where(grid_mask_valid)
    vals = intensity_grid[grid_mask_valid]
    gy_all, gx_all = np.mgrid[0:grid_n, 0:grid_n]
    try:
        intensity_grid_filled = griddata(
            (gy_pts, gx_pts), vals,
            (gy_all, gx_all), method='linear'
        )
        nan_mask2 = np.isnan(intensity_grid_filled)
        if nan_mask2.any():
            intensity_grid_filled[nan_mask2] = griddata(
                (gy_pts, gx_pts), vals,
                (gy_all[nan_mask2], gx_all[nan_mask2]), method='nearest'
            )
    except Exception:
        intensity_grid_filled = np.nan_to_num(intensity_grid, nan=float(np.nanmean(intensity_grid)))
else:
    intensity_grid_filled = np.nan_to_num(intensity_grid, nan=0.0)

# Identify anomalous cells (|deviation| > 2 sigma from mean)
grid_mean = np.nanmean(intensity_grid)
grid_std = np.nanstd(intensity_grid)
anomaly_threshold = 2.0
anomaly_map = np.abs(intensity_grid_filled - grid_mean) > anomaly_threshold * grid_std
num_anomalous_cells = int(anomaly_map.sum())
print(f'Anomalous grid cells (>2σ): {num_anomalous_cells} / {grid_n*grid_n}')

# ─────────────────────────────────────────────
# 5 (C). Dark-band characterisation
# ─────────────────────────────────────────────
# Row-averaged intensity profile
row_profile = img_norm.mean(axis=1)  # shape (H,)

# Smooth profile slightly to reduce noise
row_profile_smooth = ndimage.gaussian_filter1d(row_profile, sigma=2.0)

# Find dark bands via Otsu on inverted profile (dark = low intensity)
profile_inv = 1.0 - row_profile_smooth
otsu_val = threshold_otsu(profile_inv)
dark_mask_1d = profile_inv > otsu_val  # rows that are dark

# Label contiguous dark-band regions
dark_labels = label(dark_mask_1d.reshape(-1, 1)).ravel()
dark_regions = []
for reg_id in range(1, dark_labels.max() + 1):
    rows_in = np.where(dark_labels == reg_id)[0]
    if len(rows_in) >= 2:  # minimum width of 2 rows
        dark_regions.append({
            'row_start': int(rows_in.min()),
            'row_end':   int(rows_in.max()),
            'width_px':  int(rows_in.max() - rows_in.min() + 1),
            'mean_intensity': float(row_profile_smooth[rows_in].mean())
        })

print(f'\nDark band regions found: {len(dark_regions)}')
for dr in dark_regions:
    print(f'  rows {dr["row_start"]}-{dr["row_end"]}, width={dr["width_px"]}px')

# Physical scale estimation from nearest-neighbor distances
if len(col_x) > 10:
    from sklearn.neighbors import KDTree
    pts = np.stack([col_y, col_x], axis=1)
    kdt = KDTree(pts)
    nn_dists, _ = kdt.query(pts, k=2)
    nn_dists = nn_dists[:, 1]  # exclude self
    median_nn_px = float(np.median(nn_dists))
    # YBCO nearest-neighbour: ~0.385 nm (Cu-O plane spacing within unit cell)
    # Use Ba-Ba distance as reference ~0.385 nm in-plane
    # Typical HAADF: Ba columns ~0.38-0.39 nm apart
    ref_spacing_nm = 0.385  # nm
    nm_per_px = ref_spacing_nm / median_nn_px
else:
    median_nn_px = 6.0
    nm_per_px = 0.385 / 6.0

print(f'Estimated nm/pixel: {nm_per_px:.4f} (median NN = {median_nn_px:.2f} px)')
cell_size_nm = cell_size_px * nm_per_px
print(f'Grid cell size: {cell_size_nm:.3f} nm')

# Dark-band measurements in nm
for dr in dark_regions:
    dr['width_nm'] = dr['width_px'] * nm_per_px
    dr['center_nm'] = (dr['row_start'] + dr['width_px'] / 2) * nm_per_px

# Inter-band spacings
inter_band_spacings_nm = []
for i in range(len(dark_regions) - 1):
    gap_px = dark_regions[i+1]['row_start'] - dark_regions[i]['row_end'] - 1
    inter_band_spacings_nm.append(float(gap_px * nm_per_px))

# Fourier analysis of row-averaged profile
N = len(row_profile_smooth)
fft_vals = np.abs(fft(row_profile_smooth - row_profile_smooth.mean()))
freqs = fftfreq(N, d=nm_per_px)  # cycles / nm
fft_pos = fft_vals[:N//2]
freqs_pos = freqs[:N//2]

# Find dominant non-zero frequency
if len(freqs_pos) > 2:
    # Exclude DC (index 0)
    peak_idx = np.argmax(fft_pos[1:]) + 1
    dominant_freq = float(freqs_pos[peak_idx])
    dominant_period_nm = float(1.0 / dominant_freq) if dominant_freq > 0 else 0.0
else:
    dominant_period_nm = 0.0
    dominant_freq = 0.0

print(f'Dominant Fourier periodicity: {dominant_period_nm:.3f} nm (freq={dominant_freq:.4f} /nm)')

# YBCO c-axis layer spacings for correlation
ybco_spacings = {
    'full_unit_cell': 1.17,    # nm
    'half_unit_cell': 0.585,   # nm
    'layer_spacing': 0.38,     # nm (~c/3)
    'double_layer': 0.76,      # nm
}

# Correlation: compare dominant period to YBCO spacings
best_match_key = min(ybco_spacings, key=lambda k: abs(ybco_spacings[k] - dominant_period_nm))
best_match_val = ybco_spacings[best_match_key]
period_corr = 1.0 / (1.0 + abs(dominant_period_nm - best_match_val)) if dominant_period_nm > 0 else 0.0
print(f'Best YBCO spacing match: {best_match_key} ({best_match_val} nm), '
      f'correlation proxy: {period_corr:.3f}')

# Dark-band fraction of image area
total_dark_px = sum(dr['width_px'] for dr in dark_regions)
dark_fraction = float(total_dark_px / H)
print(f'Dark-band fraction of image rows: {dark_fraction:.3f}')

# ─────────────────────────────────────────────
# 5 (D). Strain/displacement analysis
# ─────────────────────────────────────────────
# Identify dark-band rows to exclude
dark_row_mask = dark_mask_1d.copy()  # True for dark rows

# Exclude columns in dark-band rows
clean_col_mask = np.array([
    not dark_row_mask[min(int(round(col_y[i])), H-1)]
    for i in range(len(col_y))
], dtype=bool)

print(f'\nClean columns (not in dark bands): {clean_col_mask.sum()} / {len(col_x)}')

# Fit ideal lattice to clean columns using least-squares
# Strategy: fit two lattice vectors a1, a2 via linear regression on integer indices
# Simplified: fit separate x and y lattice constants for a rectangular approximation

def fit_lattice_1d(positions, min_spacing=3.0):
    """Fit lattice constant and offsets from 1D positions."""
    positions = np.sort(positions)
    diffs = np.diff(positions)
    # Remove outliers
    med_diff = np.median(diffs)
    diffs_clean = diffs[(diffs > 0.3 * med_diff) & (diffs < 3.0 * med_diff)]
    if len(diffs_clean) == 0:
        return med_diff, positions[0]
    lattice_const = float(np.median(diffs_clean))
    # Find integer indices for each position
    n_indices = np.round((positions - positions[0]) / lattice_const).astype(int)
    if len(n_indices) > 1:
        # Linear fit: pos = a * n + offset
        from numpy.polynomial import polynomial as P
        coef = np.polyfit(n_indices, positions, 1)
        lattice_const = float(coef[0])
        offset = float(coef[1])
    else:
        offset = float(positions[0])
    return lattice_const, offset

if clean_col_mask.sum() > 20:
    cx_clean = col_x[clean_col_mask]
    cy_clean = col_y[clean_col_mask]
    ci_clean = col_int[clean_col_mask]
    
    # Estimate lattice vectors separately for x and y
    # For YBCO HAADF, columns typically in a roughly rectangular arrangement
    # Fit using KD-tree nearest-neighbor approach
    
    # Get approximate lattice constant from NN distances of clean columns
    pts_clean = np.stack([cy_clean, cx_clean], axis=1)
    from sklearn.neighbors import KDTree
    kdt_clean = KDTree(pts_clean)
    nn_d, _ = kdt_clean.query(pts_clean, k=2)
    nn_d = nn_d[:, 1]
    lat_const_px = float(np.median(nn_d))
    print(f'Lattice constant from clean columns: {lat_const_px:.2f} px = {lat_const_px*nm_per_px:.3f} nm')
    
    # Assign ideal lattice positions
    # Origin at centroid of clean columns
    cx_mean = cx_clean.mean()
    cy_mean = cy_clean.mean()
    
    # Find integer indices for each column relative to origin
    # Round to nearest lattice site
    ix = np.round((cx_clean - cx_mean) / lat_const_px).astype(int)
    iy = np.round((cy_clean - cy_mean) / lat_const_px).astype(int)
    
    # Ideal positions
    x_ideal = cx_mean + ix * lat_const_px
    y_ideal = cy_mean + iy * lat_const_px
    
    # Displacements
    dx = cx_clean - x_ideal
    dy = cy_clean - y_ideal
    disp_mag = np.sqrt(dx**2 + dy**2)
    disp_mag_nm = disp_mag * nm_per_px
    
    disp_mean = float(disp_mag_nm.mean())
    disp_std  = float(disp_mag_nm.std())
    disp_p95  = float(np.percentile(disp_mag_nm, 95))
    
    # Threshold for significant displacement: 0.12 nm
    disp_threshold_nm = 0.12
    large_disp_mask = disp_mag_nm > disp_threshold_nm
    n_large_disp = int(large_disp_mask.sum())
    
    print(f'Displacement stats (clean cols): mean={disp_mean:.4f} nm, '
          f'std={disp_std:.4f} nm, p95={disp_p95:.4f} nm')
    print(f'Columns with |disp| > {disp_threshold_nm} nm: {n_large_disp}')
else:
    print('Insufficient clean columns for displacement analysis')
    cx_clean = col_x
    cy_clean = col_y
    lat_const_px = median_nn_px
    dx = np.zeros(len(cx_clean))
    dy = np.zeros(len(cy_clean))
    disp_mag_nm = np.zeros(len(cx_clean))
    disp_mean = 0.0
    disp_std  = 0.0
    disp_p95  = 0.0
    n_large_disp = 0
    large_disp_mask = np.zeros(len(cx_clean), dtype=bool)

# Spatial distribution of large-displacement columns
if n_large_disp > 0 and clean_col_mask.sum() > 0:
    ld_x = cx_clean[large_disp_mask]
    ld_y = cy_clean[large_disp_mask]
    ld_centroid_x = float(ld_x.mean())
    ld_centroid_y = float(ld_y.mean())
else:
    ld_centroid_x = 0.0
    ld_centroid_y = 0.0

# ─────────────────────────────────────────────
# 6. Save arrays
# ─────────────────────────────────────────────
# Column positions and attributes
col_positions = np.stack([col_x, col_y, col_int, col_sl.astype(float)], axis=1)
np.save('column_positions.npy', col_positions)

# 2D intensity grid
np.save('intensity_grid.npy', intensity_grid_filled.astype(np.float32))

# Row profile
np.save('row_intensity_profile.npy', row_profile_smooth.astype(np.float32))

# Displacement field for clean columns
if clean_col_mask.sum() > 0:
    disp_array = np.stack([cx_clean, cy_clean, disp_mag_nm], axis=1)
else:
    disp_array = np.zeros((1, 3))
np.save('displacement_field.npy', disp_array.astype(np.float32))

print('Saved arrays: column_positions.npy, intensity_grid.npy, '
      'row_intensity_profile.npy, displacement_field.npy')

# ─────────────────────────────────────────────
# 7. Visualization
# ─────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(18, 12))

# Panel 1: Original image with detected columns
ax = axes[0, 0]
ax.imshow(img_norm, cmap='gray', origin='upper')
colors_sl = ['red', 'lime', 'blue', 'yellow']
for k in range(min(n_sl, 4)):
    mask_k = col_sl == k
    if mask_k.sum() > 0:
        elem = sl_stats.get(k, {}).get('element', f'SL{k}')
        ax.scatter(col_x[mask_k], col_y[mask_k], s=4,
                   c=colors_sl[k], alpha=0.6, label=elem)
ax.legend(fontsize=7, markerscale=2)
ax.set_title('Detected Columns (coloured by sublattice)', fontsize=10)
ax.set_xlabel('x (px)'); ax.set_ylabel('y (px)')

# Panel 2: 2D intensity grid with anomaly overlay
ax = axes[0, 1]
im2 = ax.imshow(intensity_grid_filled, cmap='viridis', origin='upper',
                extent=[0, W, H, 0])
plt.colorbar(im2, ax=ax, label='Norm. intensity')
# Overlay anomalies
anom_y, anom_x = np.where(anomaly_map)
cell = H / grid_n
for ay, ax2 in zip(anom_y, anom_x):
    rect = plt.Rectangle((ax2*cell, ay*cell), cell, cell,
                          fill=False, edgecolor='red', linewidth=1.0)
    ax.add_patch(rect)
ax.set_title(f'2D Intensity Map (32×32)\nRed = anomalous cells (>2σ)', fontsize=10)
ax.set_xlabel('x (px)'); ax.set_ylabel('y (px)')

# Panel 3: Row-averaged intensity profile + dark bands
ax = axes[0, 2]
y_px = np.arange(H)
y_nm = y_px * nm_per_px
ax.plot(row_profile_smooth, y_nm, 'b-', linewidth=1, label='Row avg')
for dr in dark_regions:
    ax.axhspan(dr['row_start']*nm_per_px, dr['row_end']*nm_per_px,
               alpha=0.3, color='gray', label='Dark band')
ax.set_ylabel('y (nm)'); ax.set_xlabel('Norm. intensity')
ax.set_title('Row-Averaged Intensity Profile\n(gray = dark bands)', fontsize=10)
ax.invert_yaxis()
ax.legend(fontsize=7)

# Panel 4: Fourier spectrum of row profile
ax = axes[1, 0]
ax.semilogy(freqs_pos[1:], fft_pos[1:], 'b-', linewidth=0.8)
ax.axvline(dominant_freq, color='red', linestyle='--',
           label=f'Peak: {dominant_period_nm:.2f} nm')
# Mark YBCO spacings
for name, sp in ybco_spacings.items():
    if sp > 0:
        ax.axvline(1.0/sp, color='orange', linestyle=':', alpha=0.7, linewidth=0.8)
ax.set_xlabel('Spatial frequency (1/nm)')
ax.set_ylabel('FFT amplitude (log)')
ax.set_title(f'Fourier Spectrum of Row Profile\nPeriod={dominant_period_nm:.3f} nm', fontsize=10)
ax.legend(fontsize=7)

# Panel 5: Z^1.7 comparison bar chart
ax = axes[1, 1]
n_show = min(n_sl, 4)
bar_x = np.arange(n_show)
bar_w = 0.35

obs_vals = [sl_stats.get(k, {}).get('obs_norm', 0.0) for k in range(n_show)]
pred_vals = list(Z_vals[:n_show])

b1 = ax.bar(bar_x - bar_w/2, obs_vals, bar_w, label='Observed (norm)', color='steelblue')
b2 = ax.bar(bar_x + bar_w/2, pred_vals, bar_w, label='Z^1.7 Predicted', color='darkorange')

ax.set_xticks(bar_x)
elements = [sl_stats.get(k, {}).get('element', f'SL{k}') for k in range(n_show)]
ax.set_xticklabels(elements)
ax.set_ylabel('Normalised intensity')
ax.set_title('Per-sublattice: Observed vs Z^1.7 Predicted', fontsize=10)
ax.legend(fontsize=8)

# Panel 6: Displacement magnitude map
ax = axes[1, 2]
ax.imshow(img_norm, cmap='gray', origin='upper', alpha=0.6)
if clean_col_mask.sum() > 0 and disp_mag_nm.max() > 0:
    sc = ax.scatter(cx_clean, cy_clean, c=disp_mag_nm, cmap='hot',
                    s=8, vmin=0, vmax=max(disp_p95, 0.01))
    plt.colorbar(sc, ax=ax, label='|disp| (nm)')
    # Mark large-displacement columns
    if n_large_disp > 0:
        ax.scatter(cx_clean[large_disp_mask], cy_clean[large_disp_mask],
                   s=20, facecolors='none', edgecolors='cyan', linewidth=1.0,
                   label=f'|d|>0.12nm (n={n_large_disp})')
        ax.legend(fontsize=7)
ax.set_title('Displacement Field (clean cols)\nCyan = |d|>0.12 nm', fontsize=10)
ax.set_xlabel('x (px)'); ax.set_ylabel('y (px)')

plt.suptitle('YBCO HAADF-STEM Analysis', fontsize=13, fontweight='bold')
plt.tight_layout()
plt.savefig('analysis_visualization.png', dpi=100, bbox_inches='tight')
plt.close()
print('Saved visualization: analysis_visualization.png')

# ─────────────────────────────────────────────
# 8. Compile results
# ─────────────────────────────────────────────
dark_band_info = []
for dr in dark_regions:
    dark_band_info.append({
        'row_start': dr['row_start'],
        'row_end': dr['row_end'],
        'width_nm': round(dr.get('width_nm', dr['width_px']*nm_per_px), 4),
        'center_nm': round(dr.get('center_nm', 0.0), 4)
    })

sl_stats_out = {}
for k in range(min(n_sl, 4)):
    v = sl_stats.get(k, {})
    sl_stats_out[f'SL{k}_{v.get("element","?")}'] = {
        'mean': round(v.get('mean', 0.0), 5),
        'std':  round(v.get('std',  0.0), 5),
        'count': v.get('count', 0),
        'obs_norm': round(v.get('obs_norm', 0.0), 4),
        'Z_pred':   round(v.get('Z_pred', 0.0), 4),
        'residual': round(v.get('residual', 0.0), 4)
    }

results = {
    'analysis_type': 'YBCO HAADF-STEM: sublattice intensity statistics, 2D spatial map, dark-band characterisation, strain/displacement analysis',
    'extracted_features': {
        'total_columns_detected': int(len(col_x)),
        'num_sublattices': int(n_sl),
        'per_sublattice_stats': sl_stats_out,
        'obs_intensity_ratios_SL0_SL1_SL2_SL3': [round(float(v),4) for v in obs_ratio_matrix],
        'pred_Z17_ratios_Ba_Y_Cu_O': [round(float(v),4) for v in pred_ratio_matrix],
        'residual_error_obs_minus_pred': [round(float(v),4) for v in residual_matrix],
        'intensity_grid_cell_size_nm': round(cell_size_nm, 4),
        'num_anomalous_grid_cells': num_anomalous_cells,
        'anomalous_fraction_of_grid': round(num_anomalous_cells/(grid_n*grid_n), 4),
        'dark_band_count': len(dark_regions),
        'dark_band_details_nm': dark_band_info,
        'inter_band_spacings_nm': [round(v,4) for v in inter_band_spacings_nm],
        'dominant_fourier_period_nm': round(dominant_period_nm, 4),
        'dominant_fourier_freq_per_nm': round(dominant_freq, 5),
        'best_ybco_spacing_match': best_match_key,
        'best_ybco_spacing_nm': best_match_val,
        'period_ybco_correlation_proxy': round(period_corr, 4),
        'nm_per_pixel': round(nm_per_px, 5),
        'lattice_constant_px': round(float(lat_const_px), 3),
        'lattice_constant_nm': round(float(lat_const_px) * nm_per_px, 4),
        'displacement_mean_nm': round(disp_mean, 5),
        'displacement_std_nm':  round(disp_std,  5),
        'displacement_p95_nm':  round(disp_p95,  5),
        'n_clean_columns': int(clean_col_mask.sum()),
        'n_large_displacement_cols_gt_0p12nm': n_large_disp,
        'large_disp_fraction_of_clean_cols': round(n_large_disp / max(1, int(clean_col_mask.sum())), 4),
        'large_disp_centroid_x_px': round(ld_centroid_x, 2),
        'large_disp_centroid_y_px': round(ld_centroid_y, 2),
        'dark_band_area_fraction': round(dark_fraction, 4),
        'ordered_crystal_area_fraction': round(1.0 - dark_fraction, 4)
    },
    'quality_metrics': {
        'intensity_range_raw': [float(img_min), float(img_max)],
        'column_density_per_nm2': round(len(col_x) / (H * W * nm_per_px**2), 4),
        'grid_intensity_mean': round(float(grid_mean), 5),
        'grid_intensity_std':  round(float(grid_std),  5),
        'max_residual_Z17': round(float(max(abs(v) for v in residual_matrix)), 4),
        'otsu_threshold_dark': round(float(otsu_val), 4)
    },
    'summary': (
        f'Detected {len(col_x)} atomic columns in YBCO HAADF image; '
        f'{len(dark_regions)} dark band(s) with dominant periodicity {dominant_period_nm:.3f} nm '
        f'(best match: YBCO {best_match_key}={best_match_val} nm); '
        f'Z^1.7 residuals suggest {'reasonable' if max(abs(v) for v in residual_matrix)<0.2 else 'poor'} '
        f'chemical assignment; {n_large_disp} clean columns show |displacement|>0.12 nm. '
        f'Scale estimated at {nm_per_px:.4f} nm/px. '
        f'Note: sublattice assignment uses K-means intensity clustering; '
        f'nm/px scale inferred from Ba-Ba nearest-neighbour reference (0.385 nm).'
    ),
    'saved_arrays': {
        'column_positions.npy': {
            'description': 'Detected atomic columns: columns are [x_px, y_px, norm_intensity, sublattice_id]',
            'shape': [int(len(col_x)), 4],
            'dtype': 'float64'
        },
        'intensity_grid.npy': {
            'description': '32x32 spatial intensity map (linearly interpolated), each cell ~{:.3f} nm'.format(cell_size_nm),
            'shape': [32, 32],
            'dtype': 'float32'
        },
        'row_intensity_profile.npy': {
            'description': 'Row-averaged normalised intensity profile (length=512), smoothed with sigma=2',
            'shape': [int(H)],
            'dtype': 'float32'
        },
        'displacement_field.npy': {
            'description': 'Displacement field for clean (non-dark-band) columns: [x_px, y_px, |disp|_nm]',
            'shape': [int(clean_col_mask.sum()), 3],
            'dtype': 'float32'
        }
    }
}

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