import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import least_squares

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_validator_data/v_noplot/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]

# Model: linear baseline + 2 Lorentzians
# params = [a0, a1, A1, x01, gamma1, A2, x02, gamma2]
def lorentzian(x, A, x0, gamma):
    return A / (1.0 + ((x - x0) / gamma) ** 2)

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

def residuals(params, x, y):
    return y - model(params, x)

# Estimate initial guesses from data
# Find the two highest peaks
from scipy.signal import find_peaks
peaks_idx, properties = find_peaks(y, height=np.max(y)*0.2, distance=50)
if len(peaks_idx) >= 2:
    # Sort by height descending
    sorted_peaks = sorted(peaks_idx, key=lambda i: y[i], reverse=True)
    p1_idx = sorted_peaks[0]
    p2_idx = sorted_peaks[1]
    if x[p1_idx] > x[p2_idx]:
        p1_idx, p2_idx = p2_idx, p1_idx
    x01_init = x[p1_idx]
    x02_init = x[p2_idx]
    A1_init = y[p1_idx] - np.min(y)
    A2_init = y[p2_idx] - np.min(y)
else:
    x01_init = 35.0
    x02_init = 72.0
    A1_init = 4.0
    A2_init = 3.0

# Baseline guess from edges
y_left = np.mean(y[:20])
y_right = np.mean(y[-20:])
a1_init = (y_right - y_left) / (x[-1] - x[0])
a0_init = y_left - a1_init * x[0]

gamma1_init = 2.0
gamma2_init = 1.5

p0 = [a0_init, a1_init, A1_init, x01_init, gamma1_init, A2_init, x02_init, gamma2_init]

# Bounds: [a0, a1, A1, x01, gamma1, A2, x02, gamma2]
lb = [-np.inf, -np.inf, 0.0, x01_init - 10.0, 0.01, 0.0, x02_init - 10.0, 0.01]
ub = [np.inf, np.inf, np.inf, x01_init + 10.0, 50.0, np.inf, x02_init + 10.0, 50.0]

result = least_squares(residuals, p0, args=(x, y), bounds=(lb, ub), method='trf', max_nfev=50000)

popt = result.x
a0, a1, A1, x01, gamma1, A2, x02, gamma2 = popt

# Compute parameter uncertainties from Jacobian
J = result.jac
resid = result.fun
n = len(y)
p = len(popt)
if n > p:
    s_sq = np.sum(resid**2) / (n - p)
    try:
        cov = np.linalg.inv(J.T @ J) * s_sq
        perr = np.sqrt(np.diag(cov))
    except np.linalg.LinAlgError:
        perr = np.zeros(len(popt))
else:
    perr = np.zeros(len(popt))

y_fit = model(popt, x)
resid_vals = y - y_fit

# Fit quality
ss_res = np.sum(resid_vals ** 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 = np.sqrt(np.mean(resid_vals ** 2))

# Components
baseline = a0 + a1 * x
peak1 = lorentzian(x, A1, x01, gamma1)
peak2 = lorentzian(x, A2, x02, gamma2)

# FWHM = 2 * gamma for Lorentzian
fwhm1 = 2.0 * abs(gamma1)
fwhm2 = 2.0 * abs(gamma2)
fwhm1_err = 2.0 * perr[4]
fwhm2_err = 2.0 * perr[7]

# Plot
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='Fit (2 Lorentzians + linear baseline)')
ax1.plot(x, baseline, 'b--', linewidth=1, alpha=0.7, label='Baseline')
ax1.plot(x, baseline + peak1, 'g--', linewidth=1, alpha=0.7, label=f'Peak 1 (center={x01:.2f})')
ax1.plot(x, baseline + peak2, 'm--', linewidth=1, alpha=0.7, label=f'Peak 2 (center={x02:.2f})')
ax1.fill_between(x, baseline, baseline + peak1, alpha=0.15, color='green')
ax1.fill_between(x, baseline, baseline + peak2, alpha=0.15, color='magenta')
ax1.set_ylabel('Intensity')
ax1.set_title('Spectroscopic Data: 2-Lorentzian Fit with Linear Baseline')
ax1.legend(fontsize=8)
ax1.grid(True, alpha=0.3)

ax2.plot(x, resid_vals, '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(f'Residuals (RMSE={rmse:.6f}, R²={r_squared:.6f})')
ax2.grid(True, alpha=0.3)

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

# Ensure peak1 is the one with smaller center
if x01 > x02:
    # Swap labels for output
    results = {
        "model_type": "2 Lorentzian peaks + linear baseline",
        "parameters": {
            "baseline": {
                "intercept": round(float(a0), 6),
                "intercept_err": round(float(perr[0]), 6),
                "slope": round(float(a1), 6),
                "slope_err": round(float(perr[1]), 6)
            },
            "peak_1": {
                "center": round(float(x02), 4),
                "center_err": round(float(perr[6]), 4),
                "amplitude": round(float(A2), 4),
                "amplitude_err": round(float(perr[5]), 4),
                "hwhm": round(float(abs(gamma2)), 4),
                "hwhm_err": round(float(perr[7]), 4),
                "fwhm": round(float(fwhm2), 4),
                "fwhm_err": round(float(fwhm2_err), 4)
            },
            "peak_2": {
                "center": round(float(x01), 4),
                "center_err": round(float(perr[3]), 4),
                "amplitude": round(float(A1), 4),
                "amplitude_err": round(float(perr[2]), 4),
                "hwhm": round(float(abs(gamma1)), 4),
                "hwhm_err": round(float(perr[4]), 4),
                "fwhm": round(float(fwhm1), 4),
                "fwhm_err": round(float(fwhm1_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. Peak 1 at x={x02:.2f} (FWHM={fwhm2:.2f}), Peak 2 at x={x01:.2f} (FWHM={fwhm1:.2f}). R²={r_squared:.6f}, RMSE={rmse:.6f}. Lorentzian profiles were chosen due to the sharp peaks with extended wing structure typical of spectroscopic data."
    }
else:
    results = {
        "model_type": "2 Lorentzian peaks + linear baseline",
        "parameters": {
            "baseline": {
                "intercept": round(float(a0), 6),
                "intercept_err": round(float(perr[0]), 6),
                "slope": round(float(a1), 6),
                "slope_err": round(float(perr[1]), 6)
            },
            "peak_1": {
                "center": round(float(x01), 4),
                "center_err": round(float(perr[3]), 4),
                "amplitude": round(float(A1), 4),
                "amplitude_err": round(float(perr[2]), 4),
                "hwhm": round(float(abs(gamma1)), 4),
                "hwhm_err": round(float(perr[4]), 4),
                "fwhm": round(float(fwhm1), 4),
                "fwhm_err": round(float(fwhm1_err), 4)
            },
            "peak_2": {
                "center": round(float(x02), 4),
                "center_err": round(float(perr[6]), 4),
                "amplitude": round(float(A2), 4),
                "amplitude_err": round(float(perr[5]), 4),
                "hwhm": round(float(abs(gamma2)), 4),
                "hwhm_err": round(float(perr[7]), 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. Peak 1 at x={x01:.2f} (FWHM={fwhm1:.2f}), Peak 2 at x={x02:.2f} (FWHM={fwhm2:.2f}). R²={r_squared:.6f}, RMSE={rmse:.6f}. Lorentzian profiles were chosen due to the sharp peaks with extended wing structure typical of spectroscopic data."
    }

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