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_regime_data/output/temp_spectrum_2.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: 2 Lorentzian peaks + linear baseline
# y = a0 + a1*x + A1*gamma1^2/((x-x01)^2+gamma1^2) + A2*gamma2^2/((x-x02)^2+gamma2^2)
def lorentzian(x, x0, A, gamma):
    return A * gamma**2 / ((x - x0)**2 + gamma**2)

def model_2lorentz_linear(x, a0, a1, x01, A1, gamma1, x02, A2, gamma2):
    baseline = a0 + a1 * x
    peak1 = lorentzian(x, x01, A1, gamma1)
    peak2 = lorentzian(x, x02, A2, gamma2)
    return baseline + peak1 + peak2

# Initial guesses as specified in plan
# peak1: x0=35, A=4.0, gamma=2.0
# peak2: x0=71, A=3.3, gamma=1.5
# baseline: intercept=0, slope=0.005
p0 = [0.0, 0.005, 35.0, 4.0, 2.0, 71.0, 3.3, 1.5]

# Bounds as specified:
# a0: (-inf, inf), a1: (-inf, inf)
# x01: [30, 40], A1: [0, inf], gamma1: [0.5, 10]
# x02: [66, 76], A2: [0, inf], gamma2: [0.5, 10]
lower = [-np.inf, -np.inf, 30.0, 0.0, 0.5, 66.0, 0.0, 0.5]
upper = [np.inf, np.inf, 40.0, np.inf, 10.0, 76.0, np.inf, 10.0]

# Fit using Levenberg-Marquardt (method='lm' doesn't support bounds, use 'trf' with bounds)
try:
    popt, pcov = curve_fit(model_2lorentz_linear, x, y, p0=p0, bounds=(lower, upper), method='trf', maxfev=50000)
except Exception as e:
    print(f"Fit failed with trf, trying without bounds: {e}")
    popt, pcov = curve_fit(model_2lorentz_linear, x, y, p0=p0, method='lm', maxfev=50000)

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

# Extract parameters
a0, a1, x01, A1, gamma1, x02, A2, gamma2 = popt
a0_err, a1_err, x01_err, A1_err, gamma1_err, x02_err, A2_err, gamma2_err = perr

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

# Compute 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))

# Compute individual components for visualization
baseline = a0 + a1 * x
peak1 = lorentzian(x, x01, A1, gamma1)
peak2 = lorentzian(x, x02, A2, gamma2)

# Compute FWHM for Lorentzian: FWHM = 2*gamma
fwhm1 = 2.0 * abs(gamma1)
fwhm1_err = 2.0 * gamma1_err
fwhm2 = 2.0 * abs(gamma2)
fwhm2_err = 2.0 * gamma2_err

# Save 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, alpha=0.7, label='Baseline')
ax1.plot(x, baseline + peak1, 'b-', linewidth=1.5, alpha=0.7, label=f'Peak 1 (center={x01:.2f})')
ax1.plot(x, baseline + peak2, 'g-', linewidth=1.5, alpha=0.7, label=f'Peak 2 (center={x02:.2f})')

# Fill individual peaks on baseline
ax1.fill_between(x, baseline, baseline + peak1, alpha=0.2, color='blue')
ax1.fill_between(x, baseline, baseline + peak2, alpha=0.2, color='green')

ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title('Spectroscopic Data: 2 Lorentzian Peaks + Linear Baseline', fontsize=14)
ax1.legend(fontsize=9, loc='upper right')
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', fontsize=12)
ax2.set_ylabel('Residual', fontsize=12)
ax2.set_title(f'Residuals (RMSE={rmse:.6f})', fontsize=11)
ax2.grid(True, alpha=0.3)

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

# Build results JSON
results = {
    "model_type": "2 Lorentzian peaks + linear baseline: y = a0 + a1*x + A1*gamma1^2/((x-x01)^2+gamma1^2) + A2*gamma2^2/((x-x02)^2+gamma2^2)",
    "parameters": {
        "baseline": {
            "intercept": round(float(a0), 6),
            "intercept_err": round(float(a0_err), 6),
            "slope": round(float(a1), 6),
            "slope_err": round(float(a1_err), 6)
        },
        "peak_1": {
            "center": round(float(x01), 4),
            "center_err": round(float(x01_err), 4),
            "amplitude": round(float(A1), 4),
            "amplitude_err": round(float(A1_err), 4),
            "width_gamma": round(float(gamma1), 4),
            "width_gamma_err": round(float(gamma1_err), 4),
            "fwhm": round(float(fwhm1), 4),
            "fwhm_err": round(float(fwhm1_err), 4)
        },
        "peak_2": {
            "center": round(float(x02), 4),
            "center_err": round(float(x02_err), 4),
            "amplitude": round(float(A2), 4),
            "amplitude_err": round(float(A2_err), 4),
            "width_gamma": round(float(gamma2), 4),
            "width_gamma_err": round(float(gamma2_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"Two Lorentzian peaks on a linear baseline fit to spectroscopic data (500 points, x=[0,100]). Peak 1 centered at {x01:.2f} with FWHM={fwhm1:.2f}, Peak 2 centered at {x02:.2f} with FWHM={fwhm2:.2f}. R²={r_squared:.6f}, RMSE={rmse:.6f}. Levenberg-Marquardt optimization (trf with bounds) used. Centers constrained within ±5 of initial guesses, widths in [0.5,10], amplitudes > 0."
}

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