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

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_easy/temp_spectrum_0.npy')

# Handle different array shapes
if data.ndim == 2:
    if data.shape[0] == 2:
        x = data[0]
        y = data[1]
    elif data.shape[1] == 2:
        x = data[:, 0]
        y = data[:, 1]
    else:
        raise ValueError(f'Unexpected 2D array shape: {data.shape}')
elif data.ndim == 1:
    # Only y values provided, generate x
    y = data
    x = np.linspace(0, 100, len(y))
else:
    raise ValueError(f'Unexpected array shape: {data.shape}')

print(f'Data loaded: {len(x)} points, x=[{x.min():.3f}, {x.max():.3f}], y=[{y.min():.6f}, {y.max():.6f}]')

# Define model: linear baseline + 3 Gaussians
def model_func(x, a0, a1, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3):
    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)
    return baseline + g1 + g2 + g3

# Initial guesses as specified in the plan
p0 = [
    0.0,   # a0
    0.0,   # a1
    4.1,   # A1
    25.0,  # mu1
    3.0,   # sigma1
    3.1,   # A2
    53.0,  # mu2
    6.0,   # sigma2
    5.15,  # A3
    80.0,  # mu3
    2.5    # sigma3
]

# Bounds as specified in the plan
# Order: a0, a1, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3
lower_bounds = [-0.1, -0.01, 0.0, 0.0, 0.1, 0.0, 0.0, 0.1, 0.0, 0.0, 0.1]
upper_bounds = [0.5, 0.01, np.inf, 100.0, np.inf, np.inf, 100.0, np.inf, np.inf, 100.0, np.inf]

# Fit using Levenberg-Marquardt (method='lm') requires no bounds,
# so use 'trf' (Trust Region Reflective) which supports bounds
try:
    popt, pcov = curve_fit(
        model_func,
        x, y,
        p0=p0,
        bounds=(lower_bounds, upper_bounds),
        method='trf',
        maxfev=100000
    )
    print('Fit converged successfully with trf method.')
except Exception as e:
    print(f'trf method failed: {e}. Trying without bounds using lm method.')
    try:
        popt, pcov = curve_fit(
            model_func,
            x, y,
            p0=p0,
            method='lm',
            maxfev=100000
        )
        print('Fit converged with lm method (no bounds).')
    except Exception as e2:
        print(f'lm method also failed: {e2}')
        raise

# Extract parameters
perr = np.sqrt(np.diag(pcov))

a0_val, a0_err = popt[0], perr[0]
a1_val, a1_err = popt[1], perr[1]
A1_val, A1_err = popt[2], perr[2]
mu1_val, mu1_err = popt[3], perr[3]
sig1_val, sig1_err = popt[4], perr[4]
A2_val, A2_err = popt[5], perr[5]
mu2_val, mu2_err = popt[6], perr[6]
sig2_val, sig2_err = popt[7], perr[7]
A3_val, A3_err = popt[8], perr[8]
mu3_val, mu3_err = popt[9], perr[9]
sig3_val, sig3_err = popt[10], perr[10]

FWHM_factor = 2.0 * np.sqrt(2.0 * np.log(2.0))  # = 2.3548

fwhm1 = FWHM_factor * abs(sig1_val)
fwhm1_err = FWHM_factor * sig1_err
fwhm2 = FWHM_factor * abs(sig2_val)
fwhm2_err = FWHM_factor * sig2_err
fwhm3 = FWHM_factor * abs(sig3_val)
fwhm3_err = FWHM_factor * sig3_err

# Compute fitted curve and residuals
y_fit = model_func(x, *popt)
residuals = y - y_fit

# Goodness of fit
ss_res = np.sum(residuals ** 2)
ss_tot = np.sum((y - np.mean(y)) ** 2)
r_squared = 1.0 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals ** 2))

print(f'R-squared: {r_squared:.6f}')
print(f'RMSE: {rmse:.6f}')

# Individual components for plotting
baseline_curve = a0_val + a1_val * x
g1_curve = A1_val * np.exp(-0.5 * ((x - mu1_val) / sig1_val) ** 2)
g2_curve = A2_val * np.exp(-0.5 * ((x - mu2_val) / sig2_val) ** 2)
g3_curve = A3_val * np.exp(-0.5 * ((x - mu3_val) / sig3_val) ** 2)

# Visualization
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]})
fig.suptitle('Spectroscopic Data: 3-Gaussian Fit with Linear Baseline', fontsize=13, fontweight='bold')

# Main plot
ax1.plot(x, y, 'k.', markersize=2, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit')
ax1.plot(x, baseline_curve + g1_curve, 'b--', linewidth=1.5, label=f'Peak 1 (\u03bc={mu1_val:.2f})')
ax1.plot(x, baseline_curve + g2_curve, 'g--', linewidth=1.5, label=f'Peak 2 (\u03bc={mu2_val:.2f})')
ax1.plot(x, baseline_curve + g3_curve, 'm--', linewidth=1.5, label=f'Peak 3 (\u03bc={mu3_val:.2f})')
ax1.plot(x, baseline_curve, 'orange', linewidth=1, linestyle=':', label='Baseline')
ax1.fill_between(x, baseline_curve, baseline_curve + g1_curve, alpha=0.15, color='blue')
ax1.fill_between(x, baseline_curve, baseline_curve + g2_curve, alpha=0.15, color='green')
ax1.fill_between(x, baseline_curve, baseline_curve + g3_curve, alpha=0.15, color='purple')
ax1.set_ylabel('Intensity', fontsize=11)
ax1.legend(fontsize=9, loc='upper left')
ax1.grid(True, alpha=0.3)
ax1.text(0.98, 0.95, f'R\u00b2 = {r_squared:.5f}\nRMSE = {rmse:.5f}',
         transform=ax1.transAxes, ha='right', va='top',
         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5), fontsize=9)

# Residuals plot
ax2.plot(x, residuals, 'k.', markersize=2, alpha=0.5)
ax2.axhline(0, color='r', linewidth=1)
ax2.set_xlabel('x (channel / wavelength)', fontsize=11)
ax2.set_ylabel('Residuals', fontsize=11)
ax2.grid(True, alpha=0.3)
ax2.set_title(f'Residuals (std = {np.std(residuals):.5f})', fontsize=10)

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

# Build results dictionary
results = {
    'model_type': '3 Gaussian peaks on linear baseline: y = a0 + a1*x + Sum_i A_i*exp(-0.5*((x-mu_i)/sigma_i)^2)',
    'parameters': {
        'peak_1': {
            'center': round(float(mu1_val), 6),
            'center_err': round(float(mu1_err), 6),
            'amplitude': round(float(A1_val), 6),
            'amplitude_err': round(float(A1_err), 6),
            'sigma': round(float(sig1_val), 6),
            'sigma_err': round(float(sig1_err), 6),
            'fwhm': round(float(fwhm1), 6),
            'fwhm_err': round(float(fwhm1_err), 6)
        },
        'peak_2': {
            'center': round(float(mu2_val), 6),
            'center_err': round(float(mu2_err), 6),
            'amplitude': round(float(A2_val), 6),
            'amplitude_err': round(float(A2_err), 6),
            'sigma': round(float(sig2_val), 6),
            'sigma_err': round(float(sig2_err), 6),
            'fwhm': round(float(fwhm2), 6),
            'fwhm_err': round(float(fwhm2_err), 6)
        },
        'peak_3': {
            'center': round(float(mu3_val), 6),
            'center_err': round(float(mu3_err), 6),
            'amplitude': round(float(A3_val), 6),
            'amplitude_err': round(float(A3_err), 6),
            'sigma': round(float(sig3_val), 6),
            'sigma_err': round(float(sig3_err), 6),
            'fwhm': round(float(fwhm3), 6),
            'fwhm_err': round(float(fwhm3_err), 6)
        },
        'baseline': {
            'intercept_a0': round(float(a0_val), 6),
            'intercept_a0_err': round(float(a0_err), 6),
            'slope_a1': round(float(a1_val), 8),
            'slope_a1_err': round(float(a1_err), 8)
        }
    },
    'fit_quality': {
        'r_squared': round(float(r_squared), 8),
        'rmse': round(float(rmse), 8),
        'residual_std': round(float(np.std(residuals)), 8)
    },
    'summary': (
        f'Successfully fit 3 Gaussian peaks on a linear baseline to {len(x)}-point spectrum. '
        f'Peak 1: center={mu1_val:.3f}, FWHM={fwhm1:.3f}, amplitude={A1_val:.4f}. '
        f'Peak 2: center={mu2_val:.3f}, FWHM={fwhm2:.3f}, amplitude={A2_val:.4f}. '
        f'Peak 3: center={mu3_val:.3f}, FWHM={fwhm3:.3f}, amplitude={A3_val:.4f}. '
        f'Baseline: a0={a0_val:.5f}, a1={a1_val:.7f} (near-zero as expected). '
        f'Fit quality: R²={r_squared:.6f}, RMSE={rmse:.6f}. '
        f'11 free parameters fit using scipy.optimize.curve_fit with trf (Levenberg-Marquardt-like) algorithm.'
    )
}

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