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_annealing_data/out_skill/temp_spectrum_0.npy')
if data.ndim == 2:
    if data.shape[0] == 2:
        x_data, y_data = data[0], data[1]
    elif data.shape[1] == 2:
        x_data, y_data = data[:, 0], data[:, 1]
    else:
        x_data = np.linspace(0, 100, data.shape[1])
        y_data = data[0]
else:
    y_data = data
    x_data = np.linspace(0, 100, len(y_data))

# Define model: linear baseline + 4 Gaussians (2 peaks x (narrow + broad))
def gaussian(x, amplitude, center, sigma):
    return amplitude * np.exp(-(x - center)**2 / (2 * sigma**2))

def full_model(x, baseline_intercept, baseline_slope,
               A1n, c1n, sigma1n,
               A1b, c1b, sigma1b,
               A2n, c2n, sigma2n,
               A2b, c2b, sigma2b):
    baseline = baseline_intercept + baseline_slope * x
    g1n = gaussian(x, A1n, c1n, sigma1n)
    g1b = gaussian(x, A1b, c1b, sigma1b)
    g2n = gaussian(x, A2n, c2n, sigma2n)
    g2b = gaussian(x, A2b, c2b, sigma2b)
    return baseline + g1n + g1b + g2n + g2b

# Initial guesses as specified in plan
p0 = [
    0.0, 0.005,        # baseline_intercept, baseline_slope
    3.5, 35.0, 1.0,    # A1n, c1n, sigma1n
    1.0, 35.0, 4.0,    # A1b, c1b, sigma1b
    2.8, 72.0, 1.0,    # A2n, c2n, sigma2n
    0.8, 72.0, 4.0     # A2b, c2b, sigma2b
]

# Bounds as specified in plan
# Order: baseline_intercept, baseline_slope, A1n, c1n, sigma1n, A1b, c1b, sigma1b, A2n, c2n, sigma2n, A2b, c2b, sigma2b
lower_bounds = [
    -2.0, -0.05,       # baseline
    0.0, 33.0, 0.3,    # peak1 narrow: amp>0, center within ±2 of 35, sigma [0.3,2.0]
    0.0, 33.0, 2.0,    # peak1 broad: amp>0, center within ±2 of 35, sigma [2.0,10.0]
    0.0, 70.0, 0.3,    # peak2 narrow: amp>0, center within ±2 of 72, sigma [0.3,2.0]
    0.0, 70.0, 2.0     # peak2 broad: amp>0, center within ±2 of 72, sigma [2.0,10.0]
]
upper_bounds = [
    2.0, 0.05,
    10.0, 37.0, 2.0,
    10.0, 37.0, 10.0,
    10.0, 74.0, 2.0,
    10.0, 74.0, 10.0
]

try:
    popt, pcov = curve_fit(full_model, x_data, y_data, p0=p0,
                           bounds=(lower_bounds, upper_bounds),
                           method='trf', maxfev=50000)
    perr = np.sqrt(np.diag(pcov))
except Exception as e:
    # Relax bounds and retry
    lower_bounds = [
        -5.0, -0.1,
        0.0, 25.0, 0.1,
        0.0, 25.0, 1.0,
        0.0, 60.0, 0.1,
        0.0, 60.0, 1.0
    ]
    upper_bounds = [
        5.0, 0.1,
        20.0, 45.0, 5.0,
        20.0, 45.0, 20.0,
        20.0, 84.0, 5.0,
        20.0, 84.0, 20.0
    ]
    popt, pcov = curve_fit(full_model, x_data, y_data, p0=p0,
                           bounds=(lower_bounds, upper_bounds),
                           method='trf', maxfev=100000)
    perr = np.sqrt(np.diag(pcov))

# Extract parameters
param_names = ['baseline_intercept', 'baseline_slope',
               'A1n', 'c1n', 'sigma1n',
               'A1b', 'c1b', 'sigma1b',
               'A2n', 'c2n', 'sigma2n',
               'A2b', 'c2b', 'sigma2b']

params = {name: (val, err) for name, val, err in zip(param_names, popt, perr)}

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

# Compute individual components for plotting
baseline = popt[0] + popt[1] * x_data
g1n = gaussian(x_data, popt[2], popt[3], popt[4])
g1b = gaussian(x_data, popt[5], popt[6], popt[7])
g2n = gaussian(x_data, popt[8], popt[9], popt[10])
g2b = gaussian(x_data, popt[11], popt[12], popt[13])

# Visualization
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), height_ratios=[3, 1], sharex=True)
fig.suptitle('Spectroscopic Data: 4-Gaussian + Linear Baseline Fit', fontsize=14)

# Main plot
ax1.scatter(x_data, y_data, s=1, alpha=0.5, color='gray', label='Data')
ax1.plot(x_data, y_fit, 'r-', linewidth=2, label=f'Total Fit (R²={r_squared:.6f})')
ax1.plot(x_data, baseline, 'k--', linewidth=1, alpha=0.7, label='Baseline')
ax1.plot(x_data, baseline + g1n, 'b-', linewidth=1, alpha=0.7, label=f'Peak1 Narrow (c={popt[3]:.2f})')
ax1.plot(x_data, baseline + g1b, 'b--', linewidth=1, alpha=0.7, label=f'Peak1 Broad (c={popt[6]:.2f})')
ax1.plot(x_data, baseline + g2n, 'g-', linewidth=1, alpha=0.7, label=f'Peak2 Narrow (c={popt[9]:.2f})')
ax1.plot(x_data, baseline + g2b, 'g--', linewidth=1, alpha=0.7, label=f'Peak2 Broad (c={popt[12]:.2f})')
# Combined peak 1 and peak 2
ax1.fill_between(x_data, baseline, baseline + g1n + g1b, alpha=0.15, color='blue', label='Peak 1 combined')
ax1.fill_between(x_data, baseline, baseline + g2n + g2b, alpha=0.15, color='green', label='Peak 2 combined')
ax1.set_ylabel('Intensity')
ax1.legend(fontsize=8, loc='upper right')
ax1.grid(True, alpha=0.3)

# Residuals
ax2.plot(x_data, residuals, 'k-', linewidth=0.5)
ax2.axhline(y=0, color='r', linestyle='--', linewidth=0.5)
ax2.set_xlabel('X (Channel / Wavelength)')
ax2.set_ylabel('Residuals')
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": "4 Gaussians (2 peaks x narrow+broad) + linear baseline",
    "parameters": {
        "baseline": {
            "intercept": round(float(popt[0]), 6),
            "intercept_err": round(float(perr[0]), 6),
            "slope": round(float(popt[1]), 6),
            "slope_err": round(float(perr[1]), 6)
        },
        "peak_1_narrow": {
            "center": round(float(popt[3]), 4),
            "center_err": round(float(perr[3]), 4),
            "amplitude": round(float(popt[2]), 4),
            "amplitude_err": round(float(perr[2]), 4),
            "sigma": round(float(popt[4]), 4),
            "sigma_err": round(float(perr[4]), 4)
        },
        "peak_1_broad": {
            "center": round(float(popt[6]), 4),
            "center_err": round(float(perr[6]), 4),
            "amplitude": round(float(popt[5]), 4),
            "amplitude_err": round(float(perr[5]), 4),
            "sigma": round(float(popt[7]), 4),
            "sigma_err": round(float(perr[7]), 4)
        },
        "peak_2_narrow": {
            "center": round(float(popt[9]), 4),
            "center_err": round(float(perr[9]), 4),
            "amplitude": round(float(popt[8]), 4),
            "amplitude_err": round(float(perr[8]), 4),
            "sigma": round(float(popt[10]), 4),
            "sigma_err": round(float(perr[10]), 4)
        },
        "peak_2_broad": {
            "center": round(float(popt[12]), 4),
            "center_err": round(float(perr[12]), 4),
            "amplitude": round(float(popt[11]), 4),
            "amplitude_err": round(float(perr[11]), 4),
            "sigma": round(float(popt[13]), 4),
            "sigma_err": round(float(perr[13]), 4)
        }
    },
    "fit_quality": {
        "r_squared": round(float(r_squared), 8),
        "rmse": round(float(rmse), 8)
    },
    "summary": f"Fitted 2 spectroscopic peaks, each decomposed into narrow + broad Gaussian components, on a linear baseline. Peak 1 center ~{popt[3]:.2f} (narrow sigma={popt[4]:.2f}, broad sigma={popt[7]:.2f}), Peak 2 center ~{popt[9]:.2f} (narrow sigma={popt[10]:.2f}, broad sigma={popt[13]:.2f}). R²={r_squared:.6f}, RMSE={rmse:.6f}. All components are pure Gaussians per domain rules."
}

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