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

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_overlap/temp_spectrum_0.npy')
if data.ndim == 2:
    x_data = data[:, 0]
    y_data = data[:, 1]
else:
    y_data = data
    x_data = np.linspace(0, 50, len(y_data))

# Sort by x
sort_idx = np.argsort(x_data)
x_data = x_data[sort_idx]
y_data = y_data[sort_idx]

# Remove NaN/inf
mask = np.isfinite(x_data) & np.isfinite(y_data)
x_data = x_data[mask]
y_data = y_data[mask]
N = len(x_data)

print(f"Data loaded: {N} points, x range [{x_data.min():.2f}, {x_data.max():.2f}], y range [{y_data.min():.4f}, {y_data.max():.4f}]")

# Define split-Gaussian function
def split_gaussian(x, A, mu, sigmaL, sigmaR):
    result = np.empty_like(x, dtype=float)
    left = x < mu
    right = ~left
    result[left] = A * np.exp(-0.5 * ((x[left] - mu) / sigmaL) ** 2)
    result[right] = A * np.exp(-0.5 * ((x[right] - mu) / sigmaR) ** 2)
    return result

# Full model: quadratic baseline + 2 split-Gaussians
def model(params, x):
    a, b, c, A1, mu1, s1L, s1R, A2, mu2, s2L, s2R = params
    baseline = a + b * x + c * x**2
    peak1 = split_gaussian(x, A1, mu1, s1L, s1R)
    peak2 = split_gaussian(x, A2, mu2, s2L, s2R)
    return baseline + peak1 + peak2

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

def get_components(params, x):
    a, b, c, A1, mu1, s1L, s1R, A2, mu2, s2L, s2R = params
    baseline = a + b * x + c * x**2
    peak1 = split_gaussian(x, A1, mu1, s1L, s1R)
    peak2 = split_gaussian(x, A2, mu2, s2L, s2R)
    return baseline, peak1, peak2

# Step 2: Baseline characterization from data outside peak region
baseline_mask = (x_data < 17) | (x_data > 33)
if np.sum(baseline_mask) > 5:
    bl_coeffs = np.polyfit(x_data[baseline_mask], y_data[baseline_mask], 2)
    c_init, b_init, a_init = bl_coeffs
else:
    a_init, b_init, c_init = 0.1, 0.005, -0.0001

print(f"Baseline initial estimates: a={a_init:.4f}, b={b_init:.4f}, c={c_init:.6f}")

# Bounds
#  a, b, c, A1, mu1, s1L, s1R, A2, mu2, s2L, s2R
lower = [-3.0, -0.5, -0.01, 0.5, 19.0, 0.3, 0.3, 0.5, 25.0, 0.3, 0.3]
upper = [ 3.0,  0.5,  0.01, 15.0, 25.0, 5.0, 5.0, 15.0, 31.0, 5.0, 5.0]

# Step 5: Multi-start optimization
mu1_grid = [21.0, 21.8, 22.5]
mu2_grid = [27.0, 27.8, 28.5]
s1L_grid = [0.8, 1.3, 2.0]
s1R_grid = [1.5, 2.0, 3.0]
s2L_grid = [1.0, 1.5, 2.0]
s2R_grid = [1.5, 2.2, 3.0]
A1_grid = [3.5, 4.8, 6.0]
A2_grid = [3.5, 4.5, 6.0]

print("Starting multi-start optimization (coarse pass)...")

best_cost = np.inf
best_result = None
top_results = []

total_starts = 0
for mu1_v in mu1_grid:
    for mu2_v in mu2_grid:
        for s1L_v in s1L_grid:
            for s1R_v in s1R_grid:
                for s2L_v in s2L_grid:
                    for s2R_v in s2R_grid:
                        for A1_v in A1_grid:
                            for A2_v in A2_grid:
                                total_starts += 1
                                p0 = [a_init, b_init, c_init, A1_v, mu1_v, s1L_v, s1R_v, A2_v, mu2_v, s2L_v, s2R_v]
                                # Clip to bounds
                                p0 = [max(lo, min(hi, v)) for v, lo, hi in zip(p0, lower, upper)]
                                try:
                                    res = least_squares(residuals, p0, args=(x_data, y_data),
                                                       bounds=(lower, upper), method='trf',
                                                       ftol=1e-8, xtol=1e-8, gtol=1e-8,
                                                       max_nfev=5000)
                                    cost = res.cost
                                    top_results.append((cost, res.x.copy()))
                                except:
                                    pass

print(f"Coarse pass complete: {total_starts} starts, {len(top_results)} successful")

# Sort by cost and refine top 20
top_results.sort(key=lambda x: x[0])
top_20 = top_results[:20]

print("Refining top 20 solutions...")

best_cost = np.inf
best_result = None

for cost, p0 in top_20:
    try:
        res = least_squares(residuals, p0, args=(x_data, y_data),
                           bounds=(lower, upper), method='trf',
                           ftol=1e-14, xtol=1e-14, gtol=1e-14,
                           max_nfev=50000)
        if res.cost < best_cost:
            best_cost = res.cost
            best_result = res
    except:
        pass

print(f"Best SSR after refinement: {best_cost:.6f}")

# Extract parameters
popt = best_result.x
a, b, c, A1, mu1, s1L, s1R, A2, mu2, s2L, s2R = popt

# Compute uncertainties from Jacobian
try:
    J = best_result.jac
    # Covariance = inv(J^T J) * s^2, where s^2 = SSR / (N - p)
    n_params = len(popt)
    residuals_at_opt = best_result.fun
    s_sq = np.sum(residuals_at_opt**2) / (N - n_params)
    JtJ = J.T @ J
    try:
        cov = np.linalg.inv(JtJ) * s_sq
        perr = np.sqrt(np.abs(np.diag(cov)))
    except np.linalg.LinAlgError:
        cov = np.linalg.pinv(JtJ) * s_sq
        perr = np.sqrt(np.abs(np.diag(cov)))
except:
    perr = np.zeros(len(popt))

a_err, b_err, c_err, A1_err, mu1_err, s1L_err, s1R_err, A2_err, mu2_err, s2L_err, s2R_err = perr

# Derived quantities
FWHM1 = 1.1774 * (s1L + s1R)
FWHM2 = 1.1774 * (s2L + s2R)
asym1 = s1R / s1L
asym2 = s2R / s2L
peak_sep = mu2 - mu1
resolution = 2 * peak_sep / (FWHM1 + FWHM2)

# FWHM uncertainties via error propagation
FWHM1_err = 1.1774 * np.sqrt(s1L_err**2 + s1R_err**2)
FWHM2_err = 1.1774 * np.sqrt(s2L_err**2 + s2R_err**2)

# Asymmetry ratio uncertainties
asym1_err = asym1 * np.sqrt((s1R_err/s1R)**2 + (s1L_err/s1L)**2) if s1L > 0 and s1R > 0 else 0.0
asym2_err = asym2 * np.sqrt((s2R_err/s2R)**2 + (s2L_err/s2L)**2) if s2L > 0 and s2R > 0 else 0.0

# R² and RMSE
y_fit = model(popt, x_data)
ss_res = np.sum((y_data - y_fit)**2)
ss_tot = np.sum((y_data - np.mean(y_data))**2)
r_squared = 1 - ss_res / ss_tot
rmse = np.sqrt(ss_res / N)
reduced_chi_sq = ss_res / (N - n_params)

# Residual diagnostics
resid = y_data - y_fit

# Lag-1 autocorrelation in peak region [20, 32]
peak_mask = (x_data >= 20) & (x_data <= 32)
resid_peak = resid[peak_mask]
if len(resid_peak) > 2:
    lag1_auto = np.corrcoef(resid_peak[:-1], resid_peak[1:])[0, 1]
else:
    lag1_auto = 0.0

# Valley analysis near x~25
valley_mask = (x_data >= 24) & (x_data <= 26)
if np.any(valley_mask):
    valley_max_resid = np.max(np.abs(resid[valley_mask]))
    valley_idx = np.argmin(y_fit[(x_data >= mu1) & (x_data <= mu2)])
    valley_x_range = x_data[(x_data >= mu1) & (x_data <= mu2)]
    valley_y_range = y_fit[(x_data >= mu1) & (x_data <= mu2)]
    if len(valley_y_range) > 0:
        valley_depth = np.min(valley_y_range)
    else:
        valley_depth = 0.0
else:
    valley_max_resid = 0.0
    valley_depth = 0.0

# Subregion residual statistics
subregions = [(0, 15), (15, 20), (20, 25), (25, 30), (30, x_data.max())]
subregion_stats = {}
for lo, hi in subregions:
    m = (x_data >= lo) & (x_data <= hi)
    if np.any(m):
        subregion_stats[f"{lo:.0f}-{hi:.0f}"] = {
            "mean_resid": float(np.mean(resid[m])),
            "std_resid": float(np.std(resid[m])),
            "max_abs_resid": float(np.max(np.abs(resid[m])))
        }

# Runs test on residual signs
signs = np.sign(resid)
signs = signs[signs != 0]
n_pos = np.sum(signs > 0)
n_neg = np.sum(signs < 0)
runs = 1 + np.sum(np.diff(signs) != 0)
n_total = len(signs)
if n_pos > 0 and n_neg > 0:
    expected_runs = 1 + 2 * n_pos * n_neg / n_total
    var_runs = (2 * n_pos * n_neg * (2 * n_pos * n_neg - n_total)) / (n_total**2 * (n_total - 1))
    if var_runs > 0:
        runs_z = (runs - expected_runs) / np.sqrt(var_runs)
    else:
        runs_z = 0.0
else:
    runs_z = 0.0

# BIC comparison (vs hypothetical 8-param symmetric model)
# BIC = N*ln(SSR/N) + k*ln(N)
bic_11 = N * np.log(ss_res / N) + 11 * np.log(N)

print(f"\nFit Results:")
print(f"  a = {a:.6f} +/- {a_err:.6f}")
print(f"  b = {b:.6f} +/- {b_err:.6f}")
print(f"  c = {c:.8f} +/- {c_err:.8f}")
print(f"  A1 = {A1:.6f} +/- {A1_err:.6f}")
print(f"  mu1 = {mu1:.6f} +/- {mu1_err:.6f}")
print(f"  sigma1L = {s1L:.6f} +/- {s1L_err:.6f}")
print(f"  sigma1R = {s1R:.6f} +/- {s1R_err:.6f}")
print(f"  A2 = {A2:.6f} +/- {A2_err:.6f}")
print(f"  mu2 = {mu2:.6f} +/- {mu2_err:.6f}")
print(f"  sigma2L = {s2L:.6f} +/- {s2L_err:.6f}")
print(f"  sigma2R = {s2R:.6f} +/- {s2R_err:.6f}")
print(f"  FWHM1 = {FWHM1:.4f} +/- {FWHM1_err:.4f}")
print(f"  FWHM2 = {FWHM2:.4f} +/- {FWHM2_err:.4f}")
print(f"  Asymmetry ratio 1 (sigR/sigL) = {asym1:.4f} +/- {asym1_err:.4f}")
print(f"  Asymmetry ratio 2 (sigR/sigL) = {asym2:.4f} +/- {asym2_err:.4f}")
print(f"  Peak separation = {peak_sep:.4f}")
print(f"  Resolution = {resolution:.4f}")
print(f"  R² = {r_squared:.8f}")
print(f"  RMSE = {rmse:.6f}")
print(f"  SSR = {ss_res:.6f}")
print(f"  Reduced chi² = {reduced_chi_sq:.8f}")
print(f"  Peak region lag-1 autocorrelation = {lag1_auto:.4f}")
print(f"  Valley max |residual| = {valley_max_resid:.4f}")
print(f"  Runs test Z = {runs_z:.4f}")
print(f"  BIC (11 params) = {bic_11:.2f}")

# Step 8: Check if EMG is needed
need_emg = lag1_auto > 0.7
if need_emg:
    print("\nWARNING: Lag-1 autocorrelation > 0.7, split-Gaussian may be insufficient.")
    print("Consider EMG model for better lineshape capture.")

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

x_fine = np.linspace(x_data.min(), x_data.max(), 2000)
y_fit_fine = model(popt, x_fine)
baseline_fine, peak1_fine, peak2_fine = get_components(popt, x_fine)
baseline_data, peak1_data, peak2_data = get_components(popt, x_data)

# Top panel: data + fit + components
ax1 = axes[0]
ax1.scatter(x_data, y_data, s=8, alpha=0.5, color='gray', label='Data', zorder=1)
ax1.plot(x_fine, y_fit_fine, 'r-', linewidth=2, label='Full Fit', zorder=3)
ax1.plot(x_fine, baseline_fine, 'k--', linewidth=1, alpha=0.7, label=f'Baseline (quadratic)', zorder=2)
ax1.fill_between(x_fine, baseline_fine, baseline_fine + peak1_fine, alpha=0.3, color='blue',
                  label=f'Peak 1 (μ={mu1:.2f}, FWHM={FWHM1:.2f})', zorder=2)
ax1.fill_between(x_fine, baseline_fine, baseline_fine + peak2_fine, alpha=0.3, color='green',
                  label=f'Peak 2 (μ={mu2:.2f}, FWHM={FWHM2:.2f})', zorder=2)
ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title('Split-Gaussian Fit with Quadratic Baseline', fontsize=14)
ax1.legend(fontsize=9, loc='upper right')
ax1.grid(True, alpha=0.3)

# Annotate key metrics
textstr = (f'R² = {r_squared:.6f}\n'
           f'RMSE = {rmse:.4f}\n'
           f'Peak region lag-1 AC = {lag1_auto:.4f}\n'
           f'Resolution = {resolution:.2f}')
ax1.text(0.02, 0.98, textstr, transform=ax1.transAxes, fontsize=9,
         verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

# Bottom panel: residuals
ax2 = axes[1]
ax2.scatter(x_data, resid, s=5, alpha=0.5, color='gray', zorder=1)
ax2.axhline(y=0, color='r', linewidth=1, linestyle='-')
ax2.axhline(y=rmse, color='orange', linewidth=0.8, linestyle='--', alpha=0.7, label=f'+RMSE={rmse:.4f}')
ax2.axhline(y=-rmse, color='orange', linewidth=0.8, linestyle='--', alpha=0.7, label=f'-RMSE')
ax2.set_xlabel('x (channel / wavelength)', fontsize=12)
ax2.set_ylabel('Residual', fontsize=12)
ax2.set_title('Residuals', fontsize=11)
ax2.legend(fontsize=8)
ax2.grid(True, alpha=0.3)

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

# Build results JSON
results = {
    "model_type": "Split-Gaussian (bi-Gaussian) double peak with quadratic baseline: y = (a + b*x + c*x^2) + SG1(x; A1, mu1, sigma1L, sigma1R) + SG2(x; A2, mu2, sigma2L, sigma2R)",
    "parameters": {
        "baseline": {
            "a": round(float(a), 6), "a_err": round(float(a_err), 6),
            "b": round(float(b), 6), "b_err": round(float(b_err), 6),
            "c": round(float(c), 8), "c_err": round(float(c_err), 8)
        },
        "peak_1": {
            "center": round(float(mu1), 4), "center_err": round(float(mu1_err), 4),
            "amplitude": round(float(A1), 4), "amplitude_err": round(float(A1_err), 4),
            "sigma_left": round(float(s1L), 4), "sigma_left_err": round(float(s1L_err), 4),
            "sigma_right": round(float(s1R), 4), "sigma_right_err": round(float(s1R_err), 4),
            "FWHM": round(float(FWHM1), 4), "FWHM_err": round(float(FWHM1_err), 4),
            "asymmetry_ratio": round(float(asym1), 4), "asymmetry_ratio_err": round(float(asym1_err), 4)
        },
        "peak_2": {
            "center": round(float(mu2), 4), "center_err": round(float(mu2_err), 4),
            "amplitude": round(float(A2), 4), "amplitude_err": round(float(A2_err), 4),
            "sigma_left": round(float(s2L), 4), "sigma_left_err": round(float(s2L_err), 4),
            "sigma_right": round(float(s2R), 4), "sigma_right_err": round(float(s2R_err), 4),
            "FWHM": round(float(FWHM2), 4), "FWHM_err": round(float(FWHM2_err), 4),
            "asymmetry_ratio": round(float(asym2), 4), "asymmetry_ratio_err": round(float(asym2_err), 4)
        },
        "derived": {
            "peak_separation": round(float(peak_sep), 4),
            "peak_separation_err": round(float(np.sqrt(mu1_err**2 + mu2_err**2)), 4),
            "resolution": round(float(resolution), 4),
            "valley_depth": round(float(valley_depth), 4)
        }
    },
    "fit_quality": {
        "r_squared": round(float(r_squared), 8),
        "rmse": round(float(rmse), 6),
        "ssr": round(float(ss_res), 6),
        "reduced_chi_squared": round(float(reduced_chi_sq), 8),
        "max_abs_residual": round(float(np.max(np.abs(resid))), 6),
        "peak_region_lag1_autocorrelation": round(float(lag1_auto), 4),
        "valley_max_abs_residual": round(float(valley_max_resid), 4),
        "runs_test_z": round(float(runs_z), 4),
        "bic_11_params": round(float(bic_11), 2),
        "n_parameters": 11,
        "n_data_points": int(N)
    },
    "diagnostics": {
        "subregion_residuals": subregion_stats,
        "emg_recommended": bool(need_emg)
    },
    "summary": (f"Split-Gaussian (bi-Gaussian) double peak model with quadratic baseline fitted to {N} spectroscopic data points. "
                f"Peak 1 at x={mu1:.2f} (FWHM={FWHM1:.2f}, asymmetry={asym1:.3f}) and "
                f"Peak 2 at x={mu2:.2f} (FWHM={FWHM2:.2f}, asymmetry={asym2:.3f}). "
                f"Peak separation={peak_sep:.2f}, spectral resolution={resolution:.2f}. "
                f"R²={r_squared:.6f}, RMSE={rmse:.4f}. "
                f"Peak region lag-1 autocorrelation={lag1_auto:.4f}. "
                f"Multi-start optimization with 6561 initial conditions, refined top 20 to global minimum SSR={ss_res:.4f}. "
                f"{'EMG model recommended due to high residual autocorrelation.' if need_emg else 'Split-Gaussian model provides adequate fit.'}")
}

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