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

# ── 1. Load data ──────────────────────────────────────────────────────────────
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/dbg2/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. Model definition ───────────────────────────────────────────────────────
def four_gaussian_linear(x, a0, a1,
                         A1, mu1, sig1,
                         A2, mu2, sig2,
                         A3, mu3, sig3,
                         A4, mu4, sig4):
    baseline = a0 + a1 * x
    g1 = A1 * np.exp(-0.5 * ((x - mu1) / sig1) ** 2)
    g2 = A2 * np.exp(-0.5 * ((x - mu2) / sig2) ** 2)
    g3 = A3 * np.exp(-0.5 * ((x - mu3) / sig3) ** 2)
    g4 = A4 * np.exp(-0.5 * ((x - mu4) / sig4) ** 2)
    return baseline + g1 + g2 + g3 + g4

# ── 3. Bounds & initial guesses ───────────────────────────────────────────────
p0 = [0.1, 0.002,
      3.05, 20.0, 3.0,
      2.60, 42.0, 2.5,
      4.20, 65.0, 3.0,
      2.30, 85.0, 3.5]

lower = [-1.0, -0.05,
         0.1, 10.0, 0.5,
         0.1, 32.0, 0.5,
         0.1, 55.0, 0.5,
         0.1, 75.0, 0.5]

upper = [2.0,  0.05,
         10.0, 30.0, 20.0,
         10.0, 52.0, 20.0,
         10.0, 75.0, 20.0,
         10.0, 95.0, 20.0]

bounds = (lower, upper)

# ── Helpers ───────────────────────────────────────────────────────────────────
def r_squared(y_obs, y_pred):
    ss_res = np.sum((y_obs - y_pred) ** 2)
    ss_tot = np.sum((y_obs - np.mean(y_obs)) ** 2)
    return 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0

def rmse(y_obs, y_pred):
    return float(np.sqrt(np.mean((y_obs - y_pred) ** 2)))

# ── 4. Stage 1 – Global TRF fit ───────────────────────────────────────────────
stage1_ok = False
popt = p0.copy()
pcov = np.full((14, 14), np.inf)

try:
    popt, pcov = curve_fit(
        four_gaussian_linear, x, y,
        p0=p0,
        bounds=bounds,
        method='trf',
        maxfev=50000,
        ftol=1e-10, xtol=1e-10, gtol=1e-10
    )
    y_fit = four_gaussian_linear(x, *popt)
    r2 = r_squared(y, y_fit)
    stage1_ok = (r2 >= 0.90)
except Exception as e:
    warnings.warn(f'Stage 1 failed: {e}')
    r2 = -np.inf

# ── 5. Stage 2 – Sequential warm-start fallback ───────────────────────────────
if not stage1_ok:
    warnings.warn('Stage 1 R² < 0.90 or failed – executing Stage 2 warm-start.')

    # 2a: Estimate baseline from spectral edges (bottom 5% of x range)
    edge_frac = int(0.05 * len(x))
    edge_frac = max(edge_frac, 5)
    edge_mask = np.concatenate([np.arange(edge_frac), np.arange(len(x) - edge_frac, len(x))])
    bl_coeffs = np.polyfit(x[edge_mask], y[edge_mask], 1)
    y_bl = np.polyval(bl_coeffs, x)
    y_sub = y - y_bl

    # 2b: Fit each peak individually in a narrow window
    peak_guesses = [
        (3.05, 20.0, 3.0),
        (2.60, 42.0, 2.5),
        (4.20, 65.0, 3.0),
        (2.30, 85.0, 3.5)
    ]
    warm = []
    for (Ag, mug, sigg) in peak_guesses:
        win_lo = mug - 4 * sigg
        win_hi = mug + 4 * sigg
        mask = (x >= win_lo) & (x <= win_hi)
        if mask.sum() < 5:
            warm.append((Ag, mug, sigg))
            continue
        try:
            def single_peak(xw, Aw, muw, sigw, c0, c1):
                return Aw * np.exp(-0.5 * ((xw - muw) / sigw) ** 2) + c0 + c1 * xw
            po, _ = curve_fit(
                single_peak, x[mask], y_sub[mask],
                p0=[Ag, mug, sigg, 0.0, 0.0],
                bounds=([0.1, win_lo, 0.5, -2, -0.1],
                        [10.0, win_hi, 20.0, 2, 0.1]),
                method='lm' if False else 'trf',
                maxfev=20000
            )
            warm.append((abs(po[0]), po[1], abs(po[2])))
        except Exception:
            warm.append((Ag, mug, sigg))

    p0_warm = [bl_coeffs[1], bl_coeffs[0],
               warm[0][0], warm[0][1], warm[0][2],
               warm[1][0], warm[1][1], warm[1][2],
               warm[2][0], warm[2][1], warm[2][2],
               warm[3][0], warm[3][1], warm[3][2]]

    # Clip warm-start to bounds
    p0_warm = [float(np.clip(v, lo, hi))
               for v, lo, hi in zip(p0_warm, lower, upper)]

    # 2c: Full 14-parameter TRF with warm start
    try:
        popt, pcov = curve_fit(
            four_gaussian_linear, x, y,
            p0=p0_warm,
            bounds=bounds,
            method='trf',
            maxfev=50000,
            ftol=1e-10, xtol=1e-10, gtol=1e-10
        )
        y_fit = four_gaussian_linear(x, *popt)
        r2 = r_squared(y, y_fit)
    except Exception as e2:
        warnings.warn(f'Stage 2 also failed: {e2}')
        y_fit = four_gaussian_linear(x, *popt)
        r2 = r_squared(y, y_fit)

# ── 6. Parameter uncertainties ────────────────────────────────────────────────
perr = np.sqrt(np.diag(pcov))
perr = np.where(np.isfinite(perr), perr, np.nan)

# Warn about poorly constrained parameters
names = ['a0','a1','A1','mu1','sig1','A2','mu2','sig2',
         'A3','mu3','sig3','A4','mu4','sig4']
for i, (nm, pe) in enumerate(zip(names, perr)):
    if not np.isfinite(pe):
        warnings.warn(f'Parameter {nm} is poorly constrained (pcov = inf/NaN).')

# Collinearity warnings (|corr| > 0.95)
with np.errstate(invalid='ignore'):
    std = np.sqrt(np.diag(pcov))
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        corr = pcov / np.outer(std, std)
for i in range(14):
    for j in range(i + 1, 14):
        if np.isfinite(corr[i, j]) and abs(corr[i, j]) > 0.95:
            warnings.warn(f'High correlation between {names[i]} and {names[j]}: {corr[i,j]:.3f}')

# Condition number
try:
    cond = float(np.linalg.cond(pcov))
except Exception:
    cond = float('nan')

# ── 7. Derived quantities ──────────────────────────────────────────────────────
fwhm_factor = 2.0 * np.sqrt(2.0 * np.log(2.0))

a0, a1 = popt[0], popt[1]
A1, mu1, sig1 = popt[2], popt[3], popt[4]
A2, mu2, sig2 = popt[5], popt[6], popt[7]
A3, mu3, sig3 = popt[8], popt[9], popt[10]
A4, mu4, sig4 = popt[11], popt[12], popt[13]

fwhm1 = fwhm_factor * abs(sig1)
fwhm2 = fwhm_factor * abs(sig2)
fwhm3 = fwhm_factor * abs(sig3)
fwhm4 = fwhm_factor * abs(sig4)

area1 = abs(A1) * abs(sig1) * np.sqrt(2.0 * np.pi)
area2 = abs(A2) * abs(sig2) * np.sqrt(2.0 * np.pi)
area3 = abs(A3) * abs(sig3) * np.sqrt(2.0 * np.pi)
area4 = abs(A4) * abs(sig4) * np.sqrt(2.0 * np.pi)

residuals = y - y_fit
max_abs_res = float(np.max(np.abs(residuals)))
max_abs_res_loc = float(x[np.argmax(np.abs(residuals))])

# Systematic residual structure check
threshold = 0.10 * float(np.max(y))
systematic_flag = False
min_width_pts = int(5.0 / (x[1] - x[0])) if len(x) > 1 else 5
above = np.abs(residuals) > threshold
count = 0
for val in above:
    if val:
        count += 1
        if count >= min_width_pts:
            systematic_flag = True
            break
    else:
        count = 0

rms = rmse(y, y_fit)

# ── 8. Fit quality status ─────────────────────────────────────────────────────
if r2 >= 0.90 and not systematic_flag:
    fit_status = 'CONVERGED'
elif r2 >= 0.90:
    fit_status = 'CONVERGED_WITH_SYSTEMATIC_RESIDUALS'
else:
    fit_status = 'FAILED'

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

# Individual components
bl_curve = a0 + a1 * x
g1_curve = A1 * np.exp(-0.5 * ((x - mu1) / sig1) ** 2)
g2_curve = A2 * np.exp(-0.5 * ((x - mu2) / sig2) ** 2)
g3_curve = A3 * np.exp(-0.5 * ((x - mu3) / sig3) ** 2)
g4_curve = A4 * np.exp(-0.5 * ((x - mu4) / sig4) ** 2)

ax1.plot(x, y, 'k.', ms=2, alpha=0.6, label='Data')
ax1.plot(x, y_fit, 'r-', lw=2, label=f'Total fit (R²={r2:.4f})')
ax1.plot(x, bl_curve, 'gray', lw=1, ls='--', label='Baseline')
ax1.fill_between(x, bl_curve, bl_curve + g1_curve, alpha=0.25, color='C0', label=f'Peak 1 (μ={mu1:.2f})')
ax1.fill_between(x, bl_curve, bl_curve + g2_curve, alpha=0.25, color='C1', label=f'Peak 2 (μ={mu2:.2f})')
ax1.fill_between(x, bl_curve, bl_curve + g3_curve, alpha=0.25, color='C2', label=f'Peak 3 (μ={mu3:.2f})')
ax1.fill_between(x, bl_curve, bl_curve + g4_curve, alpha=0.25, color='C3', label=f'Peak 4 (μ={mu4:.2f})')
ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title('Four-Gaussian + Linear Baseline Fit', fontsize=13)
ax1.legend(fontsize=9, loc='upper right')
ax1.set_xlim(x[0], x[-1])

ax2.plot(x, residuals, 'k-', lw=0.8, label='Residuals')
ax2.axhline(0, color='r', lw=1)
ax2.axhline(threshold, color='orange', lw=1, ls='--', alpha=0.7)
ax2.axhline(-threshold, color='orange', lw=1, ls='--', alpha=0.7, label='±10% threshold')
ax2.set_xlabel('x', fontsize=12)
ax2.set_ylabel('Residual', fontsize=12)
ax2.legend(fontsize=9)
ax2.set_xlim(x[0], x[-1])

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

# ── 10. Results JSON ──────────────────────────────────────────────────────────
def safe_float(v):
    return float(v) if np.isfinite(v) else None

results = {
    'model_type': 'Four-Gaussian + Linear Baseline (14 parameters): '
                  'y = a0 + a1*x + sum_i A_i*exp(-0.5*((x-mu_i)/sig_i)^2)',
    'parameters': {
        'baseline': {
            'a0':      safe_float(a0),
            'a0_err':  safe_float(perr[0]),
            'a1':      safe_float(a1),
            'a1_err':  safe_float(perr[1])
        },
        'peak_1': {
            'center':      safe_float(mu1),
            'center_err':  safe_float(perr[3]),
            'amplitude':   safe_float(A1),
            'amplitude_err': safe_float(perr[2]),
            'sigma':       safe_float(sig1),
            'sigma_err':   safe_float(perr[4]),
            'fwhm':        safe_float(fwhm1),
            'area':        safe_float(area1)
        },
        'peak_2': {
            'center':      safe_float(mu2),
            'center_err':  safe_float(perr[6]),
            'amplitude':   safe_float(A2),
            'amplitude_err': safe_float(perr[5]),
            'sigma':       safe_float(sig2),
            'sigma_err':   safe_float(perr[7]),
            'fwhm':        safe_float(fwhm2),
            'area':        safe_float(area2)
        },
        'peak_3': {
            'center':      safe_float(mu3),
            'center_err':  safe_float(perr[9]),
            'amplitude':   safe_float(A3),
            'amplitude_err': safe_float(perr[8]),
            'sigma':       safe_float(sig3),
            'sigma_err':   safe_float(perr[10]),
            'fwhm':        safe_float(fwhm3),
            'area':        safe_float(area3)
        },
        'peak_4': {
            'center':      safe_float(mu4),
            'center_err':  safe_float(perr[12]),
            'amplitude':   safe_float(A4),
            'amplitude_err': safe_float(perr[11]),
            'sigma':       safe_float(sig4),
            'sigma_err':   safe_float(perr[13]),
            'fwhm':        safe_float(fwhm4),
            'area':        safe_float(area4)
        }
    },
    'fit_quality': {
        'r_squared':              safe_float(r2),
        'rmse':                   safe_float(rms),
        'fit_status':             fit_status,
        'max_abs_residual':       safe_float(max_abs_res),
        'max_abs_residual_x':     safe_float(max_abs_res_loc),
        'systematic_residuals':   systematic_flag,
        'cov_condition_number':   safe_float(cond)
    },
    'summary': (
        f'Fit status: {fit_status}. '
        f'R²={r2:.4f}, RMSE={rms:.4f}. '
        f'Four Gaussian peaks identified at x≈{mu1:.2f} (FWHM={fwhm1:.2f}, area={area1:.2f}), '
        f'x≈{mu2:.2f} (FWHM={fwhm2:.2f}, area={area2:.2f}), '
        f'x≈{mu3:.2f} (FWHM={fwhm3:.2f}, area={area3:.2f}), '
        f'x≈{mu4:.2f} (FWHM={fwhm4:.2f}, area={area4:.2f}). '
        f'Linear baseline: a0={a0:.4f}, a1={a1:.6f}. '
        f'Max |residual|={max_abs_res:.4f} at x={max_abs_res_loc:.2f}. '
        f'Covariance condition number={cond:.3e}. '
        + ('Systematic residual structure detected.' if systematic_flag else 'No systematic residual structure.')
    )
}

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