import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.special import wofz
from scipy.optimize import least_squares
import json
import warnings
warnings.filterwarnings('ignore')

# ── 1. Load data ──────────────────────────────────────────────────────────────
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_series_annealing_data/output_series/temp_spectrum_0.npy')
if data.ndim == 2:
    x, y = data[:, 0], data[:, 1]
else:
    x = np.linspace(0, 100, len(data))
    y = data

# ── 2. Voigt profile ──────────────────────────────────────────────────────────
def voigt(x, mu, sigma, gamma, A):
    """Voigt profile via Faddeeva function.
    V(x) = Re[w(z)] / (sigma*sqrt(2*pi)), z = (x-mu+i*gamma)/(sigma*sqrt(2))
    Returns A * normalised_voigt(x)
    """
    sigma = max(sigma, 1e-10)
    gamma = max(gamma, 1e-10)
    z = ((x - mu) + 1j * gamma) / (sigma * np.sqrt(2))
    w = wofz(z)
    profile = np.real(w) / (sigma * np.sqrt(2 * np.pi))
    # normalise so peak value ≈ 1 when A=peak_height
    peak_z = (1j * gamma) / (sigma * np.sqrt(2))
    peak_val = np.real(wofz(peak_z)) / (sigma * np.sqrt(2 * np.pi))
    peak_val = max(peak_val, 1e-30)
    return A * profile / peak_val

# ── 3. Full model ─────────────────────────────────────────────────────────────
def model(x, params):
    """baseline + 3 Voigt peaks
    params = [a, tau, c,
              mu1, A1, sig1, gam1,
              mu2, A2, sig2, gam2,
              mu3, A3, sig3, gam3]
    baseline = a*(1 - exp(-x/tau)) + c
    """
    a, tau, c = params[0], params[1], params[2]
    tau = max(tau, 1e-3)
    baseline = a * (1 - np.exp(-x / tau)) + c
    v1 = voigt(x, params[3], params[5], params[6], params[4])
    v2 = voigt(x, params[7], params[9], params[10], params[8])
    v3 = voigt(x, params[11], params[13], params[14], params[12])
    return baseline + v1 + v2 + v3

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

# ── 4. Baseline estimation ────────────────────────────────────────────────────
bg_mask = ((x >= 0) & (x <= 20)) | ((x >= 40) & (x <= 48)) | \
           ((x >= 63) & (x <= 70)) | ((x >= 88) & (x <= 100))
xb, yb = x[bg_mask], y[bg_mask]

if len(xb) > 3:
    def exp_res(p, xb, yb):
        a, tau, c = p
        tau = max(tau, 1e-3)
        return a * (1 - np.exp(-xb / tau)) + c - yb
    try:
        res_bg = least_squares(exp_res, [yb.max(), 30, yb.min()],
                               args=(xb, yb), loss='soft_l1')
        a0, tau0, c0 = res_bg.x
    except Exception:
        a0, tau0, c0 = yb.max(), 30.0, yb.min()
else:
    a0, tau0, c0 = y.max() * 0.1, 30.0, y.min()

# ── 5. Initial peak estimates ─────────────────────────────────────────────────
# subtract baseline guess
bg_est = a0 * (1 - np.exp(-x / max(tau0, 1e-3))) + c0
y_sub = y - bg_est

mu1_0, mu2_0, mu3_0 = 30.0, 55.0, 78.0

def amp_near(x, y, center, half_win=5):
    mask = (x >= center - half_win) & (x <= center + half_win)
    return max(y[mask].max(), 0.1) if mask.any() else 0.5

A1_0 = amp_near(x, y_sub, mu1_0)
A2_0 = amp_near(x, y_sub, mu2_0)
A3_0 = amp_near(x, y_sub, mu3_0)

# start Lorentzian-heavy (gamma >> sigma)
p0 = [a0, tau0, c0,
      mu1_0, A1_0, 0.5, 3.0,
      mu2_0, A2_0, 0.5, 3.0,
      mu3_0, A3_0, 0.5, 3.0]

# ── 6. Bounds ────────────────────────────────────────────────────────────────
y_range = y.max() - y.min()
lb = [-y_range, 0.1, y.min() - y_range,
      10, 0, 1e-4, 1e-4,
      35, 0, 1e-4, 1e-4,
      60, 0, 1e-4, 1e-4]
ub = [y_range * 2, 200, y.max(),
      45, y_range * 2, 20, 20,
      70, y_range * 2, 20, 20,
      95, y_range * 2, 20, 20]

# ── 7. Fit ───────────────────────────────────────────────────────────────────
try:
    result = least_squares(residuals, p0, args=(x, y),
                           bounds=(lb, ub),
                           loss='soft_l1', f_scale=0.1,
                           max_nfev=50000, ftol=1e-12, xtol=1e-12, gtol=1e-12)
    popt = result.x
except Exception as e:
    print(f'First fit failed: {e}, retrying with linear loss')
    result = least_squares(residuals, p0, args=(x, y),
                           bounds=(lb, ub), max_nfev=50000)
    popt = result.x

# ── 8. Parameter uncertainties ───────────────────────────────────────────────
try:
    J = result.jac
    cov = np.linalg.pinv(J.T @ J)
    res_var = np.sum(result.fun ** 2) / max(len(y) - len(popt), 1)
    perr = np.sqrt(np.abs(np.diag(cov)) * res_var)
except Exception:
    perr = np.full(len(popt), np.nan)

# ── 9. Derived quantities ─────────────────────────────────────────────────────
def voigt_fwhm(sigma, gamma):
    """Approximate FWHM of Voigt (Thompson et al. 1987)"""
    fG = 2 * sigma * np.sqrt(2 * np.log(2))
    fL = 2 * gamma
    return (fG**5 + 2.69269*fG**4*fL + 2.42843*fG**3*fL**2 +
            4.47163*fG**2*fL**3 + 0.07842*fG*fL**4 + fL**5) ** 0.2

def voigt_eta(sigma, gamma):
    """pseudo-Voigt mixing parameter η"""
    fG = 2 * sigma * np.sqrt(2 * np.log(2))
    fL = 2 * gamma
    f = voigt_fwhm(sigma, gamma)
    return 1.36603*(fL/f) - 0.47719*(fL/f)**2 + 0.11116*(fL/f)**3

# ── 10. Goodness of fit ───────────────────────────────────────────────────────
y_fit = model(x, popt)
ss_res = np.sum((y - y_fit) ** 2)
ss_tot = np.sum((y - y.mean()) ** 2)
r2 = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan
rmse = np.sqrt(ss_res / len(y))

# ── 11. Plot ──────────────────────────────────────────────────────────────────
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7),
                                gridspec_kw={'height_ratios': [3, 1]})

ax1.plot(x, y, 'k.', ms=2, label='Data', alpha=0.6)
ax1.plot(x, y_fit, 'r-', lw=2, label=f'Fit (R²={r2:.4f})')

# baseline
a_p, tau_p, c_p = popt[0], popt[1], popt[2]
bg = a_p * (1 - np.exp(-x / max(tau_p, 1e-3))) + c_p
ax1.plot(x, bg, 'g--', lw=1.5, label='Baseline')

# individual peaks
colors = ['blue', 'orange', 'purple']
for i, col in enumerate(colors):
    pi = 3 + i * 4
    vi = voigt(x, popt[pi], popt[pi+2], popt[pi+3], popt[pi+1])
    ax1.fill_between(x, bg, bg + vi, alpha=0.25, color=col,
                     label=f'Peak {i+1} (μ={popt[pi]:.2f})')

ax1.set_ylabel('Intensity')
ax1.set_title('Voigt Peak Fitting with Exponential Baseline')
ax1.legend(fontsize=8)
ax1.grid(True, alpha=0.3)

ax2.plot(x, y - y_fit, 'k-', lw=0.8)
ax2.axhline(0, color='r', lw=1)
ax2.set_xlabel('x')
ax2.set_ylabel('Residuals')
ax2.grid(True, alpha=0.3)

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

# ── 12. Results ───────────────────────────────────────────────────────────────
def safe(v):
    return float(v) if np.isfinite(v) else None

peak_results = {}
for i in range(3):
    pi = 3 + i * 4
    sig = popt[pi+2]
    gam = popt[pi+3]
    fwhm = voigt_fwhm(sig, gam)
    eta = voigt_eta(sig, gam)
    peak_results[f'peak_{i+1}'] = {
        'center_mu':             safe(popt[pi]),
        'center_mu_err':         safe(perr[pi]),
        'amplitude_A':           safe(popt[pi+1]),
        'amplitude_A_err':       safe(perr[pi+1]),
        'gaussian_width_sigma':  safe(sig),
        'sigma_err':             safe(perr[pi+2]),
        'lorentzian_width_gamma':safe(gam),
        'gamma_err':             safe(perr[pi+3]),
        'fwhm':                  safe(fwhm),
        'voigt_eta':             safe(eta),
    }

results = {
    'model_type': '3 Voigt peaks + exponential baseline [a*(1-exp(-x/tau))+c]',
    'parameters': {
        **peak_results,
        'baseline': {
            'a':     safe(popt[0]),
            'a_err': safe(perr[0]),
            'tau':   safe(popt[1]),
            'tau_err': safe(perr[1]),
            'c':     safe(popt[2]),
            'c_err': safe(perr[2]),
        }
    },
    'fit_quality': {
        'r_squared': safe(r2),
        'rmse':      safe(rmse),
    },
    'summary': (
        f'Fitted 3 Voigt profiles on an exponential baseline. '
        f'Peak centres: μ1={popt[3]:.2f}, μ2={popt[7]:.2f}, μ3={popt[11]:.2f}. '
        f'R²={r2:.4f}, RMSE={rmse:.4f}. '
        f'Lorentzian (gamma) contributions: {popt[6]:.3f}, {popt[10]:.3f}, {popt[14]:.3f}. '
        f'Gaussian (sigma) contributions: {popt[5]:.3f}, {popt[9]:.3f}, {popt[13]:.3f}.'
    )
}
print(f'FIT_RESULTS_JSON:{json.dumps(results)}')
