import numpy as np
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.optimize import curve_fit
from sklearn.mixture import GaussianMixture
import json
import warnings
warnings.filterwarnings('ignore')

# Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_151653/results/analysis_YBCO_tem_ImageAnalysis_20260331_152204_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()}]")

# Normalize to 8-bit for CLAHE
img_f = img.astype(np.float64)
img_norm = ((img_f - img_f.min()) / (img_f.max() - img_f.min()) * 255).astype(np.uint8)

# Step 1: CLAHE
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
img_clahe = clahe.apply(img_norm)
img_clahe_f = img_clahe.astype(np.float64)

# Step 2: FFT on original image
fft2 = np.fft.fft2(img_f)
fft_shift = np.fft.fftshift(fft2)
power = np.abs(fft_shift) ** 2
log_power = np.log1p(power)

H, W = img.shape
cy, cx = H // 2, W // 2

# Find Bragg peaks in FFT
# Mask center DC component
power_search = log_power.copy()
yy, xx = np.mgrid[:H, :W]
dist_from_center = np.sqrt((xx - cx)**2 + (yy - cy)**2)
power_search[dist_from_center < 5] = 0

# Find peaks in FFT power spectrum
from skimage.feature import peak_local_max
fft_peaks = peak_local_max(power_search, min_distance=8, num_peaks=40, threshold_abs=np.percentile(power_search[power_search > 0], 95))

print(f"Found {len(fft_peaks)} FFT peaks")

# Convert to frequency coordinates relative to center
peak_coords = []
for p in fft_peaks:
    ky = p[0] - cy
    kx = p[1] - cx
    dist = np.sqrt(kx**2 + ky**2)
    if dist > 3:  # skip DC
        peak_coords.append((kx, ky, dist, p[1], p[0]))

peak_coords.sort(key=lambda x: x[2])

print("Top 20 FFT peaks (kx, ky, dist, real_space_period):")
for i, (kx, ky, dist, px, py) in enumerate(peak_coords[:20]):
    period = max(H, W) / dist if dist > 0 else 0
    print(f"  Peak {i}: kx={kx}, ky={ky}, dist={dist:.1f}, period={period:.2f} px")

# Select fundamental lattice vectors
# Target: a-axis ~11.37 px horizontal (kx~45, ky~0), c-axis ~17.66 px vertical (kx~0, ky~29)
# We look for peaks closest to these expected positions

def find_best_peak(peak_coords, target_kx, target_ky, tolerance=15):
    """Find peak closest to target frequency."""
    best = None
    best_dist = float('inf')
    for kx, ky, dist, px, py in peak_coords:
        d = np.sqrt((kx - target_kx)**2 + (ky - target_ky)**2)
        if d < best_dist and d < tolerance:
            best_dist = d
            best = (kx, ky, dist, px, py)
    return best

# Expected: a-axis -> kx ~ 512/11.37 ~ 45, ky ~ 0
# Expected: c-axis -> kx ~ 0, ky ~ 512/17.66 ~ 29
a_peak = find_best_peak(peak_coords, 45, 0, tolerance=20)
c_peak = find_best_peak(peak_coords, 0, 29, tolerance=20)

# If not found with tight tolerance, try broader search
if a_peak is None:
    # Look for peaks roughly along horizontal axis
    horiz_peaks = [(kx, ky, dist, px, py) for kx, ky, dist, px, py in peak_coords 
                   if abs(ky) < 10 and kx > 10]
    if horiz_peaks:
        horiz_peaks.sort(key=lambda x: x[2])
        a_peak = horiz_peaks[0]
        print(f"Selected horizontal peak via fallback: kx={a_peak[0]}, ky={a_peak[1]}")

if c_peak is None:
    # Look for peaks roughly along vertical axis
    vert_peaks = [(kx, ky, dist, px, py) for kx, ky, dist, px, py in peak_coords 
                  if abs(kx) < 10 and ky > 10]
    if vert_peaks:
        vert_peaks.sort(key=lambda x: x[2])
        c_peak = vert_peaks[0]
        print(f"Selected vertical peak via fallback: kx={c_peak[0]}, ky={c_peak[1]}")

# If still not found, use the two shortest independent peaks closest to orthogonal
if a_peak is None or c_peak is None:
    print("Using general orthogonal pair selection")
    # Select pairs that are most orthogonal
    best_pair = None
    best_score = float('inf')
    for i in range(min(len(peak_coords), 15)):
        for j in range(i+1, min(len(peak_coords), 15)):
            kx1, ky1 = peak_coords[i][0], peak_coords[i][1]
            kx2, ky2 = peak_coords[j][0], peak_coords[j][1]
            # Angle between vectors
            dot = abs(kx1*kx2 + ky1*ky2)
            mag1 = np.sqrt(kx1**2 + ky1**2)
            mag2 = np.sqrt(kx2**2 + ky2**2)
            if mag1 > 0 and mag2 > 0:
                cos_angle = dot / (mag1 * mag2)
                # Want angle close to 90 degrees (cos ~ 0)
                score = cos_angle + 0.01 * (mag1 + mag2)  # prefer shorter vectors too
                if score < best_score:
                    best_score = score
                    best_pair = (i, j)
    if best_pair:
        p1 = peak_coords[best_pair[0]]
        p2 = peak_coords[best_pair[1]]
        # Assign: more horizontal one is a, more vertical is c
        if abs(p1[0]) > abs(p1[1]):
            a_peak, c_peak = p1, p2
        else:
            a_peak, c_peak = p2, p1

# Compute lattice parameters
if a_peak is not None:
    a_period = W / a_peak[2]  # real-space period in pixels
    print(f"a-axis: kx={a_peak[0]}, ky={a_peak[1]}, period={a_period:.2f} px")
else:
    a_period = 11.37
    print("Warning: a-axis peak not found, using default")

if c_peak is not None:
    c_period = H / c_peak[2]
    print(f"c-axis: kx={c_peak[0]}, ky={c_peak[1]}, period={c_period:.2f} px")
else:
    c_period = 17.66
    print("Warning: c-axis peak not found, using default")

# Compute angle between lattice vectors
if a_peak is not None and c_peak is not None:
    dot_prod = a_peak[0]*c_peak[0] + a_peak[1]*c_peak[1]
    mag_a = np.sqrt(a_peak[0]**2 + a_peak[1]**2)
    mag_c = np.sqrt(c_peak[0]**2 + c_peak[1]**2)
    cos_angle = np.clip(dot_prod / (mag_a * mag_c), -1, 1)
    lattice_angle = np.degrees(np.arccos(abs(cos_angle)))
    # We want the angle between real-space vectors (same as reciprocal for orthogonal)
    lattice_angle_real = 180 - lattice_angle if lattice_angle > 90 else lattice_angle
    # Actually for orthogonal lattice, angle in reciprocal = angle in real space
    print(f"Lattice angle: {lattice_angle:.1f} degrees")
else:
    lattice_angle = 90.0

# Step 3: Bandpass filter on CLAHE image
from scipy.ndimage import gaussian_filter
sigma_small = 2
sigma_large = 15
bandpass = gaussian_filter(img_clahe_f, sigma_small) - gaussian_filter(img_clahe_f, sigma_large)

# Step 4: Detect atomic columns as local maxima
min_sep = 7
threshold_val = np.percentile(bandpass[bandpass > 0], 50)  # 50th percentile

peaks = peak_local_max(bandpass, min_distance=min_sep, threshold_abs=threshold_val)
print(f"Detected {len(peaks)} initial peaks")

# Step 5: Sub-pixel refinement via 2D Gaussian fitting
def gaussian_2d(coords, amp, x0, y0, sigma_x, sigma_y, offset):
    x, y = coords
    return (amp * np.exp(-((x - x0)**2 / (2 * sigma_x**2) + (y - y0)**2 / (2 * sigma_y**2))) + offset).ravel()

refined_positions = []
refined_intensities_raw = []
fit_window = 7  # 7x7 window
half_w = fit_window // 2
min_fwhm = 3.0
max_fwhm = 9.0

for peak in peaks:
    py, px = peak
    # Bounds check
    if py - half_w < 0 or py + half_w + 1 > H or px - half_w < 0 or px + half_w + 1 > W:
        continue
    
    window = img_clahe_f[py - half_w:py + half_w + 1, px - half_w:px + half_w + 1]
    
    y_grid, x_grid = np.mgrid[0:fit_window, 0:fit_window]
    
    try:
        amp_init = window.max() - window.min()
        offset_init = window.min()
        p0 = [amp_init, half_w, half_w, 2.0, 2.0, offset_init]
        bounds_lower = [0, -1, -1, 0.5, 0.5, 0]
        bounds_upper = [500, fit_window + 1, fit_window + 1, 6.0, 6.0, 300]
        
        popt, pcov = curve_fit(gaussian_2d, (x_grid, y_grid), window.ravel(),
                               p0=p0, bounds=(bounds_lower, bounds_upper), maxfev=1000)
        
        amp, x0, y0, sx, sy, offset = popt
        
        # FWHM check
        fwhm_x = 2.355 * sx
        fwhm_y = 2.355 * sy
        
        if fwhm_x < min_fwhm or fwhm_x > max_fwhm or fwhm_y < min_fwhm or fwhm_y > max_fwhm:
            continue
        
        # Residual check
        fitted = gaussian_2d((x_grid, y_grid), *popt).reshape(fit_window, fit_window)
        residual = np.std(window - fitted)
        if residual > 2 * np.std(window) * 0.5:  # relaxed slightly
            continue
        
        # Refined position in full image coordinates
        refined_x = px - half_w + x0
        refined_y = py - half_w + y0
        
        refined_positions.append([refined_x, refined_y])
        
    except (RuntimeError, ValueError):
        # If fitting fails, keep original position
        refined_positions.append([float(px), float(py)])

refined_positions = np.array(refined_positions)
print(f"Validated {len(refined_positions)} atomic columns after Gaussian fitting")

# Step 6: Nearest-neighbor distances
from scipy.spatial import cKDTree
tree = cKDTree(refined_positions)
nn_dists = []
for i in range(len(refined_positions)):
    dists, _ = tree.query(refined_positions[i], k=7)  # k=7 to get 6 nearest neighbors
    nn_dists.extend(dists[1:])  # exclude self

nn_dists = np.array(nn_dists)
# First nearest neighbor only
first_nn = []
for i in range(len(refined_positions)):
    dists, _ = tree.query(refined_positions[i], k=2)
    first_nn.append(dists[1])
first_nn = np.array(first_nn)

print(f"Nearest-neighbor distances: mean={first_nn.mean():.2f}, std={first_nn.std():.2f}")

# Step 7: Background-corrected intensities from original image
def gaussian_weight(size=3):
    """Create Gaussian weight kernel."""
    x = np.arange(size) - size // 2
    y = np.arange(size) - size // 2
    xx, yy = np.meshgrid(x, y)
    w = np.exp(-(xx**2 + yy**2) / (2 * 1.0**2))
    return w / w.sum()

gauss_kernel = gaussian_weight(3)
bg_corrected_intensities = []
raw_intensities = []

for pos in refined_positions:
    px, py = pos
    ix, iy = int(round(px)), int(round(py))
    
    # Gaussian-weighted intensity from original image
    half = 1  # 3x3 window
    if iy - half < 0 or iy + half + 1 > H or ix - half < 0 or ix + half + 1 > W:
        bg_corrected_intensities.append(0)
        raw_intensities.append(0)
        continue
    
    patch = img_f[iy - half:iy + half + 1, ix - half:ix + half + 1]
    weighted_intensity = np.sum(patch * gauss_kernel)
    raw_intensities.append(weighted_intensity)
    
    # Background: annular ring (inner=4, outer=7)
    inner_r, outer_r = 4, 7
    bg_vals = []
    for dy in range(-outer_r, outer_r + 1):
        for dx in range(-outer_r, outer_r + 1):
            r = np.sqrt(dx**2 + dy**2)
            if inner_r <= r <= outer_r:
                ny, nx = iy + dy, ix + dx
                if 0 <= ny < H and 0 <= nx < W:
                    bg_vals.append(img_f[ny, nx])
    
    if len(bg_vals) > 0:
        bg = np.mean(bg_vals)
    else:
        bg = 0
    
    bg_corrected_intensities.append(weighted_intensity - bg)

bg_corrected_intensities = np.array(bg_corrected_intensities)
raw_intensities = np.array(raw_intensities)

print(f"Background-corrected intensities: mean={bg_corrected_intensities.mean():.1f}, std={bg_corrected_intensities.std():.1f}")

# Step 8: GMM clustering
valid_mask = bg_corrected_intensities != 0
valid_intensities = bg_corrected_intensities[valid_mask].reshape(-1, 1)

# Try 2-4 components, select best BIC
best_bic = float('inf')
best_gmm = None
best_n = 2

for n_comp in [2, 3, 4]:
    gmm = GaussianMixture(n_components=n_comp, random_state=42, max_iter=200)
    gmm.fit(valid_intensities)
    bic = gmm.bic(valid_intensities)
    print(f"GMM n={n_comp}: BIC={bic:.1f}")
    if bic < best_bic:
        best_bic = bic
        best_gmm = gmm
        best_n = n_comp

print(f"Selected GMM with {best_n} components")

# Assign labels
labels = np.full(len(refined_positions), -1, dtype=int)
labels[valid_mask] = best_gmm.predict(valid_intensities)

# Sort components by mean intensity
comp_means = best_gmm.means_.flatten()
sorted_idx = np.argsort(comp_means)
label_remap = {old: new for new, old in enumerate(sorted_idx)}
labels_sorted = np.array([label_remap.get(l, l) for l in labels])

print(f"GMM component means (sorted): {comp_means[sorted_idx]}")

# Step 9: Detection coverage - quadrant analysis
quadrant_counts = np.zeros(4)
for pos in refined_positions:
    qx = 0 if pos[0] < W / 2 else 1
    qy = 0 if pos[1] < H / 2 else 1
    qi = qy * 2 + qx
    quadrant_counts[qi] += 1

print(f"Quadrant counts: TL={quadrant_counts[0]}, TR={quadrant_counts[1]}, BL={quadrant_counts[2]}, BR={quadrant_counts[3]}")

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

# 1. FFT power spectrum with lattice vectors
ax = axes[0, 0]
ax.imshow(log_power, cmap='gray', extent=[-cx, cx, -cy, cy], origin='lower', aspect='auto')
ax.set_xlim(-80, 80)
ax.set_ylim(-80, 80)
if a_peak is not None:
    ax.plot(a_peak[0], -a_peak[1], 'ro', markersize=10, label=f'a-axis ({a_period:.1f} px)')
    ax.plot(-a_peak[0], a_peak[1], 'ro', markersize=10)
if c_peak is not None:
    ax.plot(c_peak[0], -c_peak[1], 'bs', markersize=10, label=f'c-axis ({c_period:.1f} px)')
    ax.plot(-c_peak[0], c_peak[1], 'bs', markersize=10)
ax.legend(fontsize=8)
ax.set_title('FFT Power Spectrum')
ax.set_xlabel('kx')
ax.set_ylabel('ky')

# 2. Detected columns on original image, color-coded by sublattice
ax = axes[0, 1]
ax.imshow(img, cmap='gray')
colors = ['cyan', 'yellow', 'red', 'green']
for cl in range(best_n):
    mask_cl = labels_sorted == cl
    if mask_cl.any():
        ax.scatter(refined_positions[mask_cl, 0], refined_positions[mask_cl, 1],
                   c=colors[cl % len(colors)], s=8, alpha=0.7,
                   label=f'Sublattice {cl} (n={mask_cl.sum()})')
ax.legend(fontsize=7, loc='upper right')
ax.set_title(f'Detected Columns (n={len(refined_positions)})')
ax.set_xlim(0, W)
ax.set_ylim(H, 0)

# 3. Nearest-neighbor distance histogram
ax = axes[0, 2]
ax.hist(first_nn, bins=50, color='steelblue', edgecolor='black', alpha=0.7)
ax.axvline(a_period, color='red', linestyle='--', label=f'a={a_period:.1f} px')
ax.axvline(c_period, color='blue', linestyle='--', label=f'c={c_period:.1f} px')
ax.set_xlabel('Distance (pixels)')
ax.set_ylabel('Count')
ax.set_title('Nearest-Neighbor Distances')
ax.legend(fontsize=8)

# 4. Background-corrected intensity histogram with GMM
ax = axes[1, 0]
valid_int = bg_corrected_intensities[valid_mask]
ax.hist(valid_int, bins=60, density=True, color='gray', alpha=0.5, edgecolor='black')
x_range = np.linspace(valid_int.min(), valid_int.max(), 300)
total_pdf = np.zeros_like(x_range)
for i in range(best_n):
    weight = best_gmm.weights_[i]
    mean = best_gmm.means_[i, 0]
    std = np.sqrt(best_gmm.covariances_[i, 0, 0])
    component_pdf = weight * (1 / (std * np.sqrt(2 * np.pi))) * np.exp(-0.5 * ((x_range - mean) / std)**2)
    # Map to sorted label
    sorted_label = label_remap[i]
    ax.plot(x_range, component_pdf, color=colors[sorted_label % len(colors)],
            linewidth=2, label=f'Component {sorted_label} (μ={mean:.0f})')
    total_pdf += component_pdf
ax.plot(x_range, total_pdf, 'k--', linewidth=1.5, label='Total')
ax.set_xlabel('Background-corrected Intensity')
ax.set_ylabel('Density')
ax.set_title('Intensity Distribution + GMM')
ax.legend(fontsize=7)

# 5. Spatial sublattice map
ax = axes[1, 1]
ax.imshow(img, cmap='gray', alpha=0.5)
for cl in range(best_n):
    mask_cl = labels_sorted == cl
    if mask_cl.any():
        ax.scatter(refined_positions[mask_cl, 0], refined_positions[mask_cl, 1],
                   c=colors[cl % len(colors)], s=12, alpha=0.8, edgecolors='none')
ax.set_title('Sublattice Spatial Map')
ax.set_xlim(0, W)
ax.set_ylim(H, 0)

# 6. Zoomed view (central 100x100 region)
ax = axes[1, 2]
zoom_size = 100
zx1, zy1 = W // 2 - zoom_size // 2, H // 2 - zoom_size // 2
zx2, zy2 = zx1 + zoom_size, zy1 + zoom_size
ax.imshow(img[zy1:zy2, zx1:zx2], cmap='gray', extent=[zx1, zx2, zy2, zy1])
mask_zoom = ((refined_positions[:, 0] >= zx1) & (refined_positions[:, 0] < zx2) &
             (refined_positions[:, 1] >= zy1) & (refined_positions[:, 1] < zy2))
for cl in range(best_n):
    mask_cl = (labels_sorted == cl) & mask_zoom
    if mask_cl.any():
        ax.scatter(refined_positions[mask_cl, 0], refined_positions[mask_cl, 1],
                   c=colors[cl % len(colors)], s=30, alpha=0.9, edgecolors='white', linewidths=0.5,
                   label=f'Sub {cl}')
ax.legend(fontsize=7)
ax.set_title('Zoomed Center (100x100)')

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

# Save output arrays
# Positions array: [N, 2] with (x, y)
np.save('analysis_positions.npy', refined_positions)

# Labels array
np.save('analysis_labels.npy', labels_sorted)

# Background-corrected intensities
np.save('analysis_intensities.npy', bg_corrected_intensities)

print("Saved .npy files")

# Build results JSON
results = {
    "analysis_type": "FFT-based lattice characterization and atomic column detection with sublattice differentiation for YBCO STEM image",
    "extracted_features": {
        "a_axis_period_px": round(float(a_period), 2),
        "c_axis_period_px": round(float(c_period), 2),
        "lattice_angle_deg": round(float(lattice_angle), 2),
        "a_axis_bragg_peak_kx_ky": [int(a_peak[0]), int(a_peak[1])] if a_peak else None,
        "c_axis_bragg_peak_kx_ky": [int(c_peak[0]), int(c_peak[1])] if c_peak else None,
        "total_detected_columns": int(len(peaks)),
        "validated_columns_after_gaussian_fit": int(len(refined_positions)),
        "nn_distance_mean_px": round(float(first_nn.mean()), 2),
        "nn_distance_std_px": round(float(first_nn.std()), 2),
        "nn_distance_median_px": round(float(np.median(first_nn)), 2),
        "bg_corrected_intensity_mean": round(float(bg_corrected_intensities[valid_mask].mean()), 1),
        "bg_corrected_intensity_std": round(float(bg_corrected_intensities[valid_mask].std()), 1),
        "gmm_n_components": int(best_n),
        "gmm_component_means": [round(float(m), 1) for m in comp_means[sorted_idx]],
        "gmm_component_weights": [round(float(best_gmm.weights_[sorted_idx[i]]), 3) for i in range(best_n)],
        "sublattice_counts": {f"sublattice_{i}": int((labels_sorted == i).sum()) for i in range(best_n)},
        "quadrant_counts": {
            "top_left": int(quadrant_counts[0]),
            "top_right": int(quadrant_counts[1]),
            "bottom_left": int(quadrant_counts[2]),
            "bottom_right": int(quadrant_counts[3])
        },
        "detection_coverage_uniformity_cv": round(float(np.std(quadrant_counts) / np.mean(quadrant_counts)), 3) if np.mean(quadrant_counts) > 0 else None
    },
    "quality_metrics": {
        "detection_rate_per_sq_px": round(float(len(refined_positions) / (H * W)), 6),
        "gmm_bic": round(float(best_bic), 1),
        "fraction_peaks_validated": round(float(len(refined_positions) / max(len(peaks), 1)), 3),
        "bandpass_sigma_small": sigma_small,
        "bandpass_sigma_large": sigma_large,
        "detection_threshold_percentile": 50,
        "min_separation_px": min_sep,
        "fit_window_size": fit_window,
        "fwhm_range_px": [min_fwhm, max_fwhm]
    },
    "summary": f"Detected {len(refined_positions)} atomic columns in YBCO STEM image with a-axis period {a_period:.2f} px and c-axis period {c_period:.2f} px (angle {lattice_angle:.1f} deg). GMM clustering with {best_n} components identifies sublattices with mean intensities {[round(float(m),0) for m in comp_means[sorted_idx]]}. Detection coverage coefficient of variation across quadrants: {np.std(quadrant_counts)/np.mean(quadrant_counts):.3f}.",
    "saved_arrays": {
        "analysis_positions.npy": {
            "description": f"Refined atomic column positions as (x, y) pixel coordinates, {len(refined_positions)} columns",
            "shape": list(refined_positions.shape),
            "dtype": str(refined_positions.dtype)
        },
        "analysis_labels.npy": {
            "description": f"Sublattice labels (0 to {best_n-1}) for each detected column from GMM clustering of background-corrected intensities",
            "shape": list(labels_sorted.shape),
            "dtype": str(labels_sorted.dtype)
        },
        "analysis_intensities.npy": {
            "description": "Background-corrected Gaussian-weighted intensities for each detected atomic column",
            "shape": list(bg_corrected_intensities.shape),
            "dtype": str(bg_corrected_intensities.dtype)
        }
    }
}

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