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

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

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

# Model: linear baseline + 5 Gaussians
def gaussian(x, A, mu, sigma):
    return A * np.exp(-(x - mu)**2 / (2 * sigma**2))

def model(x, a, b, A1, mu1, sigma1, A2, mu2, sigma2, A3, mu3, sigma3, A4, mu4, sigma4, A5, mu5, sigma5):
    baseline = a * x + b
    g1 = gaussian(x, A1, mu1, sigma1)
    g2 = gaussian(x, A2, mu2, sigma2)
    g3 = gaussian(x, A3, mu3, sigma3)
    g4 = gaussian(x, A4, mu4, sigma4)
    g5 = gaussian(x, A5, mu5, sigma5)
    return baseline + g1 + g2 + g3 + g4 + g5

# Initial guesses as specified in plan
p0 = [
    0.004, 0.0,      # a, b (baseline)
    1.5, 33, 3,      # G1: left side of peak 1
    1.0, 37, 5,      # G2: right shoulder of peak 1
    0.5, 47, 5,      # G3: intermediate bump
    1.3, 71, 3,      # G4: left side of peak 2
    0.7, 74, 4       # G5: right shoulder of peak 2
]

# Bounds as specified: amplitudes > 0, sigmas > 0, centers ±5 of initial
lower = [
    -np.inf, -np.inf,   # a, b
    0.0, 28, 0.1,       # G1
    0.0, 32, 0.1,       # G2
    0.0, 42, 0.1,       # G3
    0.0, 66, 0.1,       # G4
    0.0, 69, 0.1        # G5
]
upper = [
    np.inf, np.inf,     # a, b
    np.inf, 38, 20,     # G1
    np.inf, 42, 20,     # G2
    np.inf, 52, 20,     # G3
    np.inf, 76, 20,     # G4
    np.inf, 79, 20      # G5
]

try:
    popt, pcov = curve_fit(model, x, y, p0=p0, bounds=(lower, upper), method='trf', maxfev=50000)
    perr = np.sqrt(np.diag(pcov))
except Exception as e:
    print(f"First attempt failed: {e}. Trying with relaxed bounds...")
    # Relax bounds
    lower_relaxed = [
        -np.inf, -np.inf,
        0.0, 20, 0.1,
        0.0, 25, 0.1,
        0.0, 35, 0.1,
        0.0, 60, 0.1,
        0.0, 60, 0.1
    ]
    upper_relaxed = [
        np.inf, np.inf,
        np.inf, 45, 30,
        np.inf, 50, 30,
        np.inf, 60, 30,
        np.inf, 85, 30,
        np.inf, 85, 30
    ]
    popt, pcov = curve_fit(model, x, y, p0=p0, bounds=(lower_relaxed, upper_relaxed), method='trf', maxfev=50000)
    perr = np.sqrt(np.diag(pcov))

# Extract parameters
a, b = popt[0], popt[1]
a_err, b_err = perr[0], perr[1]

peak_params = []
for i in range(5):
    idx = 2 + i * 3
    peak_params.append({
        'amplitude': popt[idx],
        'center': popt[idx+1],
        'sigma': popt[idx+2],
        'amplitude_err': perr[idx],
        'center_err': perr[idx+1],
        'sigma_err': perr[idx+2]
    })

# Compute fit quality
y_fit = model(x, *popt)
residuals = y - y_fit
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals**2))

# Visualization
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), gridspec_kw={'height_ratios': [3, 1]})
fig.suptitle('Spectroscopic Data Curve Fitting: 5 Gaussians + Linear Baseline', fontsize=14)

# Main plot
ax1.scatter(x, y, s=1, alpha=0.5, color='gray', label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label=f'Total Fit (R²={r_squared:.6f})')

# Plot baseline
baseline = a * x + b
ax1.plot(x, baseline, 'k--', linewidth=1, alpha=0.7, label=f'Baseline: {a:.5f}x + {b:.4f}')

# Plot individual Gaussian components
colors = ['blue', 'green', 'orange', 'purple', 'cyan']
names = ['G1 (Peak1 left)', 'G2 (Peak1 right)', 'G3 (Shoulder)', 'G4 (Peak2 left)', 'G5 (Peak2 right)']
for i, (pp, color, name) in enumerate(zip(peak_params, colors, names)):
    gi = gaussian(x, pp['amplitude'], pp['center'], pp['sigma'])
    ax1.plot(x, baseline + gi, '--', color=color, linewidth=1, alpha=0.7,
             label=f"{name}: c={pp['center']:.2f}, A={pp['amplitude']:.3f}, σ={pp['sigma']:.2f}")
    ax1.fill_between(x, baseline, baseline + gi, color=color, alpha=0.1)

ax1.set_ylabel('Intensity')
ax1.legend(fontsize=7, loc='upper left')
ax1.set_xlim(x.min(), x.max())

# Residuals plot
ax2.plot(x, residuals, '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_xlim(x.min(), x.max())

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

# Build results JSON
results = {
    "model_type": "5 Gaussian peaks + linear baseline",
    "parameters": {
        "baseline": {
            "slope_a": float(a),
            "slope_a_err": float(a_err),
            "intercept_b": float(b),
            "intercept_b_err": float(b_err)
        },
        "peak_1a": {
            "description": "Left side of peak 1",
            "center": float(peak_params[0]['center']),
            "center_err": float(peak_params[0]['center_err']),
            "amplitude": float(peak_params[0]['amplitude']),
            "amplitude_err": float(peak_params[0]['amplitude_err']),
            "sigma": float(peak_params[0]['sigma']),
            "sigma_err": float(peak_params[0]['sigma_err'])
        },
        "peak_1b": {
            "description": "Right shoulder of peak 1",
            "center": float(peak_params[1]['center']),
            "center_err": float(peak_params[1]['center_err']),
            "amplitude": float(peak_params[1]['amplitude']),
            "amplitude_err": float(peak_params[1]['amplitude_err']),
            "sigma": float(peak_params[1]['sigma']),
            "sigma_err": float(peak_params[1]['sigma_err'])
        },
        "peak_2": {
            "description": "Intermediate shoulder/bump",
            "center": float(peak_params[2]['center']),
            "center_err": float(peak_params[2]['center_err']),
            "amplitude": float(peak_params[2]['amplitude']),
            "amplitude_err": float(peak_params[2]['amplitude_err']),
            "sigma": float(peak_params[2]['sigma']),
            "sigma_err": float(peak_params[2]['sigma_err'])
        },
        "peak_3a": {
            "description": "Left side of peak 2",
            "center": float(peak_params[3]['center']),
            "center_err": float(peak_params[3]['center_err']),
            "amplitude": float(peak_params[3]['amplitude']),
            "amplitude_err": float(peak_params[3]['amplitude_err']),
            "sigma": float(peak_params[3]['sigma']),
            "sigma_err": float(peak_params[3]['sigma_err'])
        },
        "peak_3b": {
            "description": "Right shoulder of peak 2",
            "center": float(peak_params[4]['center']),
            "center_err": float(peak_params[4]['center_err']),
            "amplitude": float(peak_params[4]['amplitude']),
            "amplitude_err": float(peak_params[4]['amplitude_err']),
            "sigma": float(peak_params[4]['sigma']),
            "sigma_err": float(peak_params[4]['sigma_err'])
        },
        "R_squared": float(r_squared)
    },
    "fit_quality": {
        "r_squared": float(r_squared),
        "rmse": float(rmse)
    },
    "summary": f"Fitted spectroscopic data with 5 Gaussian peaks on a linear baseline (y = a*x + b + sum of 5 Gaussians). Peak 1 is decomposed into G1 (center~{peak_params[0]['center']:.1f}) and G2 (center~{peak_params[1]['center']:.1f}) to capture asymmetry. An intermediate shoulder G3 is centered at ~{peak_params[2]['center']:.1f}. Peak 2 is decomposed into G4 (center~{peak_params[3]['center']:.1f}) and G5 (center~{peak_params[4]['center']:.1f}). R²={r_squared:.6f}, RMSE={rmse:.6f}. All peaks use Gaussian profiles as mandated by domain rules."
}

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