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

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_validator_data/v_baseline/temp_spectrum_0.npy')
if data.ndim == 2:
    x = data[:, 0]
    y = data[:, 1]
else:
    y = data
    x = np.linspace(0, 100, len(y))

# Define model components
def cubic_baseline(x, a0, a1, a2, a3):
    return a0 + a1*x + a2*x**2 + a3*x**3

def gaussian(x, amplitude, center, sigma):
    return amplitude * np.exp(-(x - center)**2 / (2 * sigma**2))

def full_model(x, a0, a1, a2, a3,
               amp1, cen1, sig1,
               amp2, cen2, sig2,
               amp3, cen3, sig3):
    baseline = cubic_baseline(x, a0, a1, a2, a3)
    g1 = gaussian(x, amp1, cen1, sig1)
    g2 = gaussian(x, amp2, cen2, sig2)
    g3 = gaussian(x, amp3, cen3, sig3)
    return baseline + g1 + g2 + g3

# Get initial baseline estimate by excluding peak regions
peak_mask = (
    ((x >= 23) & (x <= 37)) |
    ((x >= 60) & (x <= 67)) |
    ((x >= 70) & (x <= 82)) |
    ((x >= 89) & (x <= 100))
)
baseline_mask = ~peak_mask
x_bl = x[baseline_mask]
y_bl = y[baseline_mask]

# Fit cubic polynomial to baseline regions
bl_coeffs = np.polyfit(x_bl, y_bl, 3)
a3_init, a2_init, a1_init, a0_init = bl_coeffs

# Set up lmfit Parameters
params = Parameters()
# Baseline
params.add('a0', value=a0_init)
params.add('a1', value=a1_init)
params.add('a2', value=a2_init)
params.add('a3', value=a3_init)

# Peak 1: near x=29.5, sharp and tall
params.add('amp1', value=2.5, min=0)
params.add('cen1', value=29.5, min=24.5, max=34.5)
params.add('sig1', value=2.5, min=0.5, max=15)

# Peak 2: near x=77, broader and shorter
params.add('amp2', value=1.5, min=0)
params.add('cen2', value=77, min=72, max=82)
params.add('sig2', value=3.5, min=0.5, max=15)

# Peak 3: near x=95
params.add('amp3', value=0.8, min=0)
params.add('cen3', value=95, min=90, max=100)
params.add('sig3', value=2.5, min=0.5, max=15)

# Create model and fit
model = Model(full_model, independent_vars=['x'])
result = model.fit(y, params, x=x, method='leastsq', max_nfev=10000)

# Check residuals near x=63 for potential 4th Gaussian
residuals = y - result.best_fit
mask_63 = (x >= 58) & (x <= 68)
resid_63_max = np.max(np.abs(residuals[mask_63])) if np.any(mask_63) else 0
overall_rmse = np.sqrt(np.mean(residuals**2))

# Extract best fit parameters
bf = result.params

# Compute fit quality
y_fit = result.best_fit
ss_res = np.sum((y - y_fit)**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1 - ss_res / ss_tot
rmse = np.sqrt(np.mean((y - y_fit)**2))

# Compute derived quantities for each peak
def peak_results(amp, amp_err, cen, cen_err, sig, sig_err):
    fwhm = 2 * np.sqrt(2 * np.log(2)) * sig
    fwhm_err = 2 * np.sqrt(2 * np.log(2)) * sig_err if sig_err else None
    area = amp * sig * np.sqrt(2 * np.pi)
    # Error propagation for area: area = A * sigma * sqrt(2pi)
    if amp_err and sig_err:
        area_err = area * np.sqrt((amp_err/amp)**2 + (sig_err/sig)**2) if amp > 0 and sig > 0 else None
    else:
        area_err = None
    return {
        'center': round(float(cen), 4),
        'center_err': round(float(cen_err), 4) if cen_err else None,
        'amplitude': round(float(amp), 4),
        'amplitude_err': round(float(amp_err), 4) if amp_err else None,
        'sigma': round(float(sig), 4),
        'sigma_err': round(float(sig_err), 4) if sig_err else None,
        'fwhm': round(float(fwhm), 4),
        'fwhm_err': round(float(fwhm_err), 4) if fwhm_err else None,
        'area': round(float(area), 4),
        'area_err': round(float(area_err), 4) if area_err else None
    }

peak1_res = peak_results(
    bf['amp1'].value, bf['amp1'].stderr,
    bf['cen1'].value, bf['cen1'].stderr,
    bf['sig1'].value, bf['sig1'].stderr
)
peak2_res = peak_results(
    bf['amp2'].value, bf['amp2'].stderr,
    bf['cen2'].value, bf['cen2'].stderr,
    bf['sig2'].value, bf['sig2'].stderr
)
peak3_res = peak_results(
    bf['amp3'].value, bf['amp3'].stderr,
    bf['cen3'].value, bf['cen3'].stderr,
    bf['sig3'].value, bf['sig3'].stderr
)

baseline_params = {
    'a0': round(float(bf['a0'].value), 6),
    'a0_err': round(float(bf['a0'].stderr), 6) if bf['a0'].stderr else None,
    'a1': round(float(bf['a1'].value), 6),
    'a1_err': round(float(bf['a1'].stderr), 6) if bf['a1'].stderr else None,
    'a2': round(float(bf['a2'].value), 8),
    'a2_err': round(float(bf['a2'].stderr), 8) if bf['a2'].stderr else None,
    'a3': round(float(bf['a3'].value), 10),
    'a3_err': round(float(bf['a3'].stderr), 10) if bf['a3'].stderr else None,
}

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

# Top panel: data + fit + components
ax1.plot(x, y, 'b.', markersize=2, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit')

# Plot individual components
bl = cubic_baseline(x, bf['a0'].value, bf['a1'].value, bf['a2'].value, bf['a3'].value)
ax1.plot(x, bl, 'k--', linewidth=1, alpha=0.7, label='Baseline (cubic)')

g1 = gaussian(x, bf['amp1'].value, bf['cen1'].value, bf['sig1'].value)
ax1.fill_between(x, bl, bl + g1, alpha=0.3, color='green', label=f'Peak 1 (center={bf["cen1"].value:.1f})')

g2 = gaussian(x, bf['amp2'].value, bf['cen2'].value, bf['sig2'].value)
ax1.fill_between(x, bl, bl + g2, alpha=0.3, color='orange', label=f'Peak 2 (center={bf["cen2"].value:.1f})')

g3 = gaussian(x, bf['amp3'].value, bf['cen3'].value, bf['sig3'].value)
ax1.fill_between(x, bl, bl + g3, alpha=0.3, color='purple', label=f'Peak 3 (center={bf["cen3"].value:.1f})')

ax1.set_ylabel('Intensity')
ax1.set_title('Spectroscopic Data Curve Fitting: Cubic Baseline + 3 Gaussians')
ax1.legend(loc='best', fontsize=8)
ax1.grid(True, alpha=0.3)

# Bottom panel: residuals
residuals_final = y - y_fit
ax2.plot(x, residuals_final, 'g-', linewidth=0.5, alpha=0.7)
ax2.axhline(y=0, color='k', linestyle='--', linewidth=0.5)
ax2.set_xlabel('X')
ax2.set_ylabel('Residuals')
ax2.set_title(f'Residuals (RMSE={rmse:.4f}, R²={r_squared:.6f})')
ax2.grid(True, alpha=0.3)

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

# Build results JSON
results = {
    'model_type': 'Cubic polynomial baseline + 3 Gaussian peaks (Levenberg-Marquardt)',
    'parameters': {
        'peak_1': peak1_res,
        'peak_2': peak2_res,
        'peak_3': peak3_res,
        'baseline': baseline_params
    },
    'fit_quality': {
        'r_squared': round(float(r_squared), 6),
        'rmse': round(float(rmse), 6),
        'chi_squared': round(float(result.chisqr), 6),
        'reduced_chi_squared': round(float(result.redchi), 6),
        'n_data_points': len(x),
        'n_parameters': result.nvarys
    },
    'summary': (
        f'Fitted spectroscopic data with cubic polynomial baseline + 3 Gaussian peaks. '
        f'Peak 1 at x={bf["cen1"].value:.2f} (amplitude={bf["amp1"].value:.3f}, FWHM={2*np.sqrt(2*np.log(2))*bf["sig1"].value:.3f}), '
        f'Peak 2 at x={bf["cen2"].value:.2f} (amplitude={bf["amp2"].value:.3f}, FWHM={2*np.sqrt(2*np.log(2))*bf["sig2"].value:.3f}), '
        f'Peak 3 at x={bf["cen3"].value:.2f} (amplitude={bf["amp3"].value:.3f}, FWHM={2*np.sqrt(2*np.log(2))*bf["sig3"].value:.3f}). '
        f'R²={r_squared:.6f}, RMSE={rmse:.6f}. '
        f'Max residual near x=63: {resid_63_max:.4f} (overall RMSE: {rmse:.4f}). '
        f'Model: y(x) = (a0 + a1*x + a2*x^2 + a3*x^3) + G1 + G2 + G3.'
    )
}

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