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

# Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_105905/results/analysis_YBCO_tem_ImageAnalysis_20260331_110209_001/image_0000/temp_image_0.npy'
img = np.load(path)
print(f"Image shape: {img.shape}, dtype: {img.dtype}, range: [{img.min()}, {img.max()}]")

# Step 1: Normalize to float [0,1]
img_float = (img.astype(np.float64) - img.min()) / (img.max() - img.min())

# Pixel size calibration (given in plan)
pixel_size_nm = 0.0352  # nm/pixel

# Step 2: Compute 2D FFT
fft2 = np.fft.fft2(img_float)
fft2_shifted = np.fft.fftshift(fft2)
power_spectrum = np.abs(fft2_shifted)**2
log_power = np.log1p(power_spectrum)

ny, nx = img_float.shape
center_y, center_x = ny // 2, nx // 2

# Create frequency coordinates
fy = np.fft.fftshift(np.fft.fftfreq(ny))  # in cycles/pixel
fx = np.fft.fftshift(np.fft.fftfreq(nx))
FX, FY = np.meshgrid(fx, fy)
freq_radius = np.sqrt(FX**2 + FY**2)

# Find FFT peaks (exclude DC)
# Mask out center (DC component and very low frequencies)
log_power_peaks = log_power.copy()
dc_mask_radius = 5  # pixels in frequency space
for iy in range(ny):
    for ix in range(nx):
        if (iy - center_y)**2 + (ix - center_x)**2 < dc_mask_radius**2:
            log_power_peaks[iy, ix] = 0

# Find peaks in FFT
from skimage.feature import peak_local_max
fft_peak_coords = peak_local_max(log_power_peaks, min_distance=8, threshold_rel=0.3, num_peaks=20)

# Compute spatial periodicities for each FFT peak
fft_peaks_info = []
for peak in fft_peak_coords:
    py, px = peak
    dy = py - center_y
    dx = px - center_x
    dist_pixels = np.sqrt(dy**2 + dx**2)
    if dist_pixels < 3:
        continue
    freq_cyc_per_pixel = dist_pixels / ny  # approximate (assuming square)
    spatial_period_pixels = 1.0 / freq_cyc_per_pixel if freq_cyc_per_pixel > 0 else np.inf
    spatial_period_nm = spatial_period_pixels * pixel_size_nm
    angle_deg = np.degrees(np.arctan2(dy, dx))
    fft_peaks_info.append({
        'pixel_y': int(py), 'pixel_x': int(px),
        'freq_dist_pixels': float(dist_pixels),
        'spatial_period_pixels': float(spatial_period_pixels),
        'spatial_period_nm': float(spatial_period_nm),
        'angle_deg': float(angle_deg)
    })

fft_peaks_info.sort(key=lambda p: p['freq_dist_pixels'])
print(f"Found {len(fft_peaks_info)} FFT peaks")
for i, p in enumerate(fft_peaks_info[:10]):
    print(f"  Peak {i}: period={p['spatial_period_nm']:.3f} nm, angle={p['angle_deg']:.1f} deg")

# Extract lattice parameters from FFT peaks
# Group peaks by spatial period (within tolerance)
def cluster_periods(peaks, tol_nm=0.03):
    if not peaks:
        return []
    periods = [p['spatial_period_nm'] for p in peaks]
    clusters = []
    used = set()
    for i, pi in enumerate(periods):
        if i in used:
            continue
        cluster = [peaks[i]]
        used.add(i)
        for j, pj in enumerate(periods):
            if j in used:
                continue
            if abs(pi - pj) < tol_nm:
                cluster.append(peaks[j])
                used.add(j)
        clusters.append(cluster)
    return clusters

period_clusters = cluster_periods(fft_peaks_info, tol_nm=0.04)
lattice_spacings = []
for cluster in period_clusters:
    avg_period = np.mean([p['spatial_period_nm'] for p in cluster])
    angles = [p['angle_deg'] for p in cluster]
    lattice_spacings.append({
        'period_nm': float(avg_period),
        'count': len(cluster),
        'angles': [float(a) for a in angles]
    })
lattice_spacings.sort(key=lambda x: x['period_nm'])
print(f"\nLattice spacings found: {len(lattice_spacings)}")
for ls in lattice_spacings[:6]:
    print(f"  {ls['period_nm']:.3f} nm ({ls['count']} peaks)")

# Step 3: Bandpass filter in Fourier space
# Low-frequency cutoff: remove background gradient (spatial periods > 3.5 nm)
# High-frequency cutoff: suppress noise (spatial periods < 0.3 nm)
low_freq_cutoff_nm = 3.5  # max spatial period to keep
high_freq_cutoff_nm = 0.15  # min spatial period to keep (adjusted from 0.3 to be more inclusive)

# Convert to cycles/pixel
low_freq_cutoff_cpp = pixel_size_nm / low_freq_cutoff_nm  # small freq value
high_freq_cutoff_cpp = pixel_size_nm / high_freq_cutoff_nm  # large freq value

print(f"\nBandpass: low cutoff = {low_freq_cutoff_cpp:.4f} cyc/px (>{low_freq_cutoff_nm} nm removed)")
print(f"Bandpass: high cutoff = {high_freq_cutoff_cpp:.4f} cyc/px (<{high_freq_cutoff_nm} nm removed)")

# Create bandpass filter (smooth edges with Gaussian taper)
bandpass = np.ones_like(freq_radius)
sigma_low = low_freq_cutoff_cpp * 0.3
sigma_high = high_freq_cutoff_cpp * 0.1

# High-pass (remove low frequencies)
bandpass *= 1.0 - np.exp(-0.5 * (freq_radius / low_freq_cutoff_cpp)**4)
# Low-pass (remove high frequencies)
bandpass *= np.exp(-0.5 * (freq_radius / high_freq_cutoff_cpp)**4)

# Apply bandpass
fft2_filtered = fft2_shifted * bandpass
img_filtered = np.real(np.fft.ifft2(np.fft.ifftshift(fft2_filtered)))
img_filtered = (img_filtered - img_filtered.min()) / (img_filtered.max() - img_filtered.min())

print(f"Filtered image range: [{img_filtered.min():.4f}, {img_filtered.max():.4f}]")

# Step 4: Detect atomic columns via local maximum finding
# Expected nearest-neighbor ~0.38 nm -> ~10.8 pixels, use min_distance=8
min_dist = 8

# Use peak_local_max on bandpass-filtered image
coordinates = peak_local_max(
    img_filtered, 
    min_distance=min_dist,
    threshold_abs=0.25,  # threshold on normalized filtered image
    exclude_border=10
)

print(f"\nInitial peak detection: {len(coordinates)} columns found")

# Step 5: Refine each peak position with 2D Gaussian fitting
def fit_2d_gaussian(img_patch, x0, y0):
    """Fit 2D Gaussian to a small patch. Returns refined (x, y, amplitude, sigma_x, sigma_y, offset)."""
    sy, sx = img_patch.shape
    Y, X = np.mgrid[0:sy, 0:sx]
    
    def gaussian_2d(coords, amp, xc, yc, sx, sy, offset):
        x, y = coords
        return (amp * np.exp(-0.5 * ((x - xc)**2 / sx**2 + (y - yc)**2 / sy**2)) + offset).ravel()
    
    try:
        initial_guess = [
            img_patch.max() - img_patch.min(),
            sx / 2.0, sy / 2.0,
            1.5, 1.5,
            img_patch.min()
        ]
        bounds_lower = [0, 0, 0, 0.5, 0.5, -np.inf]
        bounds_upper = [np.inf, sx, sy, sx/2, sy/2, np.inf]
        
        popt, pcov = optimize.curve_fit(
            gaussian_2d, (X, Y), img_patch.ravel(),
            p0=initial_guess, bounds=(bounds_lower, bounds_upper),
            maxfev=1000
        )
        return popt  # amp, xc, yc, sx, sy, offset
    except:
        return None

half_window = 3  # 7x7 window
refined_positions = []
fit_amplitudes = []
fit_sigmas = []
fit_offsets = []

for coord in coordinates:
    cy, cx = coord
    # Extract window from original image
    y_lo = max(0, cy - half_window)
    y_hi = min(img_float.shape[0], cy + half_window + 1)
    x_lo = max(0, cx - half_window)
    x_hi = min(img_float.shape[1], cx + half_window + 1)
    
    patch = img_float[y_lo:y_hi, x_lo:x_hi]
    
    if patch.shape[0] < 5 or patch.shape[1] < 5:
        # Window too small near edge, use raw position
        refined_positions.append([float(cx), float(cy)])
        fit_amplitudes.append(float(img_float[cy, cx]))
        fit_sigmas.append([1.0, 1.0])
        fit_offsets.append(0.0)
        continue
    
    result = fit_2d_gaussian(patch, cx - x_lo, cy - y_lo)
    
    if result is not None:
        amp, xc_local, yc_local, sx_fit, sy_fit, offset = result
        # Convert back to global coordinates
        xc_global = x_lo + xc_local
        yc_global = y_lo + yc_local
        
        # Sanity check: refined position should be close to original
        if abs(xc_global - cx) < half_window and abs(yc_global - cy) < half_window:
            refined_positions.append([float(xc_global), float(yc_global)])
            fit_amplitudes.append(float(amp))
            fit_sigmas.append([float(sx_fit), float(sy_fit)])
            fit_offsets.append(float(offset))
        else:
            refined_positions.append([float(cx), float(cy)])
            fit_amplitudes.append(float(img_float[cy, cx]))
            fit_sigmas.append([1.0, 1.0])
            fit_offsets.append(0.0)
    else:
        refined_positions.append([float(cx), float(cy)])
        fit_amplitudes.append(float(img_float[cy, cx]))
        fit_sigmas.append([1.0, 1.0])
        fit_offsets.append(0.0)

refined_positions = np.array(refined_positions)  # shape (N, 2) = (x, y)
fit_amplitudes = np.array(fit_amplitudes)
print(f"Refined {len(refined_positions)} atomic column positions")

# Step 6: Extract nearest-neighbor distances
from scipy.spatial import KDTree

tree = KDTree(refined_positions)
# Find k nearest neighbors (k=7 to get 6 neighbors + self)
k_neighbors = min(7, len(refined_positions))
dists, indices = tree.query(refined_positions, k=k_neighbors)

# Nearest neighbor distance for each column (exclude self at index 0)
nn_distances_pixels = dists[:, 1]  # distance to nearest neighbor
nn_distances_nm = nn_distances_pixels * pixel_size_nm

print(f"\nNearest-neighbor distances (nm):")
print(f"  Mean: {np.mean(nn_distances_nm):.3f}")
print(f"  Std: {np.std(nn_distances_nm):.3f}")
print(f"  Min: {np.min(nn_distances_nm):.3f}")
print(f"  Max: {np.max(nn_distances_nm):.3f}")

# All pairwise nearest neighbor distances (1st through 4th)
all_nn_dists = dists[:, 1:min(5, k_neighbors)] * pixel_size_nm
nn_hist_values, nn_hist_edges = np.histogram(all_nn_dists.ravel(), bins=50)

# Step 7: Measure integrated intensity at each column position
# Use a small circular aperture around each refined position
aperture_radius = 3  # pixels
integrated_intensities = []
for pos in refined_positions:
    cx, cy = pos
    ix, iy = int(round(cx)), int(round(cy))
    y_lo = max(0, iy - aperture_radius)
    y_hi = min(img_float.shape[0], iy + aperture_radius + 1)
    x_lo = max(0, ix - aperture_radius)
    x_hi = min(img_float.shape[1], ix + aperture_radius + 1)
    
    patch = img_float[y_lo:y_hi, x_lo:x_hi]
    # Create circular mask
    Y, X = np.mgrid[y_lo:y_hi, x_lo:x_hi]
    mask = ((X - cx)**2 + (Y - cy)**2) <= aperture_radius**2
    integrated_intensities.append(float(np.sum(patch[mask])))

integrated_intensities = np.array(integrated_intensities)

print(f"\nIntegrated intensity statistics:")
print(f"  Mean: {np.mean(integrated_intensities):.4f}")
print(f"  Std: {np.std(integrated_intensities):.4f}")

# Classify columns into heavy vs light based on intensity
# Use 2-component analysis
from sklearn.cluster import KMeans

if len(integrated_intensities) > 10:
    km = KMeans(n_clusters=2, random_state=42, n_init=10)
    intensity_labels = km.fit_predict(integrated_intensities.reshape(-1, 1))
    
    # Make label 0 = lighter, label 1 = heavier
    if km.cluster_centers_[0] > km.cluster_centers_[1]:
        intensity_labels = 1 - intensity_labels
    
    n_light = np.sum(intensity_labels == 0)
    n_heavy = np.sum(intensity_labels == 1)
    mean_light = np.mean(integrated_intensities[intensity_labels == 0])
    mean_heavy = np.mean(integrated_intensities[intensity_labels == 1])
    print(f"\nColumn classification:")
    print(f"  Light columns: {n_light} (mean intensity: {mean_light:.4f})")
    print(f"  Heavy columns: {n_heavy} (mean intensity: {mean_heavy:.4f})")
else:
    intensity_labels = np.zeros(len(integrated_intensities), dtype=int)
    n_light = len(integrated_intensities)
    n_heavy = 0
    mean_light = float(np.mean(integrated_intensities))
    mean_heavy = 0.0

# Positions in nm
positions_nm = refined_positions * pixel_size_nm

# Save arrays
np.save("analysis_positions_pixels.npy", refined_positions)
np.save("analysis_positions_nm.npy", positions_nm)
np.save("analysis_intensities.npy", integrated_intensities)
np.save("analysis_intensity_labels.npy", intensity_labels)
np.save("analysis_fft_power.npy", log_power)
np.save("analysis_filtered_image.npy", img_filtered)

# Create visualization
fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)

# 1. Original image
ax = axes[0, 0]
ax.imshow(img_float, cmap='gray')
ax.set_title('Original Image (normalized)')
ax.set_xlabel('pixels')
ax.set_ylabel('pixels')

# 2. FFT power spectrum with labeled peaks
ax = axes[0, 1]
vmax = np.percentile(log_power, 99.5)
ax.imshow(log_power, cmap='inferno', vmax=vmax)
for p in fft_peaks_info[:12]:
    ax.plot(p['pixel_x'], p['pixel_y'], 'c+', markersize=8, markeredgewidth=1.5)
    ax.annotate(f"{p['spatial_period_nm']:.2f}nm", 
                (p['pixel_x'], p['pixel_y']), 
                color='cyan', fontsize=6, 
                xytext=(3, 3), textcoords='offset points')
ax.set_title('FFT Power Spectrum')
ax.set_xlim(center_x - 80, center_x + 80)
ax.set_ylim(center_y + 80, center_y - 80)

# 3. Bandpass filtered image
ax = axes[0, 2]
ax.imshow(img_filtered, cmap='gray')
ax.set_title('Bandpass Filtered Image')

# 4. Detected atomic columns on original image
ax = axes[1, 0]
ax.imshow(img_float, cmap='gray')
colors = ['blue' if l == 0 else 'red' for l in intensity_labels]
ax.scatter(refined_positions[:, 0], refined_positions[:, 1], 
           c=colors, s=8, alpha=0.7, linewidths=0)
ax.set_title(f'Detected Columns (N={len(refined_positions)})\nBlue=light, Red=heavy')
ax.set_xlabel('pixels')
ax.set_ylabel('pixels')

# 5. Nearest-neighbor distance distribution
ax = axes[1, 1]
ax.hist(nn_distances_nm, bins=40, color='steelblue', edgecolor='black', alpha=0.8)
ax.axvline(np.mean(nn_distances_nm), color='red', linestyle='--', label=f'Mean={np.mean(nn_distances_nm):.3f} nm')
ax.set_xlabel('Nearest-neighbor distance (nm)')
ax.set_ylabel('Count')
ax.set_title('NN Distance Distribution')
ax.legend()

# 6. Column intensity distribution
ax = axes[1, 2]
ax.hist(integrated_intensities[intensity_labels == 0], bins=30, alpha=0.7, 
        color='blue', label=f'Light (N={n_light})', edgecolor='black')
ax.hist(integrated_intensities[intensity_labels == 1], bins=30, alpha=0.7, 
        color='red', label=f'Heavy (N={n_heavy})', edgecolor='black')
ax.set_xlabel('Integrated Intensity (a.u.)')
ax.set_ylabel('Count')
ax.set_title('Column Intensity Distribution')
ax.legend()

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

# Compile results
# Lattice parameter estimates from FFT
lattice_params_fft = []
for ls in lattice_spacings[:6]:
    lattice_params_fft.append({
        'spacing_nm': round(ls['period_nm'], 4),
        'num_peaks': ls['count'],
        'angles_deg': [round(a, 1) for a in ls['angles'][:4]]
    })

# NN distance histogram peaks
from scipy.signal import find_peaks as scipy_find_peaks
nn_hist_smooth = ndimage.gaussian_filter1d(nn_hist_values.astype(float), sigma=1)
nn_peaks_idx, _ = scipy_find_peaks(nn_hist_smooth, height=max(nn_hist_smooth)*0.1, distance=3)
nn_peak_distances = [round(float((nn_hist_edges[i] + nn_hist_edges[i+1])/2), 4) for i in nn_peaks_idx]

results = {
    "analysis_type": "YBCO STEM atomic column detection and lattice analysis via FFT, bandpass filtering, peak finding, and 2D Gaussian refinement",
    "extracted_features": {
        "pixel_size_nm_per_pixel": pixel_size_nm,
        "image_shape": list(img.shape),
        "total_atomic_columns_detected": int(len(refined_positions)),
        "n_heavy_columns": int(n_heavy),
        "n_light_columns": int(n_light),
        "mean_intensity_heavy": round(float(mean_heavy), 4),
        "mean_intensity_light": round(float(mean_light), 4),
        "nn_distance_mean_nm": round(float(np.mean(nn_distances_nm)), 4),
        "nn_distance_std_nm": round(float(np.std(nn_distances_nm)), 4),
        "nn_distance_min_nm": round(float(np.min(nn_distances_nm)), 4),
        "nn_distance_max_nm": round(float(np.max(nn_distances_nm)), 4),
        "nn_distance_histogram_peaks_nm": nn_peak_distances,
        "fft_lattice_spacings": lattice_params_fft,
        "integrated_intensity_mean": round(float(np.mean(integrated_intensities)), 4),
        "integrated_intensity_std": round(float(np.std(integrated_intensities)), 4),
        "bandpass_low_cutoff_nm": low_freq_cutoff_nm,
        "bandpass_high_cutoff_nm": high_freq_cutoff_nm,
        "peak_detection_min_distance_pixels": min_dist,
        "gaussian_fit_window_size": 2 * half_window + 1
    },
    "quality_metrics": {
        "mean_fit_amplitude": round(float(np.mean(fit_amplitudes)), 4),
        "std_fit_amplitude": round(float(np.std(fit_amplitudes)), 4),
        "columns_per_nm2": round(float(len(refined_positions) / (img.shape[0] * pixel_size_nm * img.shape[1] * pixel_size_nm)), 2),
        "intensity_contrast_ratio": round(float(mean_heavy / mean_light) if mean_light > 0 else 0.0, 3),
        "num_fft_peaks_found": len(fft_peaks_info)
    },
    "summary": f"Detected {len(refined_positions)} atomic columns in YBCO STEM image using FFT-guided bandpass filtering and 2D Gaussian refinement. Columns classified into {n_heavy} heavy (likely Ba/Y) and {n_light} light (likely Cu/O) based on integrated HAADF intensity. Mean nearest-neighbor distance is {np.mean(nn_distances_nm):.3f} nm. FFT analysis reveals lattice periodicities at {', '.join([f'{ls["period_nm"]:.3f}' for ls in lattice_spacings[:4]])} nm. Bandpass low-frequency cutoff set at 3.5 nm and high-frequency at 0.15 nm (adjusted from plan's 0.3 nm to retain finer lattice features). Peak detection threshold_abs=0.25 on normalized filtered image.",
    "saved_arrays": {
        "analysis_positions_pixels.npy": {
            "description": f"Refined atomic column positions in pixels (x, y), {len(refined_positions)} columns",
            "shape": list(refined_positions.shape),
            "dtype": str(refined_positions.dtype)
        },
        "analysis_positions_nm.npy": {
            "description": f"Refined atomic column positions in nm (x, y), {len(positions_nm)} columns",
            "shape": list(positions_nm.shape),
            "dtype": str(positions_nm.dtype)
        },
        "analysis_intensities.npy": {
            "description": f"Integrated intensity for each atomic column ({len(integrated_intensities)} values)",
            "shape": list(integrated_intensities.shape),
            "dtype": str(integrated_intensities.dtype)
        },
        "analysis_intensity_labels.npy": {
            "description": f"Binary labels for each column: 0=light, 1=heavy ({n_light} light, {n_heavy} heavy)",
            "shape": list(intensity_labels.shape),
            "dtype": str(intensity_labels.dtype)
        },
        "analysis_fft_power.npy": {
            "description": "Log power spectrum of 2D FFT (shifted)",
            "shape": list(log_power.shape),
            "dtype": str(log_power.dtype)
        },
        "analysis_filtered_image.npy": {
            "description": "Bandpass filtered image normalized to [0,1]",
            "shape": list(img_filtered.shape),
            "dtype": str(img_filtered.dtype)
        }
    }
}

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