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_regime_data/output/temp_spectrum_4.npy')
x = data[:, 0]
y = data[:, 1]

# Define model: 3 Lorentzian peaks + linear baseline
# Lorentzian: A * gamma^2 / ((x - x0)^2 + gamma^2)
def lorentzian(x, amp, cen, wid):
    return amp * wid**2 / ((x - cen)**2 + wid**2)

def full_model(x, baseline_intercept, baseline_slope,
               peak1a_amplitude, peak1a_center, peak1a_width,
               peak1b_amplitude, peak1b_center, peak1b_width,
               peak2_amplitude, peak2_center, peak2_width):
    baseline = baseline_intercept + baseline_slope * x
    p1a = lorentzian(x, peak1a_amplitude, peak1a_center, peak1a_width)
    p1b = lorentzian(x, peak1b_amplitude, peak1b_center, peak1b_width)
    p2 = lorentzian(x, peak2_amplitude, peak2_center, peak2_width)
    return baseline + p1a + p1b + p2

# Create model
model = Model(full_model)

# Set up parameters with initial values and constraints
params = Parameters()

# Baseline
params.add('baseline_intercept', value=0.0, min=-1.0, max=1.0)
params.add('baseline_slope', value=0.005, min=-0.05, max=0.05)

# Peak 1a (first peak, ~29)
params.add('peak1a_center', value=29.0, min=24.0, max=34.0)
params.add('peak1a_amplitude', value=2.7, min=0.0)
params.add('peak1a_width', value=2.5, min=0.5, max=10.0)

# Peak 1b (split peak, ~42)
params.add('peak1b_center', value=42.0, min=37.0, max=47.0)
params.add('peak1b_amplitude', value=2.2, min=0.0)
params.add('peak1b_width', value=2.5, min=0.5, max=10.0)

# Peak 2 (~71)
params.add('peak2_center', value=71.0, min=66.0, max=76.0)
params.add('peak2_amplitude', value=3.3, min=0.0)
params.add('peak2_width', value=1.5, min=0.5, max=10.0)

# Fit using Levenberg-Marquardt
result = model.fit(y, params, x=x, method='leastsq')

# Extract fitted parameters
best = result.best_values
stderr = {}
for name in result.params:
    stderr[name] = result.params[name].stderr if result.params[name].stderr is not None else 0.0

# 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.0 - ss_res / ss_tot
rmse = np.sqrt(np.mean((y - y_fit)**2))

# Compute individual components for plotting
baseline = best['baseline_intercept'] + best['baseline_slope'] * x
peak1a = lorentzian(x, best['peak1a_amplitude'], best['peak1a_center'], best['peak1a_width'])
peak1b = lorentzian(x, best['peak1b_amplitude'], best['peak1b_center'], best['peak1b_width'])
peak2 = lorentzian(x, best['peak2_amplitude'], best['peak2_center'], best['peak2_width'])

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

# Main plot
ax1.plot(x, y, 'ko', markersize=2, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit')
ax1.plot(x, baseline + peak1a, 'b--', linewidth=1.5, alpha=0.7, label=f'Peak 1a (cen={best["peak1a_center"]:.2f})')
ax1.plot(x, baseline + peak1b, 'g--', linewidth=1.5, alpha=0.7, label=f'Peak 1b (cen={best["peak1b_center"]:.2f})')
ax1.plot(x, baseline + peak2, 'm--', linewidth=1.5, alpha=0.7, label=f'Peak 2 (cen={best["peak2_center"]:.2f})')
ax1.plot(x, baseline, 'k:', linewidth=1, alpha=0.5, label='Baseline')
ax1.set_ylabel('Intensity')
ax1.set_title('3 Lorentzian Peaks + Linear Baseline Fit')
ax1.legend(fontsize=8)
ax1.set_xlim(x.min(), x.max())

# Residuals
residuals = y - y_fit
ax2.plot(x, residuals, 'ko', markersize=2, alpha=0.5)
ax2.axhline(0, color='r', linestyle='-', linewidth=1)
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)
plt.close()

# Build results JSON
results = {
    "model_type": "3 Lorentzian peaks + linear baseline: y = a0 + a1*x + sum_i[A_i*gamma_i^2/((x-x0_i)^2+gamma_i^2)]",
    "parameters": {
        "baseline": {
            "intercept": round(best['baseline_intercept'], 6),
            "intercept_err": round(stderr['baseline_intercept'], 6),
            "slope": round(best['baseline_slope'], 6),
            "slope_err": round(stderr['baseline_slope'], 6)
        },
        "peak_1a": {
            "center": round(best['peak1a_center'], 4),
            "center_err": round(stderr['peak1a_center'], 4),
            "amplitude": round(best['peak1a_amplitude'], 4),
            "amplitude_err": round(stderr['peak1a_amplitude'], 4),
            "width": round(best['peak1a_width'], 4),
            "width_err": round(stderr['peak1a_width'], 4)
        },
        "peak_1b": {
            "center": round(best['peak1b_center'], 4),
            "center_err": round(stderr['peak1b_center'], 4),
            "amplitude": round(best['peak1b_amplitude'], 4),
            "amplitude_err": round(stderr['peak1b_amplitude'], 4),
            "width": round(best['peak1b_width'], 4),
            "width_err": round(stderr['peak1b_width'], 4)
        },
        "peak_2": {
            "center": round(best['peak2_center'], 4),
            "center_err": round(stderr['peak2_center'], 4),
            "amplitude": round(best['peak2_amplitude'], 4),
            "amplitude_err": round(stderr['peak2_amplitude'], 4),
            "width": round(best['peak2_width'], 4),
            "width_err": round(stderr['peak2_width'], 4)
        }
    },
    "fit_quality": {
        "r_squared": round(r_squared, 6),
        "rmse": round(rmse, 6)
    },
    "summary": f"Fitted 3 Lorentzian peaks on a linear baseline using Levenberg-Marquardt optimization. Peak 1a at {best['peak1a_center']:.2f}, Peak 1b at {best['peak1b_center']:.2f} (split from original peak 1), Peak 2 at {best['peak2_center']:.2f}. R^2={r_squared:.4f}, RMSE={rmse:.4f}. This corresponds to the 500 K regime where the first peak has split into two components."
}

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