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_obvious/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

# Define model: linear baseline + 4 Gaussians (2 per peak: narrow + broad)
def gaussian(x, amplitude, center, sigma):
    return amplitude * np.exp(-(x - center)**2 / (2 * sigma**2))

def model(x, slope, intercept,
          p1n_amp, p1n_center, p1n_sigma,
          p1b_amp, p1b_center, p1b_sigma,
          p2n_amp, p2n_center, p2n_sigma,
          p2b_amp, p2b_center, p2b_sigma):
    baseline = slope * x + intercept
    g1n = gaussian(x, p1n_amp, p1n_center, p1n_sigma)
    g1b = gaussian(x, p1b_amp, p1b_center, p1b_sigma)
    g2n = gaussian(x, p2n_amp, p2n_center, p2n_sigma)
    g2b = gaussian(x, p2b_amp, p2b_center, p2b_sigma)
    return baseline + g1n + g1b + g2n + g2b

# Find approximate peak positions from the data
from scipy.signal import find_peaks as _fp
_peaks_idx, _peaks_props = _fp(y, height=np.max(y)*0.3, distance=50)
if len(_peaks_idx) >= 2:
    _sorted = np.argsort(y[_peaks_idx])[::-1]
    _p1_idx = _peaks_idx[_sorted[0]]
    _p2_idx = _peaks_idx[_sorted[1]]
    p1_center_guess = x[_p1_idx]
    p1_amp_guess = y[_p1_idx]
    p2_center_guess = x[_p2_idx]
    p2_amp_guess = y[_p2_idx]
else:
    p1_center_guess = 35.0
    p1_amp_guess = 4.0
    p2_center_guess = 72.0
    p2_amp_guess = 3.3

# Ensure p1 is the left peak
if p1_center_guess > p2_center_guess:
    p1_center_guess, p2_center_guess = p2_center_guess, p1_center_guess
    p1_amp_guess, p2_amp_guess = p2_amp_guess, p1_amp_guess

print(f"Initial guesses: peak1 at {p1_center_guess:.1f} (amp={p1_amp_guess:.2f}), peak2 at {p2_center_guess:.1f} (amp={p2_amp_guess:.2f})")

# Initial guesses: [slope, intercept, p1n_amp, p1n_center, p1n_sigma, p1b_amp, p1b_center, p1b_sigma,
#                   p2n_amp, p2n_center, p2n_sigma, p2b_amp, p2b_center, p2b_sigma]
p0 = [
    0.005, 0.0,
    p1_amp_guess * 0.8, p1_center_guess, 1.0,
    p1_amp_guess * 0.2, p1_center_guess, 4.0,
    p2_amp_guess * 0.8, p2_center_guess, 1.0,
    p2_amp_guess * 0.2, p2_center_guess, 4.0
]

# Bounds: narrow sigma [0.3, 3.0], broad sigma [2.0, 15.0], amplitudes positive
# Centers tied within ±2 of initial guess for each peak
lower = [
    -0.1, -5.0,
    0.0, p1_center_guess - 2.0, 0.3,
    0.0, p1_center_guess - 2.0, 2.0,
    0.0, p2_center_guess - 2.0, 0.3,
    0.0, p2_center_guess - 2.0, 2.0
]
upper = [
    0.1, 5.0,
    p1_amp_guess * 3, p1_center_guess + 2.0, 3.0,
    p1_amp_guess * 3, p1_center_guess + 2.0, 15.0,
    p2_amp_guess * 3, p2_center_guess + 2.0, 3.0,
    p2_amp_guess * 3, p2_center_guess + 2.0, 15.0
]

# Fit using Trust Region Reflective
try:
    popt, pcov = curve_fit(model, x, y, p0=p0, bounds=(lower, upper), method='trf', maxfev=50000)
    perr = np.sqrt(np.diag(pcov))
except Exception as e:
    print(f"Fitting failed: {e}")
    # Try with relaxed bounds
    lower_relaxed = [-1.0, -10.0] + [0.0, x.min(), 0.1, 0.0, x.min(), 0.5, 0.0, x.min(), 0.1, 0.0, x.min(), 0.5]
    upper_relaxed = [1.0, 10.0] + [20.0, x.max(), 10.0, 20.0, x.max(), 30.0, 20.0, x.max(), 10.0, 20.0, x.max(), 30.0]
    popt, pcov = curve_fit(model, x, y, p0=p0, bounds=(lower_relaxed, upper_relaxed), method='trf', maxfev=50000)
    perr = np.sqrt(np.diag(pcov))

# Extract parameters
slope, intercept = popt[0], popt[1]
p1n_amp, p1n_center, p1n_sigma = popt[2], popt[3], popt[4]
p1b_amp, p1b_center, p1b_sigma = popt[5], popt[6], popt[7]
p2n_amp, p2n_center, p2n_sigma = popt[8], popt[9], popt[10]
p2b_amp, p2b_center, p2b_sigma = popt[11], popt[12], popt[13]

slope_err, intercept_err = perr[0], perr[1]
p1n_amp_err, p1n_center_err, p1n_sigma_err = perr[2], perr[3], perr[4]
p1b_amp_err, p1b_center_err, p1b_sigma_err = perr[5], perr[6], perr[7]
p2n_amp_err, p2n_center_err, p2n_sigma_err = perr[8], perr[9], perr[10]
p2b_amp_err, p2b_center_err, p2b_sigma_err = perr[11], perr[12], perr[13]

# Compute composite peak properties
# Peak 1 effective center (amplitude-weighted)
p1_total_amp = p1n_amp + p1b_amp
p1_eff_center = (p1n_amp * p1n_center + p1b_amp * p1b_center) / p1_total_amp
# Total area = sum of areas (area of Gaussian = amplitude * sigma * sqrt(2*pi))
p1n_area = p1n_amp * p1n_sigma * np.sqrt(2 * np.pi)
p1b_area = p1b_amp * p1b_sigma * np.sqrt(2 * np.pi)
p1_total_area = p1n_area + p1b_area

p2_total_amp = p2n_amp + p2b_amp
p2_eff_center = (p2n_amp * p2n_center + p2b_amp * p2b_center) / p2_total_amp
p2n_area = p2n_amp * p2n_sigma * np.sqrt(2 * np.pi)
p2b_area = p2b_amp * p2b_sigma * np.sqrt(2 * np.pi)
p2_total_area = p2n_area + p2b_area

# FWHM of composite: evaluate numerically
x_fine = np.linspace(x.min(), x.max(), 10000)

def composite_peak1(xv):
    return gaussian(xv, p1n_amp, p1n_center, p1n_sigma) + gaussian(xv, p1b_amp, p1b_center, p1b_sigma)

def composite_peak2(xv):
    return gaussian(xv, p2n_amp, p2n_center, p2n_sigma) + gaussian(xv, p2b_amp, p2b_center, p2b_sigma)

def compute_fwhm(xv, yv):
    max_val = np.max(yv)
    half_max = max_val / 2.0
    above = np.where(yv >= half_max)[0]
    if len(above) > 0:
        return xv[above[-1]] - xv[above[0]]
    return 0.0

p1_composite_y = composite_peak1(x_fine)
p2_composite_y = composite_peak2(x_fine)
p1_fwhm = compute_fwhm(x_fine, p1_composite_y)
p2_fwhm = compute_fwhm(x_fine, p2_composite_y)

# Compute fit quality
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 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals**2))

print(f"R² = {r_squared:.6f}, RMSE = {rmse:.6f}")

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

# Main plot
ax1.plot(x, y, 'k.', markersize=1, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label=f'Total Fit (R²={r_squared:.5f})')

# Baseline
baseline_y = slope * x + intercept
ax1.plot(x, baseline_y, '--', color='gray', linewidth=1, alpha=0.7, label='Baseline')

# Individual components
g1n_y = gaussian(x, p1n_amp, p1n_center, p1n_sigma) + baseline_y
g1b_y = gaussian(x, p1b_amp, p1b_center, p1b_sigma) + baseline_y
g2n_y = gaussian(x, p2n_amp, p2n_center, p2n_sigma) + baseline_y
g2b_y = gaussian(x, p2b_amp, p2b_center, p2b_sigma) + baseline_y

# Composite peaks
peak1_y = composite_peak1(x) + baseline_y
peak2_y = composite_peak2(x) + baseline_y

ax1.fill_between(x, baseline_y, peak1_y, alpha=0.2, color='blue', label='Peak 1 (composite)')
ax1.fill_between(x, baseline_y, peak2_y, alpha=0.2, color='green', label='Peak 2 (composite)')

ax1.plot(x, g1n_y, 'b--', linewidth=1, alpha=0.6, label='Peak 1 narrow')
ax1.plot(x, g1b_y, 'b:', linewidth=1, alpha=0.6, label='Peak 1 broad')
ax1.plot(x, g2n_y, 'g--', linewidth=1, alpha=0.6, label='Peak 2 narrow')
ax1.plot(x, g2b_y, 'g:', linewidth=1, alpha=0.6, label='Peak 2 broad')

ax1.set_ylabel('Intensity')
ax1.set_title('Spectroscopic Data Fit: 4 Gaussians (2 per peak) + Linear Baseline')
ax1.legend(fontsize=8, loc='upper right')
ax1.grid(True, alpha=0.3)

# Residuals
ax2.plot(x, residuals, 'k.', markersize=1, alpha=0.5)
ax2.axhline(y=0, color='r', linestyle='-', linewidth=0.5)
ax2.set_xlabel('X')
ax2.set_ylabel('Residuals')
ax2.set_title(f'Residuals (RMSE={rmse:.5f})')
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 (2 per peak: narrow + broad) on linear baseline",
    "parameters": {
        "peak_1": {
            "center": round(float(p1_eff_center), 4),
            "center_err": round(float(np.sqrt((p1n_amp*p1n_center_err)**2 + (p1b_amp*p1b_center_err)**2) / p1_total_amp), 4),
            "narrow_amplitude": round(float(p1n_amp), 4),
            "narrow_amplitude_err": round(float(p1n_amp_err), 4),
            "narrow_sigma": round(float(p1n_sigma), 4),
            "narrow_sigma_err": round(float(p1n_sigma_err), 4),
            "broad_amplitude": round(float(p1b_amp), 4),
            "broad_amplitude_err": round(float(p1b_amp_err), 4),
            "broad_sigma": round(float(p1b_sigma), 4),
            "broad_sigma_err": round(float(p1b_sigma_err), 4),
            "total_area": round(float(p1_total_area), 4),
            "fwhm": round(float(p1_fwhm), 4)
        },
        "peak_2": {
            "center": round(float(p2_eff_center), 4),
            "center_err": round(float(np.sqrt((p2n_amp*p2n_center_err)**2 + (p2b_amp*p2b_center_err)**2) / p2_total_amp), 4),
            "narrow_amplitude": round(float(p2n_amp), 4),
            "narrow_amplitude_err": round(float(p2n_amp_err), 4),
            "narrow_sigma": round(float(p2n_sigma), 4),
            "narrow_sigma_err": round(float(p2n_sigma_err), 4),
            "broad_amplitude": round(float(p2b_amp), 4),
            "broad_amplitude_err": round(float(p2b_amp_err), 4),
            "broad_sigma": round(float(p2b_sigma), 4),
            "broad_sigma_err": round(float(p2b_sigma_err), 4),
            "total_area": round(float(p2_total_area), 4),
            "fwhm": round(float(p2_fwhm), 4)
        },
        "baseline": {
            "slope": round(float(slope), 6),
            "slope_err": round(float(slope_err), 6),
            "intercept": round(float(intercept), 6),
            "intercept_err": round(float(intercept_err), 6)
        }
    },
    "fit_quality": {
        "r_squared": round(float(r_squared), 6),
        "rmse": round(float(rmse), 6)
    },
    "summary": f"Two spectroscopic peaks fitted with 4 Gaussians (2 per peak: narrow core + broad wing) on a linear baseline. Peak 1 at {p1_eff_center:.2f} (FWHM={p1_fwhm:.2f}, area={p1_total_area:.2f}), Peak 2 at {p2_eff_center:.2f} (FWHM={p2_fwhm:.2f}, area={p2_total_area:.2f}). R²={r_squared:.5f}, RMSE={rmse:.5f}. All peaks use pure Gaussian profiles per domain rules."
}

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