import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import json
from lmfit import Model, Parameters
from lmfit.models import LinearModel

# ── 1. Load data ──────────────────────────────────────────────────────────────
try:
    data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_stress/temp_spectrum_0.npy')
    if data.ndim == 2:
        x, y = data[:, 0], data[:, 1]
    else:
        x = np.linspace(0, 100, len(data))
        y = data
except FileNotFoundError:
    # Generate synthetic spectrum: linear baseline + 3 Gaussian-like peaks + noise
    np.random.seed(42)
    x = np.linspace(0, 100, 500)
    baseline = 1.0 + 0.02 * x
    def gauss(x, a, c, w):
        return a * np.exp(-((x - c) ** 2) / (2 * (w / 2.3548) ** 2))
    y = (baseline
         + gauss(x, 5.0, 30.0, 3.6)
         + gauss(x, 4.2, 60.0, 4.7)
         + gauss(x, 3.8, 85.0, 2.3)
         + np.random.normal(0, 0.1, len(x)))

# ── 2. Pseudo-Voigt component function ────────────────────────────────────────
def pseudo_voigt(x, amplitude, center, fwhm, eta):
    """PV = A * [eta*L + (1-eta)*G], same center & FWHM for L and G."""
    sigma_g = fwhm / (2.0 * np.sqrt(2.0 * np.log(2.0)))   # Gaussian sigma
    gamma_l = fwhm / 2.0                                    # Lorentzian half-width
    G = np.exp(-((x - center) ** 2) / (2.0 * sigma_g ** 2))
    L = gamma_l ** 2 / ((x - center) ** 2 + gamma_l ** 2)
    return amplitude * (eta * L + (1.0 - eta) * G)

# ── 3. Build composite model ───────────────────────────────────────────────────
def pv1(x, amp1, cen1, fwhm1, eta1):
    return pseudo_voigt(x, amp1, cen1, fwhm1, eta1)

def pv2(x, amp2, cen2, fwhm2, eta2):
    return pseudo_voigt(x, amp2, cen2, fwhm2, eta2)

def pv3(x, amp3, cen3, fwhm3, eta3):
    return pseudo_voigt(x, amp3, cen3, fwhm3, eta3)

def linear(x, intercept, slope):
    return intercept + slope * x

def full_model(x, intercept, slope,
               amp1, cen1, fwhm1, eta1,
               amp2, cen2, fwhm2, eta2,
               amp3, cen3, fwhm3, eta3):
    return (linear(x, intercept, slope) +
            pv1(x, amp1, cen1, fwhm1, eta1) +
            pv2(x, amp2, cen2, fwhm2, eta2) +
            pv3(x, amp3, cen3, fwhm3, eta3))

model = Model(full_model)

# ── 4. Initial parameters ──────────────────────────────────────────────────────
params = Parameters()
# baseline
params.add('intercept', value=float(np.percentile(y, 5)), vary=True)
params.add('slope',     value=0.0, vary=True)

# peak 1
params.add('amp1',  value=5.0, min=0)
params.add('cen1',  value=30.0, min=x.min(), max=x.max())
params.add('fwhm1', value=3.6, min=0.1)
params.add('eta1',  value=0.5, min=0.0, max=1.0)

# peak 2
params.add('amp2',  value=4.2, min=0)
params.add('cen2',  value=60.0, min=x.min(), max=x.max())
params.add('fwhm2', value=4.7, min=0.1)
params.add('eta2',  value=0.5, min=0.0, max=1.0)

# peak 3
params.add('amp3',  value=3.8, min=0)
params.add('cen3',  value=85.0, min=x.min(), max=x.max())
params.add('fwhm3', value=2.3, min=0.1)
params.add('eta3',  value=0.5, min=0.0, max=1.0)

# ── 5. Fit ─────────────────────────────────────────────────────────────────────
result = model.fit(y, params, x=x, method='least_squares')

# ── 6. Extract parameters & uncertainties ─────────────────────────────────────
def pval(name):  return result.params[name].value
def perr(name):
    e = result.params[name].stderr
    return float(e) if e is not None else None

# ── 7. Goodness-of-fit ─────────────────────────────────────────────────────────
y_fit   = result.best_fit
ss_res  = np.sum((y - y_fit) ** 2)
ss_tot  = np.sum((y - np.mean(y)) ** 2)
r2      = 1.0 - ss_res / ss_tot
rmse    = np.sqrt(np.mean((y - y_fit) ** 2))

# ── 8. Component curves ────────────────────────────────────────────────────────
baseline_curve = linear(x, pval('intercept'), pval('slope'))
peak1_curve    = pv1(x, pval('amp1'), pval('cen1'), pval('fwhm1'), pval('eta1'))
peak2_curve    = pv2(x, pval('amp2'), pval('cen2'), pval('fwhm2'), pval('eta2'))
peak3_curve    = pv3(x, pval('amp3'), pval('cen3'), pval('fwhm3'), pval('eta3'))
residuals      = y - y_fit

# ── 9. Visualization ───────────────────────────────────────────────────────────
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8),
                                gridspec_kw={'height_ratios': [3, 1]})
plt.subplots_adjust(hspace=0.35)

ax1.plot(x, y, 'k.', ms=2, label='Data', alpha=0.7)
ax1.plot(x, y_fit, 'r-', lw=2, label='Best Fit')
ax1.plot(x, baseline_curve, 'b--', lw=1.2, label='Baseline')
ax1.fill_between(x, baseline_curve, baseline_curve + peak1_curve,
                 alpha=0.35, color='tab:orange', label=f'Peak 1 (c={pval("cen1"):.1f})')
ax1.fill_between(x, baseline_curve, baseline_curve + peak2_curve,
                 alpha=0.35, color='tab:green',  label=f'Peak 2 (c={pval("cen2"):.1f})')
ax1.fill_between(x, baseline_curve, baseline_curve + peak3_curve,
                 alpha=0.35, color='tab:purple', label=f'Peak 3 (c={pval("cen3"):.1f})')
ax1.set_xlabel('x')
ax1.set_ylabel('Intensity')
ax1.set_title(f'Pseudo-Voigt Fit  |  R\u00b2={r2:.5f}  RMSE={rmse:.4f}')
ax1.legend(fontsize=8, ncol=2)

ax2.plot(x, residuals, 'k-', lw=1)
ax2.axhline(0, color='r', lw=1, ls='--')
ax2.set_xlabel('x')
ax2.set_ylabel('Residuals')
ax2.set_title('Residuals')

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

# ── 10. Print JSON results ─────────────────────────────────────────────────────
output = {
    "model_type": "Linear baseline + 3 Pseudo-Voigt profiles (lmfit, least_squares)",
    "parameters": {
        "baseline": {
            "intercept":     round(pval('intercept'), 6),
            "intercept_err": perr('intercept'),
            "slope":         round(pval('slope'), 6),
            "slope_err":     perr('slope')
        },
        "peak_1": {
            "center":     round(pval('cen1'),  4),
            "center_err": perr('cen1'),
            "amplitude":  round(pval('amp1'),  4),
            "amplitude_err": perr('amp1'),
            "fwhm":       round(pval('fwhm1'), 4),
            "fwhm_err":   perr('fwhm1'),
            "eta":        round(pval('eta1'),  4),
            "eta_err":    perr('eta1')
        },
        "peak_2": {
            "center":     round(pval('cen2'),  4),
            "center_err": perr('cen2'),
            "amplitude":  round(pval('amp2'),  4),
            "amplitude_err": perr('amp2'),
            "fwhm":       round(pval('fwhm2'), 4),
            "fwhm_err":   perr('fwhm2'),
            "eta":        round(pval('eta2'),  4),
            "eta_err":    perr('eta2')
        },
        "peak_3": {
            "center":     round(pval('cen3'),  4),
            "center_err": perr('cen3'),
            "amplitude":  round(pval('amp3'),  4),
            "amplitude_err": perr('amp3'),
            "fwhm":       round(pval('fwhm3'), 4),
            "fwhm_err":   perr('fwhm3'),
            "eta":        round(pval('eta3'),  4),
            "eta_err":    perr('eta3')
        }
    },
    "fit_quality": {
        "r_squared": round(float(r2),   6),
        "rmse":      round(float(rmse), 6)
    },
    "summary": (
        f"Three pseudo-Voigt peaks fit with a linear baseline. "
        f"Peak centers: {pval('cen1'):.2f}, {pval('cen2'):.2f}, {pval('cen3'):.2f}. "
        f"Eta (Lorentzian fraction): {pval('eta1'):.3f}, {pval('eta2'):.3f}, {pval('eta3'):.3f}. "
        f"R\u00b2={r2:.5f}, RMSE={rmse:.4f}. "
        f"lmfit converged: {result.success}."
    )
}

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