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

# Load data
data = np.load("/Users/maxim.ziatdinov/Code/SciLink/test_series_data/output/temp_spectrum_2.npy")
if data.ndim == 2:
    x = data[:, 0]
    y = data[:, 1]
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]

# --- Step 1: Estimate baseline from wing regions ---
wing_mask = (x < 10) | (x > 90)
if np.sum(wing_mask) > 2:
    slope_res = linregress(x[wing_mask], y[wing_mask])
    m_init = slope_res.slope
    b_init = slope_res.intercept
else:
    m_init = 0.0
    b_init = np.median(y[:10])

# --- Step 2: Estimate peak parameters ---
baseline_est = m_init * x + b_init
y_sub = y - baseline_est

# Peak 1 region: x in [20, 45]
mask1 = (x >= 20) & (x <= 45)
if np.sum(mask1) > 0:
    idx1 = np.where(mask1)[0][np.argmax(y_sub[mask1])]
    mu1_init = x[idx1]
    A1_init = y_sub[idx1]
else:
    mu1_init = 30.0
    A1_init = 3.0

# Peak 2 region: x in [55, 80]
mask2 = (x >= 55) & (x <= 80)
if np.sum(mask2) > 0:
    idx2 = np.where(mask2)[0][np.argmax(y_sub[mask2])]
    mu2_init = x[idx2]
    A2_init = y_sub[idx2]
else:
    mu2_init = 65.0
    A2_init = 2.0

gamma1_init = 2.0
gamma2_init = 2.0

# Ensure positive amplitudes
A1_init = max(A1_init, 0.5)
A2_init = max(A2_init, 0.5)

print(f"Initial guesses: m={m_init:.4f}, b={b_init:.4f}")
print(f"Peak 1: A={A1_init:.3f}, mu={mu1_init:.3f}, gamma={gamma1_init:.3f}")
print(f"Peak 2: A={A2_init:.3f}, mu={mu2_init:.3f}, gamma={gamma2_init:.3f}")

# --- Step 3: Define model ---
def model_2lorentzian_linear(x, m, b, A1, mu1, gamma1, A2, mu2, gamma2):
    baseline = m * x + b
    L1 = A1 / (1.0 + ((x - mu1) / gamma1)**2)
    L2 = A2 / (1.0 + ((x - mu2) / gamma2)**2)
    return baseline + L1 + L2

# --- Step 4: Pass 1 fit ---
p0 = [m_init, b_init, A1_init, mu1_init, gamma1_init, A2_init, mu2_init, gamma2_init]

lower_bounds = [-0.05, -0.5, 0.5, mu1_init - 8, 0.3, 0.3, mu2_init - 8, 0.3]
upper_bounds = [0.05, 0.5, 15.0, mu1_init + 8, 15.0, 12.0, mu2_init + 8, 15.0]

# Clip initial guesses to be within bounds
p0_clipped = []
for i, (p, lb, ub) in enumerate(zip(p0, lower_bounds, upper_bounds)):
    p0_clipped.append(np.clip(p, lb + 1e-10, ub - 1e-10))
p0 = p0_clipped

try:
    popt1, pcov1 = curve_fit(
        model_2lorentzian_linear, x, y,
        p0=p0,
        bounds=(lower_bounds, upper_bounds),
        method='trf',
        max_nfev=15000,
        ftol=1e-12, xtol=1e-12, gtol=1e-12
    )
    print("Pass 1 converged.")
except Exception as e:
    print(f"Pass 1 failed: {e}")
    popt1 = np.array(p0)
    pcov1 = np.eye(len(p0)) * 1e10

# --- Step 5: Pass 2 polish fit ---
m_p1, b_p1, A1_p1, mu1_p1, gamma1_p1, A2_p1, mu2_p1, gamma2_p1 = popt1

lower_bounds2 = [
    max(-0.05, m_p1 - 0.02),
    max(-0.5, b_p1 - 0.3),
    max(0.5, A1_p1 * 0.6),
    mu1_p1 - 3,
    max(0.3, gamma1_p1 * 0.6),
    max(0.3, A2_p1 * 0.6),
    mu2_p1 - 3,
    max(0.3, gamma2_p1 * 0.6)
]
upper_bounds2 = [
    min(0.05, m_p1 + 0.02),
    min(0.5, b_p1 + 0.3),
    min(15.0, A1_p1 * 1.4),
    mu1_p1 + 3,
    min(15.0, gamma1_p1 * 1.4),
    min(12.0, A2_p1 * 1.4),
    mu2_p1 + 3,
    min(15.0, gamma2_p1 * 1.4)
]

# Ensure lower < upper
for i in range(len(lower_bounds2)):
    if lower_bounds2[i] >= upper_bounds2[i]:
        lower_bounds2[i] = popt1[i] - abs(popt1[i]) * 0.5 - 1e-6
        upper_bounds2[i] = popt1[i] + abs(popt1[i]) * 0.5 + 1e-6

try:
    popt2, pcov2 = curve_fit(
        model_2lorentzian_linear, x, y,
        p0=popt1,
        bounds=(lower_bounds2, upper_bounds2),
        method='trf',
        max_nfev=15000,
        ftol=1e-12, xtol=1e-12, gtol=1e-12
    )
    print("Pass 2 converged.")
    popt = popt2
    pcov = pcov2
except Exception as e:
    print(f"Pass 2 failed: {e}, using Pass 1 results.")
    popt = popt1
    pcov = pcov1

# Extract parameters
m_fit, b_fit, A1_fit, mu1_fit, gamma1_fit, A2_fit, mu2_fit, gamma2_fit = popt

# Uncertainties
perr = np.sqrt(np.diag(pcov))
m_err, b_err, A1_err, mu1_err, gamma1_err, A2_err, mu2_err, gamma2_err = perr

# --- Step 6: Compute fit quality ---
y_fit = model_2lorentzian_linear(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))

# --- Step 7: FWHM ---
FWHM_peak1 = 2.0 * abs(gamma1_fit)
FWHM_peak2 = 2.0 * abs(gamma2_fit)
FWHM_peak1_err = 2.0 * gamma1_err
FWHM_peak2_err = 2.0 * gamma2_err

# --- Step 8: Diagnostic residual analysis ---
# Peak 1 core
mask_p1_core = np.abs(x - mu1_fit) < 3
res_p1_core_mean = np.mean(residuals[mask_p1_core]) if np.sum(mask_p1_core) > 0 else 0.0
res_p1_core_std = np.std(residuals[mask_p1_core]) if np.sum(mask_p1_core) > 0 else 0.0

# Peak 2 core
mask_p2_core = np.abs(x - mu2_fit) < 3
res_p2_core_mean = np.mean(residuals[mask_p2_core]) if np.sum(mask_p2_core) > 0 else 0.0
res_p2_core_std = np.std(residuals[mask_p2_core]) if np.sum(mask_p2_core) > 0 else 0.0

# Inter-peak region
mask_inter = (x > mu1_fit + 10) & (x < mu2_fit - 10)
res_inter_mean = np.mean(residuals[mask_inter]) if np.sum(mask_inter) > 0 else 0.0
res_inter_std = np.std(residuals[mask_inter]) if np.sum(mask_inter) > 0 else 0.0

# Baseline wings
mask_baseline = (x < 15) | (x > 85)
res_baseline_mean = np.mean(residuals[mask_baseline]) if np.sum(mask_baseline) > 0 else 0.0
res_baseline_std = np.std(residuals[mask_baseline]) if np.sum(mask_baseline) > 0 else 0.0

print(f"\n--- Fit Results ---")
print(f"m = {m_fit:.6f} +/- {m_err:.6f}")
print(f"b = {b_fit:.6f} +/- {b_err:.6f}")
print(f"Peak 1: A={A1_fit:.4f}+/-{A1_err:.4f}, mu={mu1_fit:.4f}+/-{mu1_err:.4f}, gamma={gamma1_fit:.4f}+/-{gamma1_err:.4f}")
print(f"Peak 2: A={A2_fit:.4f}+/-{A2_err:.4f}, mu={mu2_fit:.4f}+/-{mu2_err:.4f}, gamma={gamma2_fit:.4f}+/-{gamma2_err:.4f}")
print(f"FWHM Peak 1: {FWHM_peak1:.4f} +/- {FWHM_peak1_err:.4f}")
print(f"FWHM Peak 2: {FWHM_peak2:.4f} +/- {FWHM_peak2_err:.4f}")
print(f"R^2 = {R_squared:.8f}")
print(f"RMSE = {RMSE:.6f}")
print(f"\n--- Diagnostic Residuals ---")
print(f"Peak 1 core: mean={res_p1_core_mean:.6f}, std={res_p1_core_std:.6f}")
print(f"Peak 2 core: mean={res_p2_core_mean:.6f}, std={res_p2_core_std:.6f}")
print(f"Inter-peak: mean={res_inter_mean:.6f}, std={res_inter_std:.6f}")
print(f"Baseline wings: mean={res_baseline_mean:.6f}, std={res_baseline_std:.6f}")

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

# Main plot
ax1.scatter(x, y, s=8, alpha=0.5, color='gray', label='Data', zorder=1)
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit (2 Lorentzian + Linear)', zorder=3)

# Plot individual components
baseline_curve = m_fit * x + b_fit
L1_curve = A1_fit / (1.0 + ((x - mu1_fit) / gamma1_fit)**2)
L2_curve = A2_fit / (1.0 + ((x - mu2_fit) / gamma2_fit)**2)

ax1.plot(x, baseline_curve, 'k--', linewidth=1, alpha=0.7, label=f'Baseline (y={m_fit:.4f}x+{b_fit:.4f})', zorder=2)
ax1.plot(x, baseline_curve + L1_curve, 'b--', linewidth=1.5, alpha=0.7,
         label=f'Peak 1 (\u03bc={mu1_fit:.2f}, A={A1_fit:.2f}, FWHM={FWHM_peak1:.2f})', zorder=2)
ax1.plot(x, baseline_curve + L2_curve, 'g--', linewidth=1.5, alpha=0.7,
         label=f'Peak 2 (\u03bc={mu2_fit:.2f}, A={A2_fit:.2f}, FWHM={FWHM_peak2:.2f})', zorder=2)

ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title(f'2-Lorentzian + Linear Baseline Fit | R\u00b2={R_squared:.6f} | RMSE={RMSE:.4f}', fontsize=13)
ax1.legend(fontsize=9, loc='upper right')
ax1.grid(True, alpha=0.3)

# Residuals plot
ax2.scatter(x, residuals, s=5, alpha=0.5, color='gray', zorder=1)
ax2.axhline(y=0, color='r', linewidth=1, linestyle='--')

# Shade diagnostic regions
for mask_d, color_d, label_d in [
    (mask_p1_core, 'blue', 'P1 core'),
    (mask_p2_core, 'green', 'P2 core'),
    (mask_inter, 'orange', 'Inter-peak'),
    (mask_baseline, 'purple', 'Baseline')
]:
    if np.sum(mask_d) > 0:
        x_region = x[mask_d]
        ax2.axvspan(x_region.min(), x_region.max(), alpha=0.08, color=color_d)

ax2.set_xlabel('x (channel/wavelength)', fontsize=12)
ax2.set_ylabel('Residuals', fontsize=12)
ax2.set_title('Residuals', fontsize=11)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('spectrum_0002_fit.png', dpi=150, bbox_inches='tight')
plt.close()
print("\nVisualization saved to spectrum_0002_fit.png")

# --- Step 10: Check for Voigt upgrade need ---
voigt_note = ""
p1_core_frac = abs(res_p1_core_mean) / A1_fit if A1_fit > 0 else 0
p2_core_frac = abs(res_p2_core_mean) / A2_fit if A2_fit > 0 else 0
if p1_core_frac > 0.01 or p2_core_frac > 0.01:
    voigt_note = f" Peak-core residual fractions: P1={p1_core_frac:.4f}, P2={p2_core_frac:.4f}. A Voigt profile upgrade may improve the fit if Gaussian core contributions are significant."
else:
    voigt_note = " Peak-core residuals are <1% of peak amplitude, confirming Lorentzian lineshape is adequate."

# --- Step 11: Output JSON ---
summary_str = (f"Two-Lorentzian + linear baseline fit to spectroscopic data with {len(x)} points. "
               f"R²={R_squared:.6f}, RMSE={RMSE:.6f}. "
               f"Peak 1 at x={mu1_fit:.2f} with FWHM={FWHM_peak1:.2f}, "
               f"Peak 2 at x={mu2_fit:.2f} with FWHM={FWHM_peak2:.2f}. "
               f"Diagnostic residual analysis - Peak1 core mean: {res_p1_core_mean:.6f}, "
               f"Peak2 core mean: {res_p2_core_mean:.6f}, "
               f"Inter-peak mean: {res_inter_mean:.6f}, "
               f"Baseline wings mean: {res_baseline_mean:.6f}.{voigt_note}")

results = {
    "model_type": "2-Lorentzian peaks + linear baseline: y = m*x + b + A1/(1+((x-mu1)/gamma1)^2) + A2/(1+((x-mu2)/gamma2)^2)",
    "parameters": {
        "baseline": {
            "slope_m": round(float(m_fit), 8),
            "slope_m_err": round(float(m_err), 8),
            "intercept_b": round(float(b_fit), 8),
            "intercept_b_err": round(float(b_err), 8)
        },
        "peak_1": {
            "center": round(float(mu1_fit), 6),
            "center_err": round(float(mu1_err), 6),
            "amplitude": round(float(A1_fit), 6),
            "amplitude_err": round(float(A1_err), 6),
            "gamma_hwhm": round(float(gamma1_fit), 6),
            "gamma_hwhm_err": round(float(gamma1_err), 6),
            "fwhm": round(float(FWHM_peak1), 6),
            "fwhm_err": round(float(FWHM_peak1_err), 6)
        },
        "peak_2": {
            "center": round(float(mu2_fit), 6),
            "center_err": round(float(mu2_err), 6),
            "amplitude": round(float(A2_fit), 6),
            "amplitude_err": round(float(A2_err), 6),
            "gamma_hwhm": round(float(gamma2_fit), 6),
            "gamma_hwhm_err": round(float(gamma2_err), 6),
            "fwhm": round(float(FWHM_peak2), 6),
            "fwhm_err": round(float(FWHM_peak2_err), 6)
        },
        "diagnostics": {
            "residual_mean_peak1_core": round(float(res_p1_core_mean), 8),
            "residual_std_peak1_core": round(float(res_p1_core_std), 8),
            "residual_mean_peak2_core": round(float(res_p2_core_mean), 8),
            "residual_std_peak2_core": round(float(res_p2_core_std), 8),
            "residual_mean_interpeak": round(float(res_inter_mean), 8),
            "residual_std_interpeak": round(float(res_inter_std), 8),
            "residual_mean_baseline": round(float(res_baseline_mean), 8),
            "residual_std_baseline": round(float(res_baseline_std), 8)
        }
    },
    "fit_quality": {
        "r_squared": round(float(R_squared), 8),
        "rmse": round(float(RMSE), 8)
    },
    "summary": summary_str
}

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