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

# ── 1. Load data ──────────────────────────────────────────────────────────────
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_components/temp_spectrum_0.npy')
if data.ndim == 2:
    x, y = data[:, 0], data[:, 1]
else:
    x = np.linspace(0, 100, len(data))
    y = data

# ── 2. Model: constant baseline + 4 Gaussians ─────────────────────────────────
def gaussian(x, mu, A, sigma):
    return A * np.exp(-0.5 * ((x - mu) / sigma) ** 2)

def model(x, baseline,
          mu1, A1, s1,
          mu2, A2, s2,
          mu3, A3, s3,
          mu4, A4, s4):
    return (baseline
            + gaussian(x, mu1, A1, s1)
            + gaussian(x, mu2, A2, s2)
            + gaussian(x, mu3, A3, s3)
            + gaussian(x, mu4, A4, s4))

# ── 3. Initial guesses & bounds ────────────────────────────────────────────────
p0 = [
    0.05,          # baseline
    20, 3.07, 2,   # peak 1
    42, 2.62, 3,   # peak 2
    65, 4.19, 2,   # peak 3
    85, 2.27, 3,   # peak 4
]

# lower / upper bounds  (baseline >= -0.1, A > 0, sigma > 0.1, 0 <= mu <= 100)
lb = [-0.1,
       0,  0,  0.1,
       0,  0,  0.1,
       0,  0,  0.1,
       0,  0,  0.1]
ub = [np.inf,
      100, np.inf, np.inf,
      100, np.inf, np.inf,
      100, np.inf, np.inf,
      100, np.inf, np.inf]

# ── 4. Fit ─────────────────────────────────────────────────────────────────────
try:
    popt, pcov = curve_fit(
        model, x, y,
        p0=p0,
        bounds=(lb, ub),
        method='trf',          # trf supports bounds; falls back gracefully
        maxfev=20000,
        ftol=1e-12, xtol=1e-12
    )
except RuntimeError:
    # Last-resort: try without tight tolerances
    popt, pcov = curve_fit(
        model, x, y,
        p0=p0,
        bounds=(lb, ub),
        method='trf',
        maxfev=50000
    )

perr = np.sqrt(np.diag(pcov))

# ── 5. Extract parameters ──────────────────────────────────────────────────────
FWHM_FACTOR = 2.0 * np.sqrt(2.0 * np.log(2.0))  # ≈ 2.3548

baseline_val  = popt[0];  baseline_err  = perr[0]
mu1, A1, s1   = popt[1], popt[2],  popt[3]
mu2, A2, s2   = popt[4], popt[5],  popt[6]
mu3, A3, s3   = popt[7], popt[8],  popt[9]
mu4, A4, s4   = popt[10], popt[11], popt[12]

e_mu1, e_A1, e_s1  = perr[1],  perr[2],  perr[3]
e_mu2, e_A2, e_s2  = perr[4],  perr[5],  perr[6]
e_mu3, e_A3, e_s3  = perr[7],  perr[8],  perr[9]
e_mu4, e_A4, e_s4  = perr[10], perr[11], perr[12]

# ── 6. Goodness of fit ─────────────────────────────────────────────────────────
y_fit    = model(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))
n_params  = len(popt)
reduced_chi2 = ss_res / (len(y) - n_params)

# ── 7. Visualization ───────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 1, figsize=(10, 7),
                          gridspec_kw={'height_ratios': [3, 1]})
fig.suptitle('4-Gaussian Fit with Constant Baseline', fontsize=13)

ax = axes[0]
ax.plot(x, y, 'k.', markersize=2, label='Data', alpha=0.6)
ax.plot(x, y_fit, 'r-', linewidth=2, label='Total fit')

# Individual components
colors = ['steelblue', 'darkorange', 'green', 'purple']
params_list = [(mu1, A1, s1), (mu2, A2, s2), (mu3, A3, s3), (mu4, A4, s4)]
for i, (mu, A, s) in enumerate(params_list):
    y_comp = baseline_val + gaussian(x, mu, A, s)
    ax.plot(x, y_comp, '--', color=colors[i], linewidth=1.5,
            label=f'Peak {i+1} (μ={mu:.1f})')

ax.axhline(baseline_val, color='gray', linestyle=':', linewidth=1,
           label=f'Baseline={baseline_val:.4f}')
ax.set_ylabel('Intensity')
ax.legend(fontsize=8, ncol=2)
ax.set_xlim(x[0], x[-1])
ax.grid(True, alpha=0.3)
ax.text(0.02, 0.97, f'R²={r_squared:.5f}   RMSE={rmse:.5f}',
        transform=ax.transAxes, va='top', fontsize=9,
        bbox=dict(boxstyle='round', fc='white', alpha=0.7))

ax2 = axes[1]
ax2.plot(x, residuals, 'k-', linewidth=0.8, alpha=0.7)
ax2.axhline(0, color='r', linewidth=1)
ax2.set_xlabel('x')
ax2.set_ylabel('Residuals')
ax2.set_xlim(x[0], x[-1])
ax2.grid(True, alpha=0.3)

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

# ── 8. Results JSON ────────────────────────────────────────────────────────────
def fmt(v):  return float(round(v, 6))

results = {
    "model_type": "4 Gaussian peaks on a constant baseline: y = C + sum_i A_i*exp(-0.5*((x-mu_i)/sigma_i)^2)",
    "parameters": {
        "baseline": {
            "value": fmt(baseline_val),
            "error": fmt(baseline_err)
        },
        "peak_1": {
            "center":     fmt(mu1),  "center_err":    fmt(e_mu1),
            "amplitude":  fmt(A1),   "amplitude_err": fmt(e_A1),
            "sigma":      fmt(s1),   "sigma_err":     fmt(e_s1),
            "fwhm":       fmt(FWHM_FACTOR * s1),
            "fwhm_err":   fmt(FWHM_FACTOR * e_s1)
        },
        "peak_2": {
            "center":     fmt(mu2),  "center_err":    fmt(e_mu2),
            "amplitude":  fmt(A2),   "amplitude_err": fmt(e_A2),
            "sigma":      fmt(s2),   "sigma_err":     fmt(e_s2),
            "fwhm":       fmt(FWHM_FACTOR * s2),
            "fwhm_err":   fmt(FWHM_FACTOR * e_s2)
        },
        "peak_3": {
            "center":     fmt(mu3),  "center_err":    fmt(e_mu3),
            "amplitude":  fmt(A3),   "amplitude_err": fmt(e_A3),
            "sigma":      fmt(s3),   "sigma_err":     fmt(e_s3),
            "fwhm":       fmt(FWHM_FACTOR * s3),
            "fwhm_err":   fmt(FWHM_FACTOR * e_s3)
        },
        "peak_4": {
            "center":     fmt(mu4),  "center_err":    fmt(e_mu4),
            "amplitude":  fmt(A4),   "amplitude_err": fmt(e_A4),
            "sigma":      fmt(s4),   "sigma_err":     fmt(e_s4),
            "fwhm":       fmt(FWHM_FACTOR * s4),
            "fwhm_err":   fmt(FWHM_FACTOR * e_s4)
        }
    },
    "fit_quality": {
        "r_squared":      fmt(r_squared),
        "rmse":           fmt(rmse),
        "reduced_chi2":   fmt(reduced_chi2),
        "n_data_points":  int(len(y)),
        "n_free_params":  int(n_params)
    },
    "summary": (
        f"Four Gaussian peaks fitted on a constant baseline (C={baseline_val:.4f}). "
        f"Peak centres: {mu1:.2f}, {mu2:.2f}, {mu3:.2f}, {mu4:.2f}. "
        f"Amplitudes: {A1:.3f}, {A2:.3f}, {A3:.3f}, {A4:.3f}. "
        f"FWHMs: {FWHM_FACTOR*s1:.2f}, {FWHM_FACTOR*s2:.2f}, "
        f"{FWHM_FACTOR*s3:.2f}, {FWHM_FACTOR*s4:.2f}. "
        f"R²={r_squared:.5f}, RMSE={rmse:.5f}."
    )
}

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