import numpy as np
import json
import warnings
from scipy.optimize import curve_fit, differential_evolution
from scipy.signal import find_peaks
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')

# ─────────────────────────────────────────────
# 1. Load data
# ─────────────────────────────────────────────
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/dbg/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. Pre-process: find local maxima
# ─────────────────────────────────────────────
peaks_idx, peak_props = find_peaks(y, height=0.3 * np.max(y), distance=10)
print("Detected peak indices:", peaks_idx)
print("Detected peak positions:", x[peaks_idx] if len(peaks_idx) > 0 else "none")
print("Detected peak amplitudes:", y[peaks_idx] if len(peaks_idx) > 0 else "none")

# ─────────────────────────────────────────────
# 3. Model definition: linear baseline + 4 Gaussians
# ─────────────────────────────────────────────
def model_4g(x, a0, a1, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3, A4, mu4, sigma4):
    baseline = a0 + a1 * x
    g1 = A1 * np.exp(-0.5 * ((x - mu1) / sigma1) ** 2)
    g2 = A2 * np.exp(-0.5 * ((x - mu2) / sigma2) ** 2)
    g3 = A3 * np.exp(-0.5 * ((x - mu3) / sigma3) ** 2)
    g4 = A4 * np.exp(-0.5 * ((x - mu4) / sigma4) ** 2)
    return baseline + g1 + g2 + g3 + g4

# ─────────────────────────────────────────────
# 4. Bounds and initial guesses
# ─────────────────────────────────────────────
p0 = [0.05, 0.001,
      3.0, 20.0, 2.0,
      2.6, 42.0, 2.0,
      4.2, 65.0, 2.0,
      2.3, 85.0, 2.0]

lower = [-1.0, -0.1,
         0.01, 10.0, 0.5,
         0.01, 35.0, 0.5,
         0.01, 58.0, 0.5,
         0.01, 78.0, 0.5]

upper = [1.0, 0.1,
         20.0, 30.0, 15.0,
         20.0, 52.0, 15.0,
         20.0, 75.0, 15.0,
         20.0, 95.0, 15.0]

bounds = (lower, upper)
de_bounds = list(zip(lower, upper))

# ─────────────────────────────────────────────
# 5. Global optimisation via differential_evolution
# ─────────────────────────────────────────────
def residual_sum_sq(params):
    try:
        fitted = model_4g(x, *params)
        return np.sum((y - fitted) ** 2)
    except Exception:
        return 1e18

print("Running differential_evolution global search ...")
de_result = differential_evolution(
    residual_sum_sq,
    de_bounds,
    seed=42,
    maxiter=2000,
    tol=1e-8,
    popsize=20,
    mutation=(0.5, 1.5),
    recombination=0.7,
    workers=1,
    polish=True
)
de_p0 = de_result.x
print("DE best RSS:", de_result.fun)
print("DE solution:", de_p0)

# ─────────────────────────────────────────────
# 6. Local refinement with curve_fit (trf) seeded from DE
# ─────────────────────────────────────────────
def try_curve_fit(seed):
    popt, pcov = curve_fit(
        model_4g, x, y,
        p0=seed,
        bounds=bounds,
        method='trf',
        maxfev=50000,
        ftol=1e-12,
        xtol=1e-12,
        gtol=1e-12
    )
    return popt, pcov

# Attempt 1: DE seed
try:
    popt_de, pcov_de = try_curve_fit(de_p0)
    y_pred_de = model_4g(x, *popt_de)
    ss_res_de = np.sum((y - y_pred_de) ** 2)
except Exception as e:
    print("DE-seeded curve_fit failed:", e)
    popt_de, pcov_de = None, None
    ss_res_de = 1e18

# Attempt 2: manual p0
try:
    popt_man, pcov_man = try_curve_fit(p0)
    y_pred_man = model_4g(x, *popt_man)
    ss_res_man = np.sum((y - y_pred_man) ** 2)
except Exception as e:
    print("Manual-p0 curve_fit failed:", e)
    popt_man, pcov_man = None, None
    ss_res_man = 1e18

# Pick best
if ss_res_de <= ss_res_man and popt_de is not None:
    popt, pcov = popt_de, pcov_de
    print("Selected DE-seeded solution.")
elif popt_man is not None:
    popt, pcov = popt_man, pcov_man
    print("Selected manual-p0 solution.")
else:
    # Fall back to DE solution directly
    popt = de_p0
    pcov = np.full((len(de_p0), len(de_p0)), np.inf)
    print("Fell back to raw DE solution (no curve_fit converged).")

# ─────────────────────────────────────────────
# 7. Uncertainties
# ─────────────────────────────────────────────
perr = np.sqrt(np.diag(pcov)) if not np.any(np.isinf(pcov)) else np.full(len(popt), np.nan)

# ─────────────────────────────────────────────
# 8. Metrics
# ─────────────────────────────────────────────
y_pred = model_4g(x, *popt)
ss_res = np.sum((y - y_pred) ** 2)
ss_tot = np.sum((y - np.mean(y)) ** 2)
r2 = 1.0 - ss_res / ss_tot
rmse = np.sqrt(np.mean((y - y_pred) ** 2))
residuals = y - y_pred

print(f"R² = {r2:.6f}")
print(f"RMSE = {rmse:.6f}")
if r2 < 0.90:
    print("WARNING: R² < 0.90 — fit may require review (5th component or non-linear baseline).")

# ─────────────────────────────────────────────
# 9. Parameter names and bound-hit check
# ─────────────────────────────────────────────
param_names = ['a0', 'a1',
               'A1', 'mu1', 'sigma1',
               'A2', 'mu2', 'sigma2',
               'A3', 'mu3', 'sigma3',
               'A4', 'mu4', 'sigma4']
bound_warnings = []
for i, (name, val) in enumerate(zip(param_names, popt)):
    tol = 1e-4
    if abs(val - lower[i]) < tol or abs(val - upper[i]) < tol:
        bound_warnings.append(f"{name}={val:.4f} is AT BOUND [{lower[i]}, {upper[i]}] — potentially unphysical.")

if bound_warnings:
    print("BOUND WARNINGS:")
    for w in bound_warnings:
        print(" ", w)

# ─────────────────────────────────────────────
# 10. Peak areas: A * sigma * sqrt(2*pi)
# ─────────────────────────────────────────────
sqrt2pi = np.sqrt(2 * np.pi)
area1 = popt[2] * popt[4] * sqrt2pi
area2 = popt[5] * popt[7] * sqrt2pi
area3 = popt[8] * popt[10] * sqrt2pi
area4 = popt[11] * popt[13] * sqrt2pi

# ─────────────────────────────────────────────
# 11. Residual diagnostics
# ─────────────────────────────────────────────
max_signal = np.max(np.abs(y))
resid_max_frac = np.max(np.abs(residuals)) / max_signal
print(f"Max residual as fraction of max signal: {resid_max_frac:.4f}")
if resid_max_frac > 0.10:
    print("WARNING: Residuals exceed 10% of max signal — systematic structure may remain.")

# ─────────────────────────────────────────────
# 12. Visualisation
# ─────────────────────────────────────────────
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), gridspec_kw={'height_ratios': [3, 1]})

# Component curves
a0_, a1_ = popt[0], popt[1]
baseline_curve = a0_ + a1_ * x
A1_, mu1_, s1_ = popt[2], popt[3], popt[4]
A2_, mu2_, s2_ = popt[5], popt[6], popt[7]
A3_, mu3_, s3_ = popt[8], popt[9], popt[10]
A4_, mu4_, s4_ = popt[11], popt[12], popt[13]
g1_curve = A1_ * np.exp(-0.5 * ((x - mu1_) / s1_) ** 2)
g2_curve = A2_ * np.exp(-0.5 * ((x - mu2_) / s2_) ** 2)
g3_curve = A3_ * np.exp(-0.5 * ((x - mu3_) / s3_) ** 2)
g4_curve = A4_ * np.exp(-0.5 * ((x - mu4_) / s4_) ** 2)

ax1.plot(x, y, 'k.', markersize=2, label='Data', alpha=0.6)
ax1.plot(x, y_pred, 'r-', linewidth=2, label=f'4-Gaussian fit (R²={r2:.4f})')
ax1.plot(x, baseline_curve + g1_curve, '--', color='steelblue', linewidth=1.2, label=f'Peak 1 (μ={mu1_:.1f})')
ax1.plot(x, baseline_curve + g2_curve, '--', color='darkorange', linewidth=1.2, label=f'Peak 2 (μ={mu2_:.1f})')
ax1.plot(x, baseline_curve + g3_curve, '--', color='forestgreen', linewidth=1.2, label=f'Peak 3 (μ={mu3_:.1f})')
ax1.plot(x, baseline_curve + g4_curve, '--', color='purple', linewidth=1.2, label=f'Peak 4 (μ={mu4_:.1f})')
ax1.plot(x, baseline_curve, ':', color='gray', linewidth=1, label='Linear baseline')
if len(peaks_idx) > 0:
    ax1.plot(x[peaks_idx], y[peaks_idx], 'v', color='gold', markersize=8, label='Detected maxima')
ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title('Spectroscopic Curve Fitting: 4-Gaussian + Linear Baseline', fontsize=13)
ax1.legend(fontsize=8, loc='upper left')
ax1.grid(True, alpha=0.3)

ax2.plot(x, residuals, 'b-', linewidth=0.8, alpha=0.7, label='Residuals')
ax2.axhline(0, color='k', linewidth=1)
ax2.fill_between(x, residuals, 0, alpha=0.2, color='blue')
ax2.set_xlabel('x', fontsize=12)
ax2.set_ylabel('Residual', fontsize=12)
ax2.set_title(f'Residuals (RMSE={rmse:.4f}, max={resid_max_frac*100:.1f}% of max signal)', fontsize=11)
ax2.grid(True, alpha=0.3)
ax2.legend(fontsize=9)

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

# ─────────────────────────────────────────────
# 13. Results JSON
# ─────────────────────────────────────────────
def _f(v): return float(round(v, 6)) if np.isfinite(v) else None

pidx = {n: i for i, n in enumerate(param_names)}

results = {
    "model_type": "Linear baseline + 4-Gaussian components (14 parameters total)",
    "parameters": {
        "baseline": {
            "intercept_a0": _f(popt[pidx['a0']]),
            "intercept_a0_err": _f(perr[pidx['a0']]),
            "slope_a1": _f(popt[pidx['a1']]),
            "slope_a1_err": _f(perr[pidx['a1']])
        },
        "peak_1": {
            "amplitude": _f(popt[pidx['A1']]),
            "amplitude_err": _f(perr[pidx['A1']]),
            "center": _f(popt[pidx['mu1']]),
            "center_err": _f(perr[pidx['mu1']]),
            "sigma": _f(popt[pidx['sigma1']]),
            "sigma_err": _f(perr[pidx['sigma1']]),
            "area": _f(area1)
        },
        "peak_2": {
            "amplitude": _f(popt[pidx['A2']]),
            "amplitude_err": _f(perr[pidx['A2']]),
            "center": _f(popt[pidx['mu2']]),
            "center_err": _f(perr[pidx['mu2']]),
            "sigma": _f(popt[pidx['sigma2']]),
            "sigma_err": _f(perr[pidx['sigma2']]),
            "area": _f(area2)
        },
        "peak_3": {
            "amplitude": _f(popt[pidx['A3']]),
            "amplitude_err": _f(perr[pidx['A3']]),
            "center": _f(popt[pidx['mu3']]),
            "center_err": _f(perr[pidx['mu3']]),
            "sigma": _f(popt[pidx['sigma3']]),
            "sigma_err": _f(perr[pidx['sigma3']]),
            "area": _f(area3)
        },
        "peak_4": {
            "amplitude": _f(popt[pidx['A4']]),
            "amplitude_err": _f(perr[pidx['A4']]),
            "center": _f(popt[pidx['mu4']]),
            "center_err": _f(perr[pidx['mu4']]),
            "sigma": _f(popt[pidx['sigma4']]),
            "sigma_err": _f(perr[pidx['sigma4']]),
            "area": _f(area4)
        }
    },
    "fit_quality": {
        "r_squared": _f(r2),
        "rmse": _f(rmse),
        "residual_max_fraction_of_signal": _f(resid_max_frac)
    },
    "optimization": {
        "global_search": "differential_evolution (popsize=20, maxiter=2000)",
        "local_refinement": "scipy.optimize.curve_fit (method=trf, maxfev=50000)",
        "de_rss": _f(de_result.fun)
    },
    "bound_warnings": bound_warnings,
    "summary": (
        f"4-Gaussian + linear baseline fit achieved R²={r2:.4f}, RMSE={rmse:.4f}. "
        f"Peaks found at x≈{popt[3]:.1f} (A={popt[2]:.3f}), x≈{popt[6]:.1f} (A={popt[5]:.3f}), "
        f"x≈{popt[9]:.1f} (A={popt[8]:.3f}), x≈{popt[12]:.1f} (A={popt[11]:.3f}). "
        f"Max residual is {resid_max_frac*100:.1f}% of max signal. "
        + ("Fit ACCEPTED (R²>=0.90)." if r2 >= 0.90 else "WARNING: R²<0.90 — consider 5th component or non-linear baseline.")
        + (f" Bound hits: {len(bound_warnings)}." if bound_warnings else " No parameters hit bounds.")
    )
}

print("FIT_RESULTS_JSON:" + json.dumps(results))
