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

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_annealing_data/out_components/temp_spectrum_0.npy')

# Handle different data shapes
if data.ndim == 2:
    if data.shape[0] == 2:
        x = data[0]
        y = data[1]
    elif data.shape[1] == 2:
        x = data[:, 0]
        y = data[:, 1]
    else:
        x = np.linspace(0, 100, data.shape[1])
        y = data[0]
elif data.ndim == 1:
    x = np.linspace(0, 100, len(data))
    y = data
else:
    raise ValueError(f"Unexpected data shape: {data.shape}")

# Sort by x
sort_idx = np.argsort(x)
x = x[sort_idx]
y = y[sort_idx]

print(f"Data loaded: {len(x)} points, x range [{x.min():.4f}, {x.max():.4f}], y range [{y.min():.6f}, {y.max():.6f}]")

# Define the model: 4 Gaussians + linear baseline
def model_func(x, a, b, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3, A4, mu4, sigma4):
    baseline = a * x + b
    g1 = A1 * np.exp(-(x - mu1)**2 / (2 * sigma1**2))
    g2 = A2 * np.exp(-(x - mu2)**2 / (2 * sigma2**2))
    g3 = A3 * np.exp(-(x - mu3)**2 / (2 * sigma3**2))
    g4 = A4 * np.exp(-(x - mu4)**2 / (2 * sigma4**2))
    return baseline + g1 + g2 + g3 + g4

# Step 1: Estimate baseline from regions away from peaks
baseline_mask = (x < 5) | ((x > 30) & (x < 35)) | ((x > 50) & (x < 55)) | ((x > 75) & (x < 80)) | (x > 95)
if np.sum(baseline_mask) > 5:
    from numpy.polynomial import polynomial as P
    baseline_coeffs = np.polyfit(x[baseline_mask], y[baseline_mask], 1)
    a_init = baseline_coeffs[0]
    b_init = baseline_coeffs[1]
    print(f"Baseline estimate: slope={a_init:.6f}, intercept={b_init:.6f}")
else:
    a_init = 0.0
    b_init = 0.1

# Step 2: Subtract baseline and identify peaks
y_residual = y - (a_init * x + b_init)

# Refine initial guesses from residual data
def find_peak_in_range(x, y, x_low, x_high):
    mask = (x >= x_low) & (x <= x_high)
    if np.sum(mask) == 0:
        return (x_low + x_high) / 2, 1.0
    idx = np.argmax(y[mask])
    return x[mask][idx], y[mask][idx]

mu1_init, A1_init = find_peak_in_range(x, y_residual, 10, 30)
mu2_init, A2_init = find_peak_in_range(x, y_residual, 33, 53)
mu3_init, A3_init = find_peak_in_range(x, y_residual, 55, 75)
mu4_init, A4_init = find_peak_in_range(x, y_residual, 77, 97)

# Ensure positive amplitudes
A1_init = max(A1_init, 0.5)
A2_init = max(A2_init, 0.5)
A3_init = max(A3_init, 0.5)
A4_init = max(A4_init, 0.5)

print(f"Initial peak estimates:")
print(f"  Peak 1: center={mu1_init:.2f}, amplitude={A1_init:.4f}")
print(f"  Peak 2: center={mu2_init:.2f}, amplitude={A2_init:.4f}")
print(f"  Peak 3: center={mu3_init:.2f}, amplitude={A3_init:.4f}")
print(f"  Peak 4: center={mu4_init:.2f}, amplitude={A4_init:.4f}")

# Initial parameter vector
# [a, b, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3, A4, mu4, sigma4]
p0 = [a_init, b_init,
      A1_init, mu1_init, 4.0,
      A2_init, mu2_init, 4.0,
      A3_init, mu3_init, 3.0,
      A4_init, mu4_init, 4.0]

# Bounds
lower_bounds = [-0.05, -5.0,
                0.0, 10.0, 0.5,
                0.0, 33.0, 0.5,
                0.0, 55.0, 0.5,
                0.0, 77.0, 0.5]
upper_bounds = [0.05, 5.0,
                20.0, 30.0, 15.0,
                20.0, 53.0, 15.0,
                20.0, 75.0, 15.0,
                20.0, 97.0, 15.0]

# Step 3: Two-stage fitting approach
# Stage 1: Fix baseline, fit Gaussians
try:
    popt, pcov = curve_fit(model_func, x, y, p0=p0,
                           bounds=(lower_bounds, upper_bounds),
                           maxfev=2000, ftol=1e-10, xtol=1e-10, gtol=1e-10,
                           method='trf')
    print("Stage 1 fit converged.")
except Exception as e:
    print(f"Stage 1 warning: {e}")
    # Try with default tolerances
    popt, pcov = curve_fit(model_func, x, y, p0=p0,
                           bounds=(lower_bounds, upper_bounds),
                           maxfev=5000, method='trf')

# Stage 2: Global refinement with tighter tolerances
try:
    popt2, pcov2 = curve_fit(model_func, x, y, p0=popt,
                             bounds=(lower_bounds, upper_bounds),
                             maxfev=5000, ftol=1e-12, xtol=1e-12, gtol=1e-12,
                             method='trf')
    popt = popt2
    pcov = pcov2
    print("Stage 2 refinement converged.")
except Exception as e:
    print(f"Stage 2 refinement note: {e}")

# Extract parameters
a_fit, b_fit = popt[0], popt[1]
A1, mu1, sigma1 = popt[2], popt[3], popt[4]
A2, mu2, sigma2 = popt[5], popt[6], popt[7]
A3, mu3, sigma3 = popt[8], popt[9], popt[10]
A4, mu4, sigma4 = popt[11], popt[12], popt[13]

# Compute uncertainties from covariance matrix
if pcov is not None and not np.any(np.isinf(pcov)):
    perr = np.sqrt(np.diag(pcov))
else:
    perr = np.zeros(len(popt))
    print("Warning: Could not compute parameter uncertainties")

# Compute fit quality
y_fit = model_func(x, *popt)
residuals = y - y_fit
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1.0 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals**2))
residual_std = np.std(residuals)

# Compute FWHM for each Gaussian
fwhm_factor = 2.0 * np.sqrt(2.0 * np.log(2.0))  # ~2.355
fwhm1 = fwhm_factor * abs(sigma1)
fwhm2 = fwhm_factor * abs(sigma2)
fwhm3 = fwhm_factor * abs(sigma3)
fwhm4 = fwhm_factor * abs(sigma4)

# FWHM uncertainties
fwhm1_err = fwhm_factor * perr[4]
fwhm2_err = fwhm_factor * perr[7]
fwhm3_err = fwhm_factor * perr[10]
fwhm4_err = fwhm_factor * perr[13]

print(f"\nFit Quality: R² = {r_squared:.6f}, RMSE = {rmse:.6f}")
print(f"Residual std = {residual_std:.6f}")

# Check for systematic residuals at known peak locations
for name, center in [('Peak1', mu1), ('Peak2', mu2), ('Peak3', mu3), ('Peak4', mu4)]:
    mask_peak = (x > center - 5) & (x < center + 5)
    if np.sum(mask_peak) > 0:
        max_res = np.max(np.abs(residuals[mask_peak]))
        print(f"  Max |residual| near {name} (x={center:.1f}): {max_res:.6f}")

# Visualization
fig, axes = plt.subplots(2, 1, figsize=(12, 8), gridspec_kw={'height_ratios': [3, 1]})

# Top panel: data + fit + components
ax1 = axes[0]
ax1.scatter(x, y, s=2, alpha=0.5, color='gray', label='Data', zorder=1)
ax1.plot(x, y_fit, 'r-', linewidth=2, label=f'Fit (R²={r_squared:.4f})', zorder=3)

# Plot individual components
baseline = a_fit * x + b_fit
ax1.plot(x, baseline, 'k--', linewidth=1, alpha=0.7, label=f'Baseline (y={a_fit:.4f}x+{b_fit:.4f})')

g1 = A1 * np.exp(-(x - mu1)**2 / (2 * sigma1**2)) + baseline
g2 = A2 * np.exp(-(x - mu2)**2 / (2 * sigma2**2)) + baseline
g3 = A3 * np.exp(-(x - mu3)**2 / (2 * sigma3**2)) + baseline
g4 = A4 * np.exp(-(x - mu4)**2 / (2 * sigma4**2)) + baseline

ax1.plot(x, g1, 'b--', linewidth=1, alpha=0.7, label=f'Peak 1 (μ={mu1:.2f})')
ax1.plot(x, g2, 'g--', linewidth=1, alpha=0.7, label=f'Peak 2 (μ={mu2:.2f})')
ax1.plot(x, g3, 'm--', linewidth=1, alpha=0.7, label=f'Peak 3 (μ={mu3:.2f})')
ax1.plot(x, g4, 'c--', linewidth=1, alpha=0.7, label=f'Peak 4 (μ={mu4:.2f})')

# Fill individual Gaussian components (without baseline offset for clarity)
g1_only = A1 * np.exp(-(x - mu1)**2 / (2 * sigma1**2))
g2_only = A2 * np.exp(-(x - mu2)**2 / (2 * sigma2**2))
g3_only = A3 * np.exp(-(x - mu3)**2 / (2 * sigma3**2))
g4_only = A4 * np.exp(-(x - mu4)**2 / (2 * sigma4**2))

ax1.fill_between(x, baseline, g1, alpha=0.15, color='blue')
ax1.fill_between(x, baseline, g2, alpha=0.15, color='green')
ax1.fill_between(x, baseline, g3, alpha=0.15, color='magenta')
ax1.fill_between(x, baseline, g4, alpha=0.15, color='cyan')

ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title('Spectroscopic Data: 4-Gaussian + Linear Baseline Fit', fontsize=14)
ax1.legend(loc='upper right', fontsize=8)
ax1.grid(True, alpha=0.3)

# Bottom panel: residuals
ax2 = axes[1]
ax2.scatter(x, residuals, s=2, alpha=0.5, color='gray')
ax2.axhline(y=0, color='r', linewidth=1)
ax2.axhline(y=2*residual_std, color='r', linewidth=0.5, linestyle='--', alpha=0.5)
ax2.axhline(y=-2*residual_std, color='r', linewidth=0.5, linestyle='--', alpha=0.5)
ax2.set_xlabel('x', fontsize=12)
ax2.set_ylabel('Residuals', fontsize=12)
ax2.set_title(f'Residuals (std={residual_std:.6f})', fontsize=11)
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")

# Build results JSON
results = {
    "model_type": "4 Gaussians + linear baseline: y = (a*x + b) + sum_i[A_i * exp(-(x-mu_i)^2 / (2*sigma_i^2))]",
    "parameters": {
        "baseline": {
            "slope": float(f"{a_fit:.6f}"),
            "slope_err": float(f"{perr[0]:.6f}"),
            "intercept": float(f"{b_fit:.6f}"),
            "intercept_err": float(f"{perr[1]:.6f}")
        },
        "peak_1": {
            "center": float(f"{mu1:.4f}"),
            "center_err": float(f"{perr[3]:.4f}"),
            "amplitude": float(f"{A1:.4f}"),
            "amplitude_err": float(f"{perr[2]:.4f}"),
            "sigma": float(f"{sigma1:.4f}"),
            "sigma_err": float(f"{perr[4]:.4f}"),
            "fwhm": float(f"{fwhm1:.4f}"),
            "fwhm_err": float(f"{fwhm1_err:.4f}")
        },
        "peak_2": {
            "center": float(f"{mu2:.4f}"),
            "center_err": float(f"{perr[6]:.4f}"),
            "amplitude": float(f"{A2:.4f}"),
            "amplitude_err": float(f"{perr[5]:.4f}"),
            "sigma": float(f"{sigma2:.4f}"),
            "sigma_err": float(f"{perr[7]:.4f}"),
            "fwhm": float(f"{fwhm2:.4f}"),
            "fwhm_err": float(f"{fwhm2_err:.4f}")
        },
        "peak_3": {
            "center": float(f"{mu3:.4f}"),
            "center_err": float(f"{perr[9]:.4f}"),
            "amplitude": float(f"{A3:.4f}"),
            "amplitude_err": float(f"{perr[8]:.4f}"),
            "sigma": float(f"{sigma3:.4f}"),
            "sigma_err": float(f"{perr[10]:.4f}"),
            "fwhm": float(f"{fwhm3:.4f}"),
            "fwhm_err": float(f"{fwhm3_err:.4f}")
        },
        "peak_4": {
            "center": float(f"{mu4:.4f}"),
            "center_err": float(f"{perr[12]:.4f}"),
            "amplitude": float(f"{A4:.4f}"),
            "amplitude_err": float(f"{perr[11]:.4f}"),
            "sigma": float(f"{sigma4:.4f}"),
            "sigma_err": float(f"{perr[13]:.4f}"),
            "fwhm": float(f"{fwhm4:.4f}"),
            "fwhm_err": float(f"{fwhm4_err:.4f}")
        }
    },
    "fit_quality": {
        "r_squared": float(f"{r_squared:.6f}"),
        "rmse": float(f"{rmse:.6f}"),
        "residual_std": float(f"{residual_std:.6f}")
    },
    "summary": (
        f"Four Gaussian peaks on a linear baseline fitted to spectroscopic data (1000 points). "
        f"R²={r_squared:.4f}, RMSE={rmse:.6f}. "
        f"Peak centers at x={mu1:.2f}, {mu2:.2f}, {mu3:.2f}, {mu4:.2f} with "
        f"FWHM={fwhm1:.2f}, {fwhm2:.2f}, {fwhm3:.2f}, {fwhm4:.2f}. "
        f"Baseline slope={a_fit:.5f}. "
        f"Two-stage bounded nonlinear least-squares (TRF) optimization with "
        f"initial guesses refined from baseline-subtracted peak detection."
    )
}

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