import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy.special import wofz

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_series/temp_spectrum_3.npy')
x = data[:, 0]
y = data[:, 1]

# First, let's visualize the data characteristics to understand why single Gaussian failed
# The locked model got R²=0.7384, which suggests the peak shape is fundamentally different
# Possibilities: two overlapping peaks, or a non-Gaussian lineshape (e.g., Lorentzian or Voigt)

# Strategy: Try multiple models and pick the simplest one that achieves R² >= 0.99
# 1. Single Voigt profile (different lineshape than Gaussian)
# 2. Two Gaussian peaks
# 3. Single Gaussian + linear/sloped baseline

# Let's first examine the data more carefully
y_max = np.max(y)
x_at_max = x[np.argmax(y)]

# Estimate baseline from edges
baseline_region = np.concatenate([y[x < 5], y[x > 45]])
if len(baseline_region) > 0:
    baseline_est = np.mean(baseline_region)
else:
    baseline_est = np.min(y)

print(f"Data inspection: x_at_max={x_at_max:.2f}, y_max={y_max:.4f}, baseline_est={baseline_est:.4f}")
print(f"Y range: [{np.min(y):.4f}, {np.max(y):.4f}]")

# Check for asymmetry or multiple peaks
# Look at the peak region
peak_amp = y_max - baseline_est
half_max = baseline_est + peak_amp / 2

# Find half-max positions on each side
above_half = x[y > half_max]
if len(above_half) > 0:
    left_hm = above_half[0]
    right_hm = above_half[-1]
    fwhm_est = right_hm - left_hm
    center_of_fwhm = (left_hm + right_hm) / 2
    print(f"FWHM estimate: {fwhm_est:.2f}, center of FWHM: {center_of_fwhm:.2f}")
    print(f"Peak center vs FWHM center offset: {x_at_max - center_of_fwhm:.2f}")
else:
    fwhm_est = 5.0

# Check for second derivative to find multiple peaks
from scipy.ndimage import gaussian_filter1d
y_smooth = gaussian_filter1d(y, sigma=3)
y_deriv2 = np.gradient(np.gradient(y_smooth, x), x)

# Find local minima in second derivative (peak positions)
from scipy.signal import find_peaks
neg_d2 = -y_deriv2
peaks_d2, _ = find_peaks(neg_d2, height=0, prominence=0.001)
print(f"Potential peak positions from 2nd derivative: {x[peaks_d2]}")

# Also check residuals with a single Gaussian to understand the failure mode
def single_gaussian(x, a, A, mu, sigma):
    return a + A * np.exp(-0.5 * ((x - mu) / sigma) ** 2)

try:
    p0_sg = [baseline_est, peak_amp, x_at_max, fwhm_est / 2.355]
    popt_sg, _ = curve_fit(single_gaussian, x, y, p0=p0_sg, maxfev=10000)
    y_fit_sg = single_gaussian(x, *popt_sg)
    ss_res_sg = np.sum((y - y_fit_sg) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r2_sg = 1 - ss_res_sg / ss_tot
    print(f"Single Gaussian R² = {r2_sg:.4f}")
    
    # Look at residuals
    residuals_sg = y - y_fit_sg
    print(f"Residual range: [{np.min(residuals_sg):.4f}, {np.max(residuals_sg):.4f}]")
    # Check if residuals show systematic pattern
    residual_peaks, _ = find_peaks(np.abs(residuals_sg), height=0.1*peak_amp)
    print(f"Significant residual peak positions: {x[residual_peaks]}")
except:
    r2_sg = 0
    print("Single Gaussian fit failed")

# ============================================================
# Model 1: Two Gaussian peaks on constant baseline
# ============================================================
def two_gaussians(x, a, A1, mu1, sigma1, A2, mu2, sigma2):
    g1 = A1 * np.exp(-0.5 * ((x - mu1) / sigma1) ** 2)
    g2 = A2 * np.exp(-0.5 * ((x - mu2) / sigma2) ** 2)
    return a + g1 + g2

# ============================================================
# Model 2: Single Voigt profile on constant baseline
# ============================================================
def voigt_profile(x, a, A, mu, sigma, gamma):
    """Voigt profile: convolution of Gaussian and Lorentzian"""
    z = ((x - mu) + 1j * gamma) / (sigma * np.sqrt(2))
    v = np.real(wofz(z)) / (sigma * np.sqrt(2 * np.pi))
    # Normalize so A is the amplitude
    z0 = (1j * gamma) / (sigma * np.sqrt(2))
    v0 = np.real(wofz(z0)) / (sigma * np.sqrt(2 * np.pi))
    if v0 > 0:
        return a + A * v / v0
    else:
        return a + A * v

# ============================================================
# Model 3: Single Lorentzian on constant baseline
# ============================================================
def single_lorentzian(x, a, A, mu, gamma):
    return a + A / (1 + ((x - mu) / gamma) ** 2)

# ============================================================
# Model 4: Pseudo-Voigt (linear combination of Gaussian and Lorentzian)
# ============================================================
def pseudo_voigt(x, a, A, mu, sigma, eta):
    """eta=0 pure Gaussian, eta=1 pure Lorentzian"""
    # Use same width parameter for both
    gauss = np.exp(-0.5 * ((x - mu) / sigma) ** 2)
    gamma = sigma * np.sqrt(2 * np.log(2))  # match FWHM
    lorentz = 1 / (1 + ((x - mu) / gamma) ** 2)
    return a + A * ((1 - eta) * gauss + eta * lorentz)

# Try all models and collect results
results_all = {}

# --- Try Single Lorentzian ---
try:
    p0_lor = [baseline_est, peak_amp, x_at_max, fwhm_est / 2]
    bounds_lor = ([0, 0, 0, 0.1], [1, 10, 50, 20])
    popt_lor, pcov_lor = curve_fit(single_lorentzian, x, y, p0=p0_lor, bounds=bounds_lor, maxfev=20000)
    y_fit_lor = single_lorentzian(x, *popt_lor)
    ss_res = np.sum((y - y_fit_lor) ** 2)
    r2_lor = 1 - ss_res / ss_tot
    rmse_lor = np.sqrt(np.mean((y - y_fit_lor) ** 2))
    print(f"Single Lorentzian R² = {r2_lor:.6f}, RMSE = {rmse_lor:.6f}")
    results_all['lorentzian'] = {'r2': r2_lor, 'rmse': rmse_lor, 'popt': popt_lor, 'pcov': pcov_lor}
except Exception as e:
    print(f"Lorentzian fit failed: {e}")
    results_all['lorentzian'] = {'r2': -1}

# --- Try Pseudo-Voigt ---
try:
    p0_pv = [baseline_est, peak_amp, x_at_max, fwhm_est / 2.355, 0.5]
    bounds_pv = ([0, 0, 0, 0.1, 0], [1, 10, 50, 20, 1])
    popt_pv, pcov_pv = curve_fit(pseudo_voigt, x, y, p0=p0_pv, bounds=bounds_pv, maxfev=20000)
    y_fit_pv = pseudo_voigt(x, *popt_pv)
    ss_res = np.sum((y - y_fit_pv) ** 2)
    r2_pv = 1 - ss_res / ss_tot
    rmse_pv = np.sqrt(np.mean((y - y_fit_pv) ** 2))
    print(f"Pseudo-Voigt R² = {r2_pv:.6f}, RMSE = {rmse_pv:.6f}")
    results_all['pseudo_voigt'] = {'r2': r2_pv, 'rmse': rmse_pv, 'popt': popt_pv, 'pcov': pcov_pv}
except Exception as e:
    print(f"Pseudo-Voigt fit failed: {e}")
    results_all['pseudo_voigt'] = {'r2': -1}

# --- Try Voigt profile ---
try:
    p0_v = [baseline_est, peak_amp, x_at_max, fwhm_est / 2.355, fwhm_est / 2]
    bounds_v = ([0, 0, 0, 0.1, 0.01], [1, 10, 50, 20, 20])
    popt_v, pcov_v = curve_fit(voigt_profile, x, y, p0=p0_v, bounds=bounds_v, maxfev=20000)
    y_fit_v = voigt_profile(x, *popt_v)
    ss_res = np.sum((y - y_fit_v) ** 2)
    r2_v = 1 - ss_res / ss_tot
    rmse_v = np.sqrt(np.mean((y - y_fit_v) ** 2))
    print(f"Voigt R² = {r2_v:.6f}, RMSE = {rmse_v:.6f}")
    results_all['voigt'] = {'r2': r2_v, 'rmse': rmse_v, 'popt': popt_v, 'pcov': pcov_v}
except Exception as e:
    print(f"Voigt fit failed: {e}")
    results_all['voigt'] = {'r2': -1}

# --- Try Two Gaussians ---
try:
    # Guess two peaks: one near the main peak, one offset
    # Try several initial guesses
    best_r2_2g = -1
    best_popt_2g = None
    best_pcov_2g = None
    
    # Generate multiple starting points for the second peak
    offsets = [-8, -5, -3, 3, 5, 8]
    amp_fracs = [0.3, 0.5, 0.7]
    
    for offset in offsets:
        for frac in amp_fracs:
            try:
                mu2_guess = x_at_max + offset
                if mu2_guess < 0 or mu2_guess > 50:
                    continue
                p0_2g = [baseline_est, peak_amp * (1-frac), x_at_max, fwhm_est / 2.355,
                         peak_amp * frac, mu2_guess, fwhm_est / 2.355]
                bounds_2g = ([0, 0, 0, 0.1, 0, 0, 0.1],
                            [1, 10, 50, 20, 10, 50, 20])
                popt_2g_try, pcov_2g_try = curve_fit(two_gaussians, x, y, p0=p0_2g, 
                                                      bounds=bounds_2g, maxfev=30000)
                y_fit_2g_try = two_gaussians(x, *popt_2g_try)
                ss_res = np.sum((y - y_fit_2g_try) ** 2)
                r2_2g_try = 1 - ss_res / ss_tot
                if r2_2g_try > best_r2_2g:
                    best_r2_2g = r2_2g_try
                    best_popt_2g = popt_2g_try
                    best_pcov_2g = pcov_2g_try
            except:
                continue
    
    if best_popt_2g is not None:
        y_fit_2g = two_gaussians(x, *best_popt_2g)
        rmse_2g = np.sqrt(np.mean((y - y_fit_2g) ** 2))
        print(f"Two Gaussians R² = {best_r2_2g:.6f}, RMSE = {rmse_2g:.6f}")
        print(f"  Peak 1: A={best_popt_2g[1]:.3f}, mu={best_popt_2g[2]:.3f}, sigma={best_popt_2g[3]:.3f}")
        print(f"  Peak 2: A={best_popt_2g[4]:.3f}, mu={best_popt_2g[5]:.3f}, sigma={best_popt_2g[6]:.3f}")
        results_all['two_gaussians'] = {'r2': best_r2_2g, 'rmse': rmse_2g, 'popt': best_popt_2g, 'pcov': best_pcov_2g}
    else:
        results_all['two_gaussians'] = {'r2': -1}
except Exception as e:
    print(f"Two Gaussians fit failed: {e}")
    results_all['two_gaussians'] = {'r2': -1}

# --- Try Single Gaussian + linear baseline ---
def gaussian_linear_baseline(x, a, b, A, mu, sigma):
    return a + b * x + A * np.exp(-0.5 * ((x - mu) / sigma) ** 2)

try:
    p0_glb = [baseline_est, 0.0, peak_amp, x_at_max, fwhm_est / 2.355]
    bounds_glb = ([-1, -0.1, 0, 0, 0.1], [2, 0.1, 10, 50, 20])
    popt_glb, pcov_glb = curve_fit(gaussian_linear_baseline, x, y, p0=p0_glb, bounds=bounds_glb, maxfev=20000)
    y_fit_glb = gaussian_linear_baseline(x, *popt_glb)
    ss_res = np.sum((y - y_fit_glb) ** 2)
    r2_glb = 1 - ss_res / ss_tot
    rmse_glb = np.sqrt(np.mean((y - y_fit_glb) ** 2))
    print(f"Gaussian + linear baseline R² = {r2_glb:.6f}, RMSE = {rmse_glb:.6f}")
    results_all['gauss_linear'] = {'r2': r2_glb, 'rmse': rmse_glb, 'popt': popt_glb, 'pcov': pcov_glb}
except Exception as e:
    print(f"Gaussian + linear baseline fit failed: {e}")
    results_all['gauss_linear'] = {'r2': -1}

# ============================================================
# Select the best model (simplest that achieves R² >= 0.99)
# ============================================================
# Preference order for parsimony (fewest parameters first):
# single_lorentzian (4 params) < gauss_linear (5 params) < pseudo_voigt (5 params) < voigt (5 params) < two_gaussians (7 params)

print("\n=== Model Selection ===")
model_priority = ['lorentzian', 'gauss_linear', 'pseudo_voigt', 'voigt', 'two_gaussians']
selected_model = None
for model_name in model_priority:
    if model_name in results_all and results_all[model_name]['r2'] >= 0.99:
        selected_model = model_name
        print(f"Selected: {model_name} (R² = {results_all[model_name]['r2']:.6f})")
        break

# If none meets threshold, pick the best
if selected_model is None:
    best_model = max(results_all.keys(), key=lambda k: results_all[k]['r2'])
    selected_model = best_model
    print(f"No model met R²>=0.99. Best: {selected_model} (R² = {results_all[selected_model]['r2']:.6f})")

print(f"\nFinal selected model: {selected_model}")

# ============================================================
# Generate final results based on selected model
# ============================================================
sel = results_all[selected_model]
r2_final = sel['r2']
rmse_final = sel['rmse']
popt_final = sel['popt']
pcov_final = sel['pcov']
perr_final = np.sqrt(np.diag(pcov_final)) if pcov_final is not None else np.zeros_like(popt_final)

if selected_model == 'lorentzian':
    y_fit_final = single_lorentzian(x, *popt_final)
    model_desc = "Single Lorentzian peak on constant baseline: y = a + A / (1 + ((x-mu)/gamma)^2)"
    fwhm_val = 2 * popt_final[3]
    fwhm_err = 2 * perr_final[3]
    params = {
        "peak_1": {
            "center": float(popt_final[2]),
            "center_err": float(perr_final[2]),
            "amplitude": float(popt_final[1]),
            "amplitude_err": float(perr_final[1]),
            "gamma": float(popt_final[3]),
            "gamma_err": float(perr_final[3]),
            "FWHM": float(fwhm_val),
            "FWHM_err": float(fwhm_err)
        },
        "baseline": {
            "offset": float(popt_final[0]),
            "offset_err": float(perr_final[0])
        }
    }
    components = None

elif selected_model == 'pseudo_voigt':
    y_fit_final = pseudo_voigt(x, *popt_final)
    model_desc = f"Pseudo-Voigt peak on constant baseline: y = a + A*[(1-eta)*G(x) + eta*L(x)], eta={popt_final[4]:.3f}"
    sigma_val = popt_final[3]
    fwhm_gauss = 2.3548 * sigma_val
    # For pseudo-Voigt, the FWHM is approximately the same for both components since we matched them
    fwhm_val = fwhm_gauss  # approximate
    fwhm_err = 2.3548 * perr_final[3]
    params = {
        "peak_1": {
            "center": float(popt_final[2]),
            "center_err": float(perr_final[2]),
            "amplitude": float(popt_final[1]),
            "amplitude_err": float(perr_final[1]),
            "sigma": float(popt_final[3]),
            "sigma_err": float(perr_final[3]),
            "eta": float(popt_final[4]),
            "eta_err": float(perr_final[4]),
            "FWHM": float(fwhm_val),
            "FWHM_err": float(fwhm_err)
        },
        "baseline": {
            "offset": float(popt_final[0]),
            "offset_err": float(perr_final[0])
        }
    }
    components = None

elif selected_model == 'voigt':
    y_fit_final = voigt_profile(x, *popt_final)
    model_desc = "Voigt profile on constant baseline"
    sigma_v = popt_final[3]
    gamma_v = popt_final[4]
    # Approximate FWHM for Voigt
    fwhm_g = 2.3548 * sigma_v
    fwhm_l = 2 * gamma_v
    fwhm_val = 0.5346 * fwhm_l + np.sqrt(0.2166 * fwhm_l**2 + fwhm_g**2)
    params = {
        "peak_1": {
            "center": float(popt_final[2]),
            "center_err": float(perr_final[2]),
            "amplitude": float(popt_final[1]),
            "amplitude_err": float(perr_final[1]),
            "sigma": float(popt_final[3]),
            "sigma_err": float(perr_final[3]),
            "gamma": float(popt_final[4]),
            "gamma_err": float(perr_final[4]),
            "FWHM": float(fwhm_val),
            "FWHM_err": 0.0
        },
        "baseline": {
            "offset": float(popt_final[0]),
            "offset_err": float(perr_final[0])
        }
    }
    components = None

elif selected_model == 'two_gaussians':
    y_fit_final = two_gaussians(x, *popt_final)
    model_desc = "Two Gaussian peaks on constant baseline"
    fwhm1 = 2.3548 * popt_final[3]
    fwhm2 = 2.3548 * popt_final[6]
    fwhm1_err = 2.3548 * perr_final[3]
    fwhm2_err = 2.3548 * perr_final[6]
    # Sort peaks by center position
    if popt_final[2] <= popt_final[5]:
        pk1_idx, pk2_idx = (1,2,3), (4,5,6)
    else:
        pk1_idx, pk2_idx = (4,5,6), (1,2,3)
    params = {
        "peak_1": {
            "center": float(popt_final[pk1_idx[1]]),
            "center_err": float(perr_final[pk1_idx[1]]),
            "amplitude": float(popt_final[pk1_idx[0]]),
            "amplitude_err": float(perr_final[pk1_idx[0]]),
            "sigma": float(popt_final[pk1_idx[2]]),
            "sigma_err": float(perr_final[pk1_idx[2]]),
            "FWHM": float(2.3548 * popt_final[pk1_idx[2]]),
            "FWHM_err": float(2.3548 * perr_final[pk1_idx[2]])
        },
        "peak_2": {
            "center": float(popt_final[pk2_idx[1]]),
            "center_err": float(perr_final[pk2_idx[1]]),
            "amplitude": float(popt_final[pk2_idx[0]]),
            "amplitude_err": float(perr_final[pk2_idx[0]]),
            "sigma": float(popt_final[pk2_idx[2]]),
            "sigma_err": float(perr_final[pk2_idx[2]]),
            "FWHM": float(2.3548 * popt_final[pk2_idx[2]]),
            "FWHM_err": float(2.3548 * perr_final[pk2_idx[2]])
        },
        "baseline": {
            "offset": float(popt_final[0]),
            "offset_err": float(perr_final[0])
        }
    }
    # Compute components for plotting
    comp1 = popt_final[0] + popt_final[1] * np.exp(-0.5 * ((x - popt_final[2]) / popt_final[3]) ** 2)
    comp2 = popt_final[0] + popt_final[4] * np.exp(-0.5 * ((x - popt_final[5]) / popt_final[6]) ** 2)
    components = [comp1, comp2]

elif selected_model == 'gauss_linear':
    y_fit_final = gaussian_linear_baseline(x, *popt_final)
    model_desc = "Single Gaussian peak on linear baseline: y = a + b*x + A*exp(-0.5*((x-mu)/sigma)^2)"
    fwhm_val = 2.3548 * popt_final[4]
    fwhm_err = 2.3548 * perr_final[4]
    params = {
        "peak_1": {
            "center": float(popt_final[3]),
            "center_err": float(perr_final[3]),
            "amplitude": float(popt_final[2]),
            "amplitude_err": float(perr_final[2]),
            "sigma": float(popt_final[4]),
            "sigma_err": float(perr_final[4]),
            "FWHM": float(fwhm_val),
            "FWHM_err": float(fwhm_err)
        },
        "baseline": {
            "offset": float(popt_final[0]),
            "offset_err": float(perr_final[0]),
            "slope": float(popt_final[1]),
            "slope_err": float(perr_final[1])
        }
    }
    components = None

# ============================================================
# Visualization
# ============================================================
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), height_ratios=[3, 1], sharex=True)
fig.suptitle(f'Spectrum 4/5 Fit: {selected_model}', fontsize=14)

# Main plot
ax1.scatter(x, y, s=5, alpha=0.5, color='blue', label='Data', zorder=2)
ax1.plot(x, y_fit_final, 'r-', linewidth=2, label=f'Fit (R²={r2_final:.6f})', zorder=3)

if selected_model == 'two_gaussians' and components is not None:
    baseline_val = popt_final[0]
    # Show individual Gaussian components (on baseline)
    g1 = baseline_val + popt_final[1] * np.exp(-0.5 * ((x - popt_final[2]) / popt_final[3]) ** 2)
    g2 = baseline_val + popt_final[4] * np.exp(-0.5 * ((x - popt_final[5]) / popt_final[6]) ** 2)
    ax1.plot(x, g1, '--', color='green', alpha=0.7, label='Peak 1')
    ax1.plot(x, g2, '--', color='orange', alpha=0.7, label='Peak 2')
    ax1.axhline(y=baseline_val, color='gray', linestyle=':', alpha=0.5, label='Baseline')

ax1.set_ylabel('Intensity')
ax1.legend(loc='best')
ax1.grid(True, alpha=0.3)

# Residuals
residuals = y - y_fit_final
ax2.scatter(x, residuals, s=3, alpha=0.5, color='green')
ax2.axhline(y=0, color='red', linestyle='--', linewidth=1)
ax2.set_xlabel('x')
ax2.set_ylabel('Residuals')
ax2.grid(True, alpha=0.3)

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

# ============================================================
# Summary
# ============================================================
summary_parts = [f"Spectrum 4/5 (position=4): locked single Gaussian model failed (R²=0.7384)."]
summary_parts.append(f"Best alternative: {selected_model} with R²={r2_final:.6f}, RMSE={rmse_final:.6f}.")

if selected_model == 'two_gaussians':
    summary_parts.append(f"Two overlapping Gaussian peaks detected, suggesting a composite feature not captured by single Gaussian.")
elif selected_model == 'lorentzian':
    summary_parts.append(f"Lorentzian lineshape provides significantly better fit than Gaussian, indicating broader tails characteristic of lifetime broadening.")
elif selected_model == 'pseudo_voigt':
    eta_val = popt_final[4]
    summary_parts.append(f"Pseudo-Voigt with eta={eta_val:.3f} ({'more Lorentzian' if eta_val > 0.5 else 'more Gaussian'}) provides better fit, indicating mixed broadening mechanisms.")
elif selected_model == 'voigt':
    summary_parts.append(f"Voigt profile (Gaussian+Lorentzian convolution) indicates both instrumental and intrinsic broadening contribute.")
elif selected_model == 'gauss_linear':
    summary_parts.append(f"Linear baseline correction was needed; flat baseline assumption of locked model was inadequate.")

for name in results_all:
    if results_all[name]['r2'] > 0:
        summary_parts.append(f"{name}: R²={results_all[name]['r2']:.6f}")

summary = " ".join(summary_parts)

results = {
    "model_type": model_desc,
    "parameters": params,
    "fit_quality": {
        "r_squared": float(r2_final),
        "rmse": float(rmse_final)
    },
    "summary": summary
}

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