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

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]
else:
    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 the model: linear baseline + 4 Gaussians
def model_func(x, a0, a1, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3, A4, mu4, sigma4):
    baseline = a0 + a1 * x
    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

# Initial guesses
p0 = [
    0.1, 0.0,       # a0, a1
    3.0, 20.0, 4.0, # A1, mu1, sigma1
    2.6, 42.0, 4.0, # A2, mu2, sigma2
    4.2, 65.0, 3.0, # A3, mu3, sigma3
    2.2, 88.0, 3.0  # A4, mu4, sigma4
]

# Bounds as specified
lower_bounds = [
    -np.inf, -np.inf,   # a0, a1
    0.01, 10.0, 0.5,    # A1, mu1, sigma1
    0.01, 32.0, 0.5,    # A2, mu2, sigma2
    0.01, 55.0, 0.5,    # A3, mu3, sigma3
    0.01, 78.0, 0.5     # A4, mu4, sigma4
]

upper_bounds = [
    np.inf, np.inf,      # a0, a1
    np.inf, 30.0, 20.0,  # A1, mu1, sigma1
    np.inf, 52.0, 20.0,  # A2, mu2, sigma2
    np.inf, 75.0, 20.0,  # A3, mu3, sigma3
    np.inf, 98.0, 20.0   # A4, mu4, sigma4
]

bounds = (lower_bounds, upper_bounds)

# Stage 1: fit with Trust Region Reflective (handles bounds robustly)
try:
    popt, pcov = curve_fit(model_func, x, y, p0=p0, bounds=bounds, method='trf', maxfev=50000)
except Exception:
    # Fallback: try with Levenberg-Marquardt (no bounds)
    try:
        popt, pcov = curve_fit(model_func, x, y, p0=p0, method='lm', maxfev=50000)
    except Exception:
        popt = np.array(p0)
        pcov = np.zeros((len(p0), len(p0)))

# Stage 2: refine all 14 parameters simultaneously using stage 1 result as initial guess
try:
    popt2, pcov2 = curve_fit(model_func, x, y, p0=popt, bounds=bounds, method='trf', maxfev=100000)
    popt = popt2
    pcov = pcov2
except Exception:
    pass  # keep stage 1 result

# Extract parameter uncertainties
perr = np.sqrt(np.diag(np.abs(pcov)))

# Compute fit and residuals
y_fit = model_func(x, *popt)
residuals = y - y_fit

# R-squared
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0

# RMSE
rmse = np.sqrt(np.mean(residuals**2))

# Extract parameters
a0, a1 = popt[0], popt[1]
a0_err, a1_err = perr[0], perr[1]

peak_params = []
peak_names = ['peak_1', 'peak_2', 'peak_3', 'peak_4']
for i in range(4):
    idx = 2 + i * 3
    A = popt[idx]
    mu = popt[idx + 1]
    sigma = abs(popt[idx + 2])
    A_err = perr[idx]
    mu_err = perr[idx + 1]
    sigma_err = perr[idx + 2]
    area = A * sigma * np.sqrt(2 * np.pi)
    # Error propagation for area: area = A * sigma * sqrt(2pi)
    # d(area)/dA = sigma * sqrt(2pi), d(area)/dsigma = A * sqrt(2pi)
    area_err = np.sqrt((sigma * np.sqrt(2 * np.pi) * A_err)**2 + (A * np.sqrt(2 * np.pi) * sigma_err)**2)
    peak_params.append({
        'name': peak_names[i],
        'amplitude': float(A),
        'amplitude_err': float(A_err),
        'center': float(mu),
        'center_err': float(mu_err),
        'sigma': float(sigma),
        'sigma_err': float(sigma_err),
        'fwhm': float(2.3548 * sigma),
        'area': float(area),
        'area_err': float(area_err)
    })

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

# Main plot
ax1.scatter(x, y, s=2, alpha=0.5, color='gray', label='Data', zorder=1)
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit', zorder=5)

# Plot baseline
baseline = a0 + a1 * x
ax1.plot(x, baseline, 'k--', linewidth=1, alpha=0.7, label='Baseline')

# Plot individual Gaussian components
colors = ['blue', 'green', 'orange', 'purple']
for i in range(4):
    idx = 2 + i * 3
    A = popt[idx]
    mu = popt[idx + 1]
    sigma = popt[idx + 2]
    g_i = A * np.exp(-(x - mu)**2 / (2 * sigma**2))
    ax1.plot(x, baseline + g_i, '--', color=colors[i], linewidth=1.5, alpha=0.8, 
             label=f'Peak {i+1} (center={mu:.1f})')
    ax1.fill_between(x, baseline, baseline + g_i, alpha=0.15, color=colors[i])

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

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

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

# Build results JSON
params_dict = {
    'baseline': {
        'a0': float(a0), 'a0_err': float(a0_err),
        'a1': float(a1), 'a1_err': float(a1_err)
    }
}
for pp in peak_params:
    params_dict[pp['name']] = {
        'center': pp['center'],
        'center_err': pp['center_err'],
        'amplitude': pp['amplitude'],
        'amplitude_err': pp['amplitude_err'],
        'sigma': pp['sigma'],
        'sigma_err': pp['sigma_err'],
        'fwhm': pp['fwhm'],
        'area': pp['area'],
        'area_err': pp['area_err']
    }

# Determine dominant peak
areas = [pp['area'] for pp in peak_params]
max_area_idx = np.argmax(areas)
dominant = peak_params[max_area_idx]

summary = (f"4-Gaussian + linear baseline fit achieved R²={r_squared:.6f}, RMSE={rmse:.6f}. "
           f"Peak centers at {peak_params[0]['center']:.2f}, {peak_params[1]['center']:.2f}, "
           f"{peak_params[2]['center']:.2f}, {peak_params[3]['center']:.2f}. "
           f"Dominant peak is Peak {max_area_idx+1} (center={dominant['center']:.2f}, "
           f"amplitude={dominant['amplitude']:.3f}, area={dominant['area']:.3f}). "
           f"Relative peak areas: {', '.join([f'Peak {i+1}={a:.3f}' for i, a in enumerate(areas)])}. "
           f"{'Fit quality is excellent (R²>0.94).' if r_squared > 0.94 else 'Fit quality is below target (R²<0.94); consider additional components or non-linear baseline.'}")

results = {
    'model_type': 'Linear baseline + 4 Gaussian peaks (y = a0+a1*x + sum_i A_i*exp(-(x-mu_i)^2/(2*sigma_i^2)))',
    'parameters': params_dict,
    'fit_quality': {
        'r_squared': float(r_squared),
        'rmse': float(rmse)
    },
    'summary': summary
}

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