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_easy/temp_spectrum_0.npy')

# Determine x and y
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:
        # Assume first row/col is x, second is y
        x = data[0]
        y = data[1]
else:
    # 1D array, create x from index
    y = data
    x = np.linspace(0, 100, len(y))

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

# Define model: linear baseline + 2 Gaussians
def model(x, m, b, A1, mu1, sigma1, A2, mu2, sigma2):
    baseline = m * x + b
    g1 = A1 * np.exp(-(x - mu1)**2 / (2 * sigma1**2))
    g2 = A2 * np.exp(-(x - mu2)**2 / (2 * sigma2**2))
    return baseline + g1 + g2

# Initial guesses
p0 = [0.01, 0.0, 2.2, 40.0, 5.0, 1.8, 70.0, 6.0]

# Bounds: [m, b, A1, mu1, sigma1, A2, mu2, sigma2]
lower = [-0.05, -np.inf, 0.0, 30.0, 0.01, 0.0, 60.0, 0.01]
upper = [0.05, np.inf, np.inf, 50.0, np.inf, np.inf, 80.0, np.inf]

# Fit
try:
    popt, pcov = curve_fit(model, x, y, p0=p0, bounds=(lower, upper), maxfev=50000)
    perr = np.sqrt(np.diag(pcov))
except Exception as e:
    print(f"Fitting failed: {e}")
    # Try with looser bounds
    popt, pcov = curve_fit(model, x, y, p0=p0, maxfev=50000)
    perr = np.sqrt(np.diag(pcov))

m_fit, b_fit, A1_fit, mu1_fit, sigma1_fit, A2_fit, mu2_fit, sigma2_fit = popt
m_err, b_err, A1_err, mu1_err, sigma1_err, A2_err, mu2_err, sigma2_err = perr

# Compute fitted values
y_fit = model(x, *popt)

# Compute components
baseline = m_fit * x + b_fit
gauss1 = A1_fit * np.exp(-(x - mu1_fit)**2 / (2 * sigma1_fit**2))
gauss2 = A2_fit * np.exp(-(x - mu2_fit)**2 / (2 * sigma2_fit**2))

# Residuals
residuals = y - y_fit

# R² and RMSE
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))

# FWHM = 2*sqrt(2*ln(2))*sigma ≈ 2.3548*sigma
fwhm_factor = 2.0 * np.sqrt(2.0 * np.log(2.0))
fwhm1 = fwhm_factor * abs(sigma1_fit)
fwhm1_err = fwhm_factor * sigma1_err
fwhm2 = fwhm_factor * abs(sigma2_fit)
fwhm2_err = fwhm_factor * sigma2_err

# Visualization
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]}, sharex=True)

# Main plot
ax1.scatter(x, y, s=5, alpha=0.5, color='gray', label='Data', zorder=1)
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit', zorder=3)
ax1.plot(x, baseline, 'k--', linewidth=1, label=f'Baseline (y={m_fit:.4f}x+{b_fit:.4f})', zorder=2)
ax1.plot(x, baseline + gauss1, 'b--', linewidth=1.5, label=f'Peak 1 (center={mu1_fit:.2f})', zorder=2)
ax1.plot(x, baseline + gauss2, 'g--', linewidth=1.5, label=f'Peak 2 (center={mu2_fit:.2f})', zorder=2)
# Fill individual Gaussians on baseline
ax1.fill_between(x, baseline, baseline + gauss1, alpha=0.2, color='blue')
ax1.fill_between(x, baseline, baseline + gauss2, alpha=0.2, color='green')
ax1.set_ylabel('Intensity')
ax1.set_title('Spectroscopic Data: 2 Gaussian Peaks + Linear Baseline')
ax1.legend(loc='best', fontsize=9)
ax1.grid(True, alpha=0.3)

# Residuals plot
ax2.scatter(x, residuals, s=3, alpha=0.5, color='gray')
ax2.axhline(y=0, color='r', linestyle='-', linewidth=1)
ax2.set_xlabel('X')
ax2.set_ylabel('Residuals')
ax2.set_title(f'Residuals (R²={r_squared:.6f}, RMSE={rmse:.6f})')
ax2.grid(True, alpha=0.3)

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

# Build results JSON
results = {
    "model_type": "2 Gaussian peaks + linear baseline: y = (m*x + b) + A1*exp(-(x-mu1)^2/(2*sigma1^2)) + A2*exp(-(x-mu2)^2/(2*sigma2^2))",
    "parameters": {
        "baseline": {
            "slope_m": round(float(m_fit), 6),
            "slope_m_err": round(float(m_err), 6),
            "intercept_b": round(float(b_fit), 6),
            "intercept_b_err": round(float(b_err), 6)
        },
        "peak_1": {
            "center": round(float(mu1_fit), 4),
            "center_err": round(float(mu1_err), 4),
            "amplitude": round(float(A1_fit), 4),
            "amplitude_err": round(float(A1_err), 4),
            "sigma": round(float(sigma1_fit), 4),
            "sigma_err": round(float(sigma1_err), 4),
            "fwhm": round(float(fwhm1), 4),
            "fwhm_err": round(float(fwhm1_err), 4)
        },
        "peak_2": {
            "center": round(float(mu2_fit), 4),
            "center_err": round(float(mu2_err), 4),
            "amplitude": round(float(A2_fit), 4),
            "amplitude_err": round(float(A2_err), 4),
            "sigma": round(float(sigma2_fit), 4),
            "sigma_err": round(float(sigma2_err), 4),
            "fwhm": round(float(fwhm2), 4),
            "fwhm_err": round(float(fwhm2_err), 4)
        }
    },
    "fit_quality": {
        "r_squared": round(float(r_squared), 6),
        "rmse": round(float(rmse), 6)
    },
    "summary": f"Fitted 2 Gaussian peaks on a linear baseline to 500-point spectroscopic data. Peak 1 centered at {mu1_fit:.2f} (FWHM={fwhm1:.2f}), Peak 2 centered at {mu2_fit:.2f} (FWHM={fwhm2:.2f}). Linear baseline slope={m_fit:.5f}. R²={r_squared:.6f}, RMSE={rmse:.6f}."
}

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