import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import stats, ndimage
from skimage import measure
import os
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# Step 1: Load data
# ============================================================
image_path = '/Users/maxim.ziatdinov/Code/SciLink/test_tier2_handoff/image_0000/temp_image_0.npy'
image = np.load(image_path)

# Try to load Tier 1 outputs
label_path = os.path.join(os.path.dirname(image_path), 'analysis_labels.npy')
features_path = os.path.join(os.path.dirname(image_path), 'grain_features.npy')

labels = None
grain_features_raw = None

if os.path.exists(label_path):
    labels = np.load(label_path)
    print(f"Loaded label map from {label_path}, shape={labels.shape}, unique labels={len(np.unique(labels))}")
else:
    print(f"Label map not found at {label_path}, will run segmentation.")

if os.path.exists(features_path):
    grain_features_raw = np.load(features_path, allow_pickle=True)
    print(f"Loaded grain features from {features_path}")

# ============================================================
# Step 1b: If no label map, run SAM segmentation
# ============================================================
if labels is None:
    from scilink.tools.sam import run_sam_analysis
    # Image is 512x512 uint8 grayscale
    # Estimate grain sizes: typical microstructure grains ~30-150px diameter -> area ~700-17000
    result = run_sam_analysis(image, params={
        "sam_parameters": "default",
        "min_area": 500,
        "max_area": 80000,
        "pruning_iou_threshold": 0.5
    })
    # Build label map
    labels = np.zeros(image.shape[:2], dtype=np.int32)
    for i, p in enumerate(result["particles"]):
        mask = np.array(p["mask"], dtype=bool)
        labels[mask] = i + 1
    print(f"SAM segmentation: {result['total_count']} particles detected")

# ============================================================
# Step 2: Extract per-grain properties using regionprops
# ============================================================
props = measure.regionprops(labels, intensity_image=image)

# Build structured grain data
grain_data = []
for p in props:
    grain_data.append({
        'label': p.label,
        'area': p.area,
        'centroid_x': p.centroid[1],  # column = x
        'centroid_y': p.centroid[0],  # row = y
        'solidity': p.solidity,
        'perimeter': p.perimeter,
        'major_axis': p.major_axis_length,
        'minor_axis': p.minor_axis_length,
        'equivalent_diameter': p.equivalent_diameter,
        'aspect_ratio': p.major_axis_length / max(p.minor_axis_length, 1e-6),
        'mean_intensity': p.mean_intensity,
        'bbox': p.bbox
    })

print(f"Total grains from regionprops: {len(grain_data)}")

# ============================================================
# Step 3: Clean dataset
# ============================================================
MIN_AREA = 2500  # pixels^2
MIN_SOLIDITY = 0.80

excluded_ids = []
cleaned_grains = []
for g in grain_data:
    exclude = False
    if g['area'] < MIN_AREA:
        exclude = True
    if g['solidity'] < MIN_SOLIDITY:
        exclude = True
    if exclude:
        excluded_ids.append(g['label'])
    else:
        cleaned_grains.append(g)

print(f"Cleaned grains: {len(cleaned_grains)} (excluded {len(excluded_ids)})")

# If cleaning is too aggressive (removes too many), relax thresholds
if len(cleaned_grains) < 5:
    print("Warning: Too few grains after cleaning. Relaxing thresholds.")
    MIN_AREA = 500
    MIN_SOLIDITY = 0.60
    excluded_ids = []
    cleaned_grains = []
    for g in grain_data:
        exclude = False
        if g['area'] < MIN_AREA:
            exclude = True
        if g['solidity'] < MIN_SOLIDITY:
            exclude = True
        if exclude:
            excluded_ids.append(g['label'])
        else:
            cleaned_grains.append(g)
    print(f"After relaxation - Cleaned grains: {len(cleaned_grains)} (excluded {len(excluded_ids)})")

# Extract arrays for analysis
areas = np.array([g['area'] for g in cleaned_grains])
centroid_x = np.array([g['centroid_x'] for g in cleaned_grains])
centroid_y = np.array([g['centroid_y'] for g in cleaned_grains])
equiv_diameters = np.array([g['equivalent_diameter'] for g in cleaned_grains])
aspect_ratios = np.array([g['aspect_ratio'] for g in cleaned_grains])
solidities = np.array([g['solidity'] for g in cleaned_grains])

# Also get raw (uncleaned) arrays for comparison
raw_areas = np.array([g['area'] for g in grain_data])
raw_centroid_x = np.array([g['centroid_x'] for g in grain_data])

# ============================================================
# Step 4: Spatial zone analysis
# ============================================================
img_width = image.shape[1]  # 512
zone_boundary_1 = img_width / 3.0  # ~170.67
zone_boundary_2 = 2 * img_width / 3.0  # ~341.33

# Assign zones
zones = []
for cx in centroid_x:
    if cx < zone_boundary_1:
        zones.append('left')
    elif cx < zone_boundary_2:
        zones.append('center')
    else:
        zones.append('right')
zones = np.array(zones)

left_areas = areas[zones == 'left']
center_areas = areas[zones == 'center']
right_areas = areas[zones == 'right']

# Zone-wise statistics
def zone_stats(arr):
    if len(arr) == 0:
        return {'count': 0, 'mean': None, 'median': None, 'iqr': None}
    return {
        'count': int(len(arr)),
        'mean': float(np.mean(arr)),
        'median': float(np.median(arr)),
        'iqr': float(np.percentile(arr, 75) - np.percentile(arr, 25))
    }

left_stats = zone_stats(left_areas)
center_stats = zone_stats(center_areas)
right_stats = zone_stats(right_areas)

# Kruskal-Wallis test
zone_groups = [a for a in [left_areas, center_areas, right_areas] if len(a) >= 1]
if len(zone_groups) >= 2:
    kw_stat, kw_p = stats.kruskal(*zone_groups)
else:
    kw_stat, kw_p = float('nan'), float('nan')

# Pairwise Mann-Whitney U tests with Bonferroni correction
pairwise_results = {}
pairs = [('left', 'center', left_areas, center_areas),
         ('left', 'right', left_areas, right_areas),
         ('center', 'right', center_areas, right_areas)]
for name1, name2, arr1, arr2 in pairs:
    key = f"{name1}_vs_{name2}"
    if len(arr1) >= 1 and len(arr2) >= 1:
        u_stat, u_p = stats.mannwhitneyu(arr1, arr2, alternative='two-sided')
        pairwise_results[key] = float(min(u_p * 3, 1.0))  # Bonferroni correction
    else:
        pairwise_results[key] = float('nan')

# ============================================================
# Step 5: Linear regression on cleaned data with bootstrap CI
# ============================================================
# Raw regression
if len(raw_areas) >= 3:
    raw_slope, raw_intercept, raw_r, raw_p, raw_se = stats.linregress(raw_centroid_x, raw_areas)
    raw_R2 = raw_r ** 2
else:
    raw_slope, raw_intercept, raw_R2 = float('nan'), float('nan'), float('nan')

# Cleaned regression
if len(areas) >= 3:
    clean_slope, clean_intercept, clean_r, clean_p, clean_se = stats.linregress(centroid_x, areas)
    clean_R2 = clean_r ** 2
else:
    clean_slope, clean_intercept, clean_R2 = float('nan'), float('nan'), float('nan')

# Bootstrap CI for slope
n_bootstrap = 1000
bootstrap_slopes = []
if len(areas) >= 3:
    rng = np.random.RandomState(42)
    n = len(areas)
    for _ in range(n_bootstrap):
        idx = rng.choice(n, size=n, replace=True)
        bs_slope, _, _, _, _ = stats.linregress(centroid_x[idx], areas[idx])
        bootstrap_slopes.append(bs_slope)
    bootstrap_slopes = np.array(bootstrap_slopes)
    slope_ci_low = float(np.percentile(bootstrap_slopes, 2.5))
    slope_ci_high = float(np.percentile(bootstrap_slopes, 97.5))
else:
    slope_ci_low, slope_ci_high = float('nan'), float('nan')

# ============================================================
# Step 6: Anomalous grain subpopulation
# ============================================================
ASPECT_RATIO_THRESH = 2.0
SOLIDITY_THRESH = 0.85

anom_mask = (aspect_ratios > ASPECT_RATIO_THRESH) | (solidities < SOLIDITY_THRESH)
anom_grains = [g for g, m in zip(cleaned_grains, anom_mask) if m]
anom_count = int(np.sum(anom_mask))
anom_positions = [(g['centroid_x'], g['centroid_y']) for g in anom_grains]

# Test if anomalous grain centroids differ from overall population
if anom_count >= 2 and len(centroid_x) >= 3:
    anom_cx = np.array([g['centroid_x'] for g in anom_grains])
    pop_mean_cx = np.mean(centroid_x)
    t_stat_anom, p_anom = stats.ttest_1samp(anom_cx, pop_mean_cx)
    p_anom = float(p_anom)
else:
    p_anom = float('nan')

# ============================================================
# Step 7: Grain size distribution fitting
# ============================================================
# Equivalent diameters
ed = equiv_diameters[equiv_diameters > 0]
log_ed = np.log(ed)

# Fit log-normal distribution (MLE)
lognorm_shape, lognorm_loc, lognorm_scale = stats.lognorm.fit(ed, floc=0)
lognorm_mu = np.log(lognorm_scale)
lognorm_sigma = lognorm_shape

# Fit normal distribution (MLE)
norm_mu, norm_sigma = stats.norm.fit(ed)

# Shapiro-Wilk test on log-transformed diameters (test for log-normality)
if len(log_ed) >= 3:
    sw_W, sw_p = stats.shapiro(log_ed)
else:
    sw_W, sw_p = float('nan'), float('nan')

# Anderson-Darling tests
if len(ed) >= 3:
    ad_lognorm = stats.anderson(log_ed, dist='norm')  # log-normal = normal after log transform
    ad_normal = stats.anderson(ed, dist='norm')
    ad_stat_lognorm = float(ad_lognorm.statistic)
    ad_stat_normal = float(ad_normal.statistic)
else:
    ad_stat_lognorm = float('nan')
    ad_stat_normal = float('nan')

# AIC/BIC computation
n_data = len(ed)

# Log-likelihood for log-normal
ll_lognorm = np.sum(stats.lognorm.logpdf(ed, lognorm_shape, loc=0, scale=lognorm_scale))
k_lognorm = 2  # mu, sigma
AIC_lognorm = 2 * k_lognorm - 2 * ll_lognorm
BIC_lognorm = k_lognorm * np.log(n_data) - 2 * ll_lognorm

# Log-likelihood for normal
ll_normal = np.sum(stats.norm.logpdf(ed, loc=norm_mu, scale=norm_sigma))
k_normal = 2  # mu, sigma
AIC_normal = 2 * k_normal - 2 * ll_normal
BIC_normal = k_normal * np.log(n_data) - 2 * ll_normal

# ============================================================
# Step 8: Visualization
# ============================================================
fig, axes = plt.subplots(2, 2, figsize=(16, 14))

# Panel (a): Spatial map of grains colored by area with zone boundaries
ax = axes[0, 0]
ax.imshow(image, cmap='gray', alpha=0.7)

if len(cleaned_grains) > 0:
    sc = ax.scatter(centroid_x, centroid_y, c=areas, cmap='viridis', 
                    s=30, edgecolors='white', linewidth=0.5, alpha=0.8)
    plt.colorbar(sc, ax=ax, label='Grain Area (px²)', shrink=0.7)

# Mark anomalous grains
if anom_count > 0:
    anom_x = [p[0] for p in anom_positions]
    anom_y = [p[1] for p in anom_positions]
    ax.scatter(anom_x, anom_y, s=80, facecolors='none', edgecolors='red', linewidth=2, label='Anomalous')

# Zone boundaries
ax.axvline(x=zone_boundary_1, color='cyan', linestyle='--', linewidth=1.5, label=f'Zone boundary ({zone_boundary_1:.0f})')
ax.axvline(x=zone_boundary_2, color='cyan', linestyle='--', linewidth=1.5, label=f'Zone boundary ({zone_boundary_2:.0f})')
ax.set_title('(a) Spatial Map: Grains by Area')
ax.set_xlabel('X position (px)')
ax.set_ylabel('Y position (px)')
ax.legend(loc='upper right', fontsize=8)

# Panel (b): Box plots by zone
ax = axes[0, 1]
zone_data = [left_areas, center_areas, right_areas]
zone_labels_plot = [f'Left\n(n={len(left_areas)})', f'Center\n(n={len(center_areas)})', f'Right\n(n={len(right_areas)})']
bp = ax.boxplot([d if len(d) > 0 else [0] for d in zone_data], labels=zone_labels_plot, patch_artist=True)
colors_box = ['#3498db', '#2ecc71', '#e74c3c']
for patch, color in zip(bp['boxes'], colors_box):
    patch.set_facecolor(color)
    patch.set_alpha(0.6)
ax.set_title(f'(b) Grain Area by Zone\nKW H={kw_stat:.2f}, p={kw_p:.4f}')
ax.set_ylabel('Grain Area (px²)')
# Annotate pairwise p-values
y_max = max([np.max(d) if len(d) > 0 else 0 for d in zone_data]) * 1.1
for i, (key, pval) in enumerate(pairwise_results.items()):
    ax.text(i + 1.5, y_max * (0.9 + 0.05 * i), f'{key}: p={pval:.4f}', fontsize=7, ha='center')

# Panel (c): Scatter + regression
ax = axes[1, 0]
if len(raw_areas) > 0:
    ax.scatter(raw_centroid_x, raw_areas, alpha=0.3, s=15, c='gray', label=f'Raw (R²={raw_R2:.3f})')
if len(areas) > 0:
    ax.scatter(centroid_x, areas, alpha=0.5, s=20, c='blue', label=f'Cleaned (R²={clean_R2:.3f})')
    # Regression line
    x_line = np.linspace(0, img_width, 100)
    y_line = clean_slope * x_line + clean_intercept
    ax.plot(x_line, y_line, 'r-', linewidth=2, label=f'Slope={clean_slope:.2f}')
    # Bootstrap CI
    if len(bootstrap_slopes) > 0:
        y_low = slope_ci_low * x_line + clean_intercept
        y_high = slope_ci_high * x_line + clean_intercept
        ax.fill_between(x_line, y_low, y_high, alpha=0.2, color='red', label=f'95% CI [{slope_ci_low:.2f}, {slope_ci_high:.2f}]')
ax.set_title('(c) Area vs. X-Position with Regression')
ax.set_xlabel('Centroid X (px)')
ax.set_ylabel('Grain Area (px²)')
ax.legend(fontsize=8)

# Panel (d): Histogram + fitted distributions
ax = axes[1, 1]
if len(ed) > 0:
    n_bins = min(30, max(10, int(np.sqrt(len(ed)))))
    counts, bin_edges, patches = ax.hist(ed, bins=n_bins, density=True, alpha=0.6, color='steelblue', label='Data')
    
    x_fit = np.linspace(max(ed.min() * 0.5, 0.1), ed.max() * 1.3, 200)
    
    # Log-normal PDF
    y_lognorm = stats.lognorm.pdf(x_fit, lognorm_shape, loc=0, scale=lognorm_scale)
    ax.plot(x_fit, y_lognorm, 'r-', linewidth=2, label=f'Log-normal (μ={lognorm_mu:.2f}, σ={lognorm_sigma:.2f})\nAIC={AIC_lognorm:.1f}')
    
    # Normal PDF
    y_normal = stats.norm.pdf(x_fit, loc=norm_mu, scale=norm_sigma)
    ax.plot(x_fit, y_normal, 'g--', linewidth=2, label=f'Normal (μ={norm_mu:.2f}, σ={norm_sigma:.2f})\nAIC={AIC_normal:.1f}')
    
    ax.set_title(f'(d) Equivalent Diameter Distribution\nSW(log): W={sw_W:.4f}, p={sw_p:.4f}')
    ax.set_xlabel('Equivalent Diameter (px)')
    ax.set_ylabel('Density')
    ax.legend(fontsize=7)

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

# ============================================================
# Save outputs
# ============================================================
np.save('analysis_labels.npy', labels)

# Save cleaned grain features
cleaned_features = np.array([(g['label'], g['area'], g['centroid_x'], g['centroid_y'],
                               g['equivalent_diameter'], g['aspect_ratio'], g['solidity'])
                              for g in cleaned_grains],
                             dtype=[('label', 'i4'), ('area', 'f8'), ('centroid_x', 'f8'),
                                    ('centroid_y', 'f8'), ('equiv_diameter', 'f8'),
                                    ('aspect_ratio', 'f8'), ('solidity', 'f8')])
np.save('grain_features.npy', cleaned_features)

# ============================================================
# Print results as JSON
# ============================================================
results = {
    "analysis_type": "Tier 2 integrated spatial gradient validation and grain size distribution fitting",
    "extracted_features": {
        "cleaned_grain_count": int(len(cleaned_grains)),
        "excluded_grain_count": int(len(excluded_ids)),
        "zone_median_areas_left_center_right": [
            left_stats['median'], center_stats['median'], right_stats['median']
        ],
        "zone_counts_left_center_right": [
            left_stats['count'], center_stats['count'], right_stats['count']
        ],
        "kruskal_wallis_H_statistic": float(kw_stat),
        "kruskal_wallis_p_value": float(kw_p),
        "pairwise_mannwhitney_p_values": pairwise_results,
        "cleaned_regression_slope_px2_per_px": float(clean_slope) if not np.isnan(clean_slope) else None,
        "cleaned_regression_R2": float(clean_R2) if not np.isnan(clean_R2) else None,
        "raw_regression_R2": float(raw_R2) if not np.isnan(raw_R2) else None,
        "bootstrap_slope_95CI": [slope_ci_low, slope_ci_high],
        "lognormal_mu": float(lognorm_mu),
        "lognormal_sigma": float(lognorm_sigma),
        "shapiro_wilk_W_on_log_diameters": float(sw_W),
        "shapiro_wilk_p_value": float(sw_p),
        "anderson_darling_statistic_lognormal": ad_stat_lognorm,
        "anderson_darling_statistic_normal": ad_stat_normal,
        "AIC_lognormal": float(AIC_lognorm),
        "AIC_normal": float(AIC_normal),
        "BIC_lognormal": float(BIC_lognorm),
        "BIC_normal": float(BIC_normal),
        "anomalous_grain_count": anom_count,
        "anomalous_grain_spatial_positions": anom_positions[:20],  # limit for JSON size
        "anomalous_vs_population_centroid_x_pvalue": p_anom
    },
    "quality_metrics": {
        "total_grains_before_cleaning": len(grain_data),
        "cleaning_min_area_threshold": MIN_AREA,
        "cleaning_min_solidity_threshold": MIN_SOLIDITY,
        "anomalous_aspect_ratio_threshold": ASPECT_RATIO_THRESH,
        "anomalous_solidity_threshold": SOLIDITY_THRESH,
        "preferred_distribution": "lognormal" if AIC_lognorm < AIC_normal else "normal",
        "bootstrap_iterations": n_bootstrap
    },
    "summary": f"Tier 2 analysis of {len(cleaned_grains)} cleaned grains (from {len(grain_data)} total). Spatial gradient regression slope={clean_slope:.2f} px²/px (R²={clean_R2:.3f}, 95% CI [{slope_ci_low:.2f}, {slope_ci_high:.2f}]). Kruskal-Wallis p={kw_p:.4f}. Log-normal fit: μ={lognorm_mu:.2f}, σ={lognorm_sigma:.2f} (AIC={AIC_lognorm:.1f} vs normal AIC={AIC_normal:.1f}). {anom_count} anomalous grains identified.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map from segmentation, {len(grain_data)} grains labeled 1-{len(grain_data)}, background=0",
            "shape": list(labels.shape),
            "dtype": str(labels.dtype)
        },
        "grain_features.npy": {
            "description": f"Structured array of {len(cleaned_grains)} cleaned grain features (label, area, centroid_x, centroid_y, equiv_diameter, aspect_ratio, solidity)",
            "shape": list(cleaned_features.shape),
            "dtype": str(cleaned_features.dtype)
        }
    }
}

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