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_validator_data/v_subtle/temp_spectrum_0.npy')
x = data[:, 0] if data.ndim == 2 else np.linspace(0, 100, len(data))
y = data[:, 1] if data.ndim == 2 else data

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

# ============================================================
# Model: Two Gaussian peaks on a linear baseline
# y(x) = a0 + a1*x + G1(x) + G2(x)
# where Gi(x) = amp_i * exp(-(x - cen_i)^2 / (2 * sig_i^2))
# ============================================================

def gaussian(x, amp, cen, sig):
    return amp * np.exp(-(x - cen)**2 / (2 * sig**2))

def model_2gauss_linear(x, a0, a1, amp1, cen1, sig1, amp2, cen2, sig2):
    baseline = a0 + a1 * x
    g1 = gaussian(x, amp1, cen1, sig1)
    g2 = gaussian(x, amp2, cen2, sig2)
    return baseline + g1 + g2

# --- Initial guess estimation ---
# Estimate baseline from endpoints
y_start = np.mean(y[:20])
y_end = np.mean(y[-20:])
a1_guess = (y_end - y_start) / (x[-1] - x[0])
a0_guess = y_start - a1_guess * x[0]

# Subtract baseline to find peaks
baseline_guess = a0_guess + a1_guess * x
y_detrended = y - baseline_guess

# Find the two highest peaks in the detrended data
from scipy.signal import find_peaks as _find_peaks

# Smooth slightly for peak finding
from scipy.ndimage import uniform_filter1d
y_smooth = uniform_filter1d(y_detrended, size=20)

peak_indices, peak_props = _find_peaks(y_smooth, height=0, distance=50)
if len(peak_indices) >= 2:
    heights = y_smooth[peak_indices]
    top2 = np.argsort(heights)[-2:]
    idx1, idx2 = sorted(peak_indices[top2])
else:
    # Fallback: split domain in two halves
    mid = len(x) // 2
    idx1 = np.argmax(y_detrended[:mid])
    idx2 = mid + np.argmax(y_detrended[mid:])

cen1_guess = x[idx1]
cen2_guess = x[idx2]
amp1_guess = y_detrended[idx1]
amp2_guess = y_detrended[idx2]
sig1_guess = (x[-1] - x[0]) / 10
sig2_guess = (x[-1] - x[0]) / 10

p0 = [a0_guess, a1_guess, amp1_guess, cen1_guess, sig1_guess, amp2_guess, cen2_guess, sig2_guess]

# Bounds
lb = [-np.inf, -np.inf, 0, x[0], 0.1, 0, x[0], 0.1]
ub = [np.inf, np.inf, np.inf, x[-1], (x[-1]-x[0]), np.inf, x[-1], (x[-1]-x[0])]

try:
    popt, pcov = curve_fit(model_2gauss_linear, x, y, p0=p0, bounds=(lb, ub), maxfev=50000)
    perr = np.sqrt(np.diag(pcov))
except Exception as e:
    print(f"First fit attempt failed: {e}, trying without bounds")
    popt, pcov = curve_fit(model_2gauss_linear, x, y, p0=p0, maxfev=50000)
    perr = np.sqrt(np.diag(pcov))

a0_f, a1_f, amp1_f, cen1_f, sig1_f, amp2_f, cen2_f, sig2_f = popt
a0_e, a1_e, amp1_e, cen1_e, sig1_e, amp2_e, cen2_e, sig2_e = perr

y_fit = model_2gauss_linear(x, *popt)
residuals = y - y_fit

# --- Check if quadratic baseline improves the fit ---
def model_2gauss_quad(x, a0, a1, a2, amp1, cen1, sig1, amp2, cen2, sig2):
    baseline = a0 + a1 * x + a2 * x**2
    g1 = gaussian(x, amp1, cen1, sig1)
    g2 = gaussian(x, amp2, cen2, sig2)
    return baseline + g1 + g2

# Check for systematic curvature in residuals
resid_poly = np.polyfit(x, residuals, 2)
resid_curvature = abs(resid_poly[0])

use_quadratic = False
if resid_curvature > 1e-6:
    p0_quad = [a0_f, a1_f, 0.0, amp1_f, cen1_f, sig1_f, amp2_f, cen2_f, sig2_f]
    lb_q = [-np.inf, -np.inf, -np.inf, 0, x[0], 0.1, 0, x[0], 0.1]
    ub_q = [np.inf, np.inf, np.inf, np.inf, x[-1], (x[-1]-x[0]), np.inf, x[-1], (x[-1]-x[0])]
    try:
        popt_q, pcov_q = curve_fit(model_2gauss_quad, x, y, p0=p0_quad, bounds=(lb_q, ub_q), maxfev=50000)
        perr_q = np.sqrt(np.diag(pcov_q))
        y_fit_q = model_2gauss_quad(x, *popt_q)
        ss_res_q = np.sum((y - y_fit_q)**2)
        ss_res_l = np.sum(residuals**2)
        # Use quadratic if it reduces SS_res by more than ~1%
        if ss_res_q < ss_res_l * 0.99:
            use_quadratic = True
    except Exception:
        pass

if use_quadratic:
    popt_final = popt_q
    perr_final = perr_q
    y_fit = y_fit_q
    residuals = y - y_fit
    a0_f, a1_f, a2_f, amp1_f, cen1_f, sig1_f, amp2_f, cen2_f, sig2_f = popt_final
    a0_e, a1_e, a2_e, amp1_e, cen1_e, sig1_e, amp2_e, cen2_e, sig2_e = perr_final
    model_desc = "two_gaussians_quadratic_baseline"
else:
    popt_final = popt
    perr_final = perr
    a2_f = 0.0
    a2_e = 0.0
    model_desc = "two_gaussians_linear_baseline"

# Compute fit quality
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals**2))

# Ensure peak_1 is the one with smaller center
if cen1_f > cen2_f:
    amp1_f, amp2_f = amp2_f, amp1_f
    cen1_f, cen2_f = cen2_f, cen1_f
    sig1_f, sig2_f = sig2_f, sig1_f
    amp1_e, amp2_e = amp2_e, amp1_e
    cen1_e, cen2_e = cen2_e, cen1_e
    sig1_e, sig2_e = sig2_e, sig1_e

# FWHM = 2*sqrt(2*ln2)*sigma ~ 2.3548*sigma
fwhm1 = 2.3548200 * sig1_f
fwhm1_err = 2.3548200 * sig1_e
fwhm2 = 2.3548200 * sig2_f
fwhm2_err = 2.3548200 * sig2_e

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

ax1.plot(x, y, 'k.', markersize=1, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit')

# Plot baseline
if use_quadratic:
    baseline_curve = a0_f + a1_f * x + a2_f * x**2
else:
    baseline_curve = a0_f + a1_f * x
ax1.plot(x, baseline_curve, 'b--', linewidth=1, label='Baseline')

# Plot individual Gaussians on baseline
g1_curve = gaussian(x, amp1_f, cen1_f, sig1_f)
g2_curve = gaussian(x, amp2_f, cen2_f, sig2_f)
ax1.fill_between(x, baseline_curve, baseline_curve + g1_curve, alpha=0.3, color='green', label=f'Peak 1 (cen={cen1_f:.2f})')
ax1.fill_between(x, baseline_curve, baseline_curve + g2_curve, alpha=0.3, color='orange', label=f'Peak 2 (cen={cen2_f:.2f})')

ax1.set_ylabel('Intensity')
ax1.set_title(f'Spectroscopic Data Fit — R²={r_squared:.6f}, RMSE={rmse:.6f}')
ax1.legend(fontsize=9)

ax2.plot(x, residuals, 'k.', markersize=1, alpha=0.5)
ax2.axhline(0, color='r', linestyle='--', linewidth=1)
ax2.set_xlabel('X')
ax2.set_ylabel('Residuals')
ax2.set_title('Residuals')

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

# ============================================================
# Results JSON
# ============================================================
params = {
    "baseline": {
        "intercept_a0": round(float(a0_f), 8),
        "intercept_a0_err": round(float(a0_e), 8),
        "slope_a1": round(float(a1_f), 8),
        "slope_a1_err": round(float(a1_e), 8)
    },
    "peak_1": {
        "center": round(float(cen1_f), 6),
        "center_err": round(float(cen1_e), 6),
        "amplitude": round(float(amp1_f), 6),
        "amplitude_err": round(float(amp1_e), 6),
        "sigma": round(float(sig1_f), 6),
        "sigma_err": round(float(sig1_e), 6),
        "fwhm": round(float(fwhm1), 6),
        "fwhm_err": round(float(fwhm1_err), 6)
    },
    "peak_2": {
        "center": round(float(cen2_f), 6),
        "center_err": round(float(cen2_e), 6),
        "amplitude": round(float(amp2_f), 6),
        "amplitude_err": round(float(amp2_e), 6),
        "sigma": round(float(sig2_f), 6),
        "sigma_err": round(float(sig2_e), 6),
        "fwhm": round(float(fwhm2), 6),
        "fwhm_err": round(float(fwhm2_err), 6)
    }
}

if use_quadratic:
    params["baseline"]["quadratic_a2"] = round(float(a2_f), 10)
    params["baseline"]["quadratic_a2_err"] = round(float(a2_e), 10)

baseline_type = "quadratic" if use_quadratic else "linear"
summary = (f"Fitted two Gaussian peaks on a {baseline_type} baseline. "
           f"Peak 1 centered at {cen1_f:.2f} (FWHM={fwhm1:.2f}), "
           f"Peak 2 centered at {cen2_f:.2f} (FWHM={fwhm2:.2f}). "
           f"R²={r_squared:.6f}, RMSE={rmse:.6f}. "
           f"{'Quadratic baseline used due to curvature in residuals from linear fit.' if use_quadratic else 'Linear baseline sufficient; no systematic curvature detected in residuals.'}")

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

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