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

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_regime_data/output/temp_spectrum_3.npy')
x_data = data[:, 0]
y_data = data[:, 1]

# The locked model used 2 Lorentzians + linear baseline and got R²=0.7469.
# The plan says to try a DIFFERENT functional form.
# Neighbor spectrum [4] used 3 Lorentzian peaks successfully.
# But parsimony says use simplest model achieving R²>=0.995.
# Let's try 2 Voigt profiles + linear baseline first (different functional form than Lorentzian).
# Voigt profiles have an extra shape parameter that can interpolate between Gaussian and Lorentzian.

def voigt_profile(x, center, amplitude, sigma, gamma):
    """Voigt profile using the real part of the Faddeeva function."""
    z = ((x - center) + 1j * gamma) / (sigma * np.sqrt(2))
    w = wofz(z)
    profile = np.real(w) / (sigma * np.sqrt(2 * np.pi))
    # Normalize so that amplitude is the peak height approximately
    z0 = (1j * gamma) / (sigma * np.sqrt(2))
    w0 = wofz(z0)
    peak_val = np.real(w0) / (sigma * np.sqrt(2 * np.pi))
    if peak_val > 0:
        return amplitude * profile / peak_val
    else:
        return amplitude * profile

def model_2voigt_linear(x, a0, a1, 
                         A1, x01, sigma1, gamma1,
                         A2, x02, sigma2, gamma2):
    baseline = a0 + a1 * x
    peak1 = voigt_profile(x, x01, A1, sigma1, gamma1)
    peak2 = voigt_profile(x, x02, A2, sigma2, gamma2)
    return baseline + peak1 + peak2

def model_3voigt_linear(x, a0, a1,
                         A1, x01, sigma1, gamma1,
                         A2, x02, sigma2, gamma2,
                         A3, x03, sigma3, gamma3):
    baseline = a0 + a1 * x
    peak1 = voigt_profile(x, x01, A1, sigma1, gamma1)
    peak2 = voigt_profile(x, x02, A2, sigma2, gamma2)
    peak3 = voigt_profile(x, x03, A3, sigma3, gamma3)
    return baseline + peak1 + peak2 + peak3

# --- Strategy: Try 2 Voigt peaks first (parsimony), then 3 if needed ---

# First, let's analyze the data to find peaks
from scipy.signal import find_peaks
peaks_idx, props = find_peaks(y_data, height=0.5, distance=20, prominence=0.3)
peak_positions = x_data[peaks_idx]
peak_heights = y_data[peaks_idx]
print(f"Detected peaks at x = {peak_positions}, heights = {peak_heights}")

# Try 2 Voigt + linear baseline
gmodel2 = Model(model_2voigt_linear)
params2 = Parameters()

# Baseline
params2.add('a0', value=0.0, min=-1, max=2)
params2.add('a1', value=0.005, min=-0.1, max=0.1)

# Peak 1 around x=35
params2.add('A1', value=4.0, min=0.01, max=20)
params2.add('x01', value=35, min=20, max=50)
params2.add('sigma1', value=2.0, min=0.1, max=15)
params2.add('gamma1', value=2.0, min=0.1, max=15)

# Peak 2 around x=71
params2.add('A2', value=3.3, min=0.01, max=20)
params2.add('x02', value=71, min=55, max=85)
params2.add('sigma2', value=1.5, min=0.1, max=15)
params2.add('gamma2', value=1.5, min=0.1, max=15)

result2 = gmodel2.fit(y_data, params2, x=x_data, method='leastsq', max_nfev=10000)

ss_res2 = np.sum((y_data - result2.best_fit) ** 2)
ss_tot = np.sum((y_data - np.mean(y_data)) ** 2)
r2_2voigt = 1 - ss_res2 / ss_tot
rmse_2voigt = np.sqrt(np.mean((y_data - result2.best_fit) ** 2))
print(f"2 Voigt R² = {r2_2voigt:.6f}, RMSE = {rmse_2voigt:.6f}")

# Also try with broader initial guesses and different methods
best_result = result2
best_r2 = r2_2voigt
best_model_name = '2_voigt'

# Try multiple initializations for 2 Voigt
for sigma_init in [1.0, 3.0, 5.0, 8.0]:
    for gamma_init in [0.5, 2.0, 5.0, 8.0]:
        try:
            p = Parameters()
            p.add('a0', value=0.0, min=-1, max=2)
            p.add('a1', value=0.005, min=-0.1, max=0.1)
            p.add('A1', value=4.0, min=0.01, max=20)
            p.add('x01', value=35, min=20, max=50)
            p.add('sigma1', value=sigma_init, min=0.1, max=20)
            p.add('gamma1', value=gamma_init, min=0.1, max=20)
            p.add('A2', value=3.3, min=0.01, max=20)
            p.add('x02', value=71, min=55, max=85)
            p.add('sigma2', value=sigma_init, min=0.1, max=20)
            p.add('gamma2', value=gamma_init, min=0.1, max=20)
            r = gmodel2.fit(y_data, p, x=x_data, method='leastsq', max_nfev=10000)
            ss_r = np.sum((y_data - r.best_fit) ** 2)
            r2 = 1 - ss_r / ss_tot
            if r2 > best_r2:
                best_r2 = r2
                best_result = r
                best_model_name = '2_voigt'
                print(f"  Improved 2-Voigt: R²={r2:.6f} with sigma={sigma_init}, gamma={gamma_init}")
        except:
            pass

print(f"Best 2-Voigt R² = {best_r2:.6f}")

# If 2 Voigt doesn't reach threshold, try 3 Voigt
if best_r2 < 0.995:
    print("2 Voigt insufficient, trying 3 Voigt peaks...")
    gmodel3 = Model(model_3voigt_linear)
    
    # Determine initial guesses for 3 peaks
    # Use detected peaks or reasonable guesses
    if len(peak_positions) >= 3:
        pk_centers = sorted(peak_positions[:3])
        pk_heights_sorted = [peak_heights[list(peak_positions).index(c)] for c in pk_centers]
    elif len(peak_positions) == 2:
        # Maybe there's a shoulder; split the first peak region
        pk_centers = [29, 42, 71]
        pk_heights_sorted = [2.7, 2.2, 3.3]
    else:
        pk_centers = [29, 42, 71]
        pk_heights_sorted = [2.7, 2.2, 3.3]
    
    best_r2_3 = -1
    best_result_3 = None
    
    for sigma_init in [1.0, 2.0, 4.0, 6.0]:
        for gamma_init in [0.5, 1.5, 3.0, 5.0]:
            try:
                p3 = Parameters()
                p3.add('a0', value=0.0, min=-2, max=2)
                p3.add('a1', value=0.005, min=-0.1, max=0.1)
                p3.add('A1', value=pk_heights_sorted[0], min=0.01, max=20)
                p3.add('x01', value=pk_centers[0], min=pk_centers[0]-10, max=pk_centers[0]+10)
                p3.add('sigma1', value=sigma_init, min=0.1, max=20)
                p3.add('gamma1', value=gamma_init, min=0.1, max=20)
                p3.add('A2', value=pk_heights_sorted[1], min=0.01, max=20)
                p3.add('x02', value=pk_centers[1], min=pk_centers[1]-10, max=pk_centers[1]+10)
                p3.add('sigma2', value=sigma_init, min=0.1, max=20)
                p3.add('gamma2', value=gamma_init, min=0.1, max=20)
                p3.add('A3', value=pk_heights_sorted[2], min=0.01, max=20)
                p3.add('x03', value=pk_centers[2], min=pk_centers[2]-10, max=pk_centers[2]+10)
                p3.add('sigma3', value=sigma_init, min=0.1, max=20)
                p3.add('gamma3', value=gamma_init, min=0.1, max=20)
                r3 = gmodel3.fit(y_data, p3, x=x_data, method='leastsq', max_nfev=20000)
                ss_r3 = np.sum((y_data - r3.best_fit) ** 2)
                r2_3 = 1 - ss_r3 / ss_tot
                if r2_3 > best_r2_3:
                    best_r2_3 = r2_3
                    best_result_3 = r3
                    print(f"  Improved 3-Voigt: R²={r2_3:.6f} with sigma={sigma_init}, gamma={gamma_init}")
            except:
                pass
    
    if best_r2_3 > best_r2:
        best_r2 = best_r2_3
        best_result = best_result_3
        best_model_name = '3_voigt'
        print(f"Best 3-Voigt R² = {best_r2_3:.6f}")

# Also try 2 Pseudo-Voigt (eta mixing) if still below threshold
if best_r2 < 0.995:
    print("Trying 2 Pseudo-Voigt peaks...")
    def pseudo_voigt(x, center, amplitude, width, eta):
        """Pseudo-Voigt: eta*Lorentzian + (1-eta)*Gaussian"""
        lorentz = amplitude * (width**2) / ((x - center)**2 + width**2)
        gauss = amplitude * np.exp(-np.log(2) * ((x - center) / width)**2)
        return eta * lorentz + (1 - eta) * gauss
    
    def model_2pv(x, a0, a1, A1, x01, w1, eta1, A2, x02, w2, eta2):
        return a0 + a1*x + pseudo_voigt(x, x01, A1, w1, eta1) + pseudo_voigt(x, x02, A2, w2, eta2)
    
    gpv = Model(model_2pv)
    for w_init in [1.5, 3.0, 5.0, 8.0]:
        for eta_init in [0.0, 0.3, 0.5, 0.7, 1.0]:
            try:
                ppv = Parameters()
                ppv.add('a0', value=0.0, min=-2, max=2)
                ppv.add('a1', value=0.005, min=-0.1, max=0.1)
                ppv.add('A1', value=4.0, min=0.01, max=20)
                ppv.add('x01', value=35, min=20, max=50)
                ppv.add('w1', value=w_init, min=0.1, max=30)
                ppv.add('eta1', value=eta_init, min=0, max=1)
                ppv.add('A2', value=3.3, min=0.01, max=20)
                ppv.add('x02', value=71, min=55, max=85)
                ppv.add('w2', value=w_init, min=0.1, max=30)
                ppv.add('eta2', value=eta_init, min=0, max=1)
                rpv = gpv.fit(y_data, ppv, x=x_data, method='leastsq', max_nfev=10000)
                ss_rpv = np.sum((y_data - rpv.best_fit) ** 2)
                r2_pv = 1 - ss_rpv / ss_tot
                if r2_pv > best_r2:
                    best_r2 = r2_pv
                    best_result = rpv
                    best_model_name = '2_pseudo_voigt'
                    print(f"  Improved 2-PV: R²={r2_pv:.6f} with w={w_init}, eta={eta_init}")
            except:
                pass

# Also try 2 Gaussian peaks + linear baseline
if best_r2 < 0.995:
    print("Trying 2 Gaussian peaks...")
    def model_2gauss(x, a0, a1, A1, x01, sig1, A2, x02, sig2):
        return a0 + a1*x + A1*np.exp(-0.5*((x-x01)/sig1)**2) + A2*np.exp(-0.5*((x-x02)/sig2)**2)
    
    gg = Model(model_2gauss)
    for sig_init in [1.5, 3.0, 5.0, 8.0, 12.0]:
        try:
            pg = Parameters()
            pg.add('a0', value=0.0, min=-2, max=2)
            pg.add('a1', value=0.005, min=-0.1, max=0.1)
            pg.add('A1', value=4.0, min=0.01, max=20)
            pg.add('x01', value=35, min=20, max=50)
            pg.add('sig1', value=sig_init, min=0.1, max=30)
            pg.add('A2', value=3.3, min=0.01, max=20)
            pg.add('x02', value=71, min=55, max=85)
            pg.add('sig2', value=sig_init, min=0.1, max=30)
            rg = gg.fit(y_data, pg, x=x_data, method='leastsq', max_nfev=10000)
            ss_rg = np.sum((y_data - rg.best_fit) ** 2)
            r2_g = 1 - ss_rg / ss_tot
            if r2_g > best_r2:
                best_r2 = r2_g
                best_result = rg
                best_model_name = '2_gaussian'
                print(f"  Improved 2-Gauss: R²={r2_g:.6f} with sig={sig_init}")
        except:
            pass

print(f"\nFinal best model: {best_model_name}, R² = {best_r2:.6f}")

# Extract results
final_result = best_result
final_fit = final_result.best_fit
residuals = y_data - final_fit
ss_res_final = np.sum(residuals ** 2)
r2_final = 1 - ss_res_final / ss_tot
rmse_final = np.sqrt(np.mean(residuals ** 2))

# Build parameter dict
param_dict = {}
for name, par in final_result.params.items():
    param_dict[name] = {
        'value': float(par.value),
        'stderr': float(par.stderr) if par.stderr is not None else None
    }

# Organize into structured output
results_params = {}
if best_model_name == '2_voigt':
    model_desc = '2 Voigt profiles + linear baseline'
    results_params['peak_1'] = {
        'center': float(final_result.params['x01'].value),
        'center_err': float(final_result.params['x01'].stderr) if final_result.params['x01'].stderr else None,
        'amplitude': float(final_result.params['A1'].value),
        'amplitude_err': float(final_result.params['A1'].stderr) if final_result.params['A1'].stderr else None,
        'sigma': float(final_result.params['sigma1'].value),
        'sigma_err': float(final_result.params['sigma1'].stderr) if final_result.params['sigma1'].stderr else None,
        'gamma': float(final_result.params['gamma1'].value),
        'gamma_err': float(final_result.params['gamma1'].stderr) if final_result.params['gamma1'].stderr else None,
    }
    results_params['peak_2'] = {
        'center': float(final_result.params['x02'].value),
        'center_err': float(final_result.params['x02'].stderr) if final_result.params['x02'].stderr else None,
        'amplitude': float(final_result.params['A2'].value),
        'amplitude_err': float(final_result.params['A2'].stderr) if final_result.params['A2'].stderr else None,
        'sigma': float(final_result.params['sigma2'].value),
        'sigma_err': float(final_result.params['sigma2'].stderr) if final_result.params['sigma2'].stderr else None,
        'gamma': float(final_result.params['gamma2'].value),
        'gamma_err': float(final_result.params['gamma2'].stderr) if final_result.params['gamma2'].stderr else None,
    }
    results_params['baseline'] = {
        'intercept': float(final_result.params['a0'].value),
        'intercept_err': float(final_result.params['a0'].stderr) if final_result.params['a0'].stderr else None,
        'slope': float(final_result.params['a1'].value),
        'slope_err': float(final_result.params['a1'].stderr) if final_result.params['a1'].stderr else None,
    }
elif best_model_name == '3_voigt':
    model_desc = '3 Voigt profiles + linear baseline'
    for i in range(1, 4):
        results_params[f'peak_{i}'] = {
            'center': float(final_result.params[f'x0{i}'].value),
            'center_err': float(final_result.params[f'x0{i}'].stderr) if final_result.params[f'x0{i}'].stderr else None,
            'amplitude': float(final_result.params[f'A{i}'].value),
            'amplitude_err': float(final_result.params[f'A{i}'].stderr) if final_result.params[f'A{i}'].stderr else None,
            'sigma': float(final_result.params[f'sigma{i}'].value),
            'sigma_err': float(final_result.params[f'sigma{i}'].stderr) if final_result.params[f'sigma{i}'].stderr else None,
            'gamma': float(final_result.params[f'gamma{i}'].value),
            'gamma_err': float(final_result.params[f'gamma{i}'].stderr) if final_result.params[f'gamma{i}'].stderr else None,
        }
    results_params['baseline'] = {
        'intercept': float(final_result.params['a0'].value),
        'intercept_err': float(final_result.params['a0'].stderr) if final_result.params['a0'].stderr else None,
        'slope': float(final_result.params['a1'].value),
        'slope_err': float(final_result.params['a1'].stderr) if final_result.params['a1'].stderr else None,
    }
elif best_model_name == '2_pseudo_voigt':
    model_desc = '2 Pseudo-Voigt profiles + linear baseline'
    results_params['peak_1'] = {
        'center': float(final_result.params['x01'].value),
        'center_err': float(final_result.params['x01'].stderr) if final_result.params['x01'].stderr else None,
        'amplitude': float(final_result.params['A1'].value),
        'amplitude_err': float(final_result.params['A1'].stderr) if final_result.params['A1'].stderr else None,
        'width': float(final_result.params['w1'].value),
        'width_err': float(final_result.params['w1'].stderr) if final_result.params['w1'].stderr else None,
        'eta': float(final_result.params['eta1'].value),
        'eta_err': float(final_result.params['eta1'].stderr) if final_result.params['eta1'].stderr else None,
    }
    results_params['peak_2'] = {
        'center': float(final_result.params['x02'].value),
        'center_err': float(final_result.params['x02'].stderr) if final_result.params['x02'].stderr else None,
        'amplitude': float(final_result.params['A2'].value),
        'amplitude_err': float(final_result.params['A2'].stderr) if final_result.params['A2'].stderr else None,
        'width': float(final_result.params['w2'].value),
        'width_err': float(final_result.params['w2'].stderr) if final_result.params['w2'].stderr else None,
        'eta': float(final_result.params['eta2'].value),
        'eta_err': float(final_result.params['eta2'].stderr) if final_result.params['eta2'].stderr else None,
    }
    results_params['baseline'] = {
        'intercept': float(final_result.params['a0'].value),
        'intercept_err': float(final_result.params['a0'].stderr) if final_result.params['a0'].stderr else None,
        'slope': float(final_result.params['a1'].value),
        'slope_err': float(final_result.params['a1'].stderr) if final_result.params['a1'].stderr else None,
    }
elif best_model_name == '2_gaussian':
    model_desc = '2 Gaussian profiles + linear baseline'
    results_params['peak_1'] = {
        'center': float(final_result.params['x01'].value),
        'center_err': float(final_result.params['x01'].stderr) if final_result.params['x01'].stderr else None,
        'amplitude': float(final_result.params['A1'].value),
        'amplitude_err': float(final_result.params['A1'].stderr) if final_result.params['A1'].stderr else None,
        'sigma': float(final_result.params['sig1'].value),
        'sigma_err': float(final_result.params['sig1'].stderr) if final_result.params['sig1'].stderr else None,
    }
    results_params['peak_2'] = {
        'center': float(final_result.params['x02'].value),
        'center_err': float(final_result.params['x02'].stderr) if final_result.params['x02'].stderr else None,
        'amplitude': float(final_result.params['A2'].value),
        'amplitude_err': float(final_result.params['A2'].stderr) if final_result.params['A2'].stderr else None,
        'sigma': float(final_result.params['sig2'].value),
        'sigma_err': float(final_result.params['sig2'].stderr) if final_result.params['sig2'].stderr else None,
    }
    results_params['baseline'] = {
        'intercept': float(final_result.params['a0'].value),
        'intercept_err': float(final_result.params['a0'].stderr) if final_result.params['a0'].stderr else None,
        'slope': float(final_result.params['a1'].value),
        'slope_err': float(final_result.params['a1'].stderr) if final_result.params['a1'].stderr else None,
    }
else:
    model_desc = best_model_name
    results_params['all_params'] = param_dict

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

ax1.plot(x_data, y_data, 'b.', markersize=2, alpha=0.5, label='Data')
ax1.plot(x_data, final_fit, 'r-', linewidth=2, label=f'Fit ({model_desc})')

# Plot components
params_vals = final_result.params
baseline_vals = params_vals['a0'].value + params_vals['a1'].value * x_data
ax1.plot(x_data, baseline_vals, 'k--', linewidth=1, alpha=0.5, label='Baseline')

if best_model_name == '2_voigt':
    comp1 = baseline_vals + voigt_profile(x_data, params_vals['x01'].value, params_vals['A1'].value, params_vals['sigma1'].value, params_vals['gamma1'].value)
    comp2 = baseline_vals + voigt_profile(x_data, params_vals['x02'].value, params_vals['A2'].value, params_vals['sigma2'].value, params_vals['gamma2'].value)
    ax1.plot(x_data, comp1, 'g--', linewidth=1, alpha=0.7, label='Peak 1')
    ax1.plot(x_data, comp2, 'm--', linewidth=1, alpha=0.7, label='Peak 2')
elif best_model_name == '3_voigt':
    for i in range(1, 4):
        comp = baseline_vals + voigt_profile(x_data, params_vals[f'x0{i}'].value, params_vals[f'A{i}'].value, params_vals[f'sigma{i}'].value, params_vals[f'gamma{i}'].value)
        ax1.plot(x_data, comp, ['g--', 'm--', 'c--'][i-1], linewidth=1, alpha=0.7, label=f'Peak {i}')
elif best_model_name == '2_pseudo_voigt':
    def pseudo_voigt_eval(x, center, amplitude, width, eta):
        lorentz = amplitude * (width**2) / ((x - center)**2 + width**2)
        gauss = amplitude * np.exp(-np.log(2) * ((x - center) / width)**2)
        return eta * lorentz + (1 - eta) * gauss
    comp1 = baseline_vals + pseudo_voigt_eval(x_data, params_vals['x01'].value, params_vals['A1'].value, params_vals['w1'].value, params_vals['eta1'].value)
    comp2 = baseline_vals + pseudo_voigt_eval(x_data, params_vals['x02'].value, params_vals['A2'].value, params_vals['w2'].value, params_vals['eta2'].value)
    ax1.plot(x_data, comp1, 'g--', linewidth=1, alpha=0.7, label='Peak 1')
    ax1.plot(x_data, comp2, 'm--', linewidth=1, alpha=0.7, label='Peak 2')
elif best_model_name == '2_gaussian':
    comp1 = baseline_vals + params_vals['A1'].value * np.exp(-0.5*((x_data - params_vals['x01'].value)/params_vals['sig1'].value)**2)
    comp2 = baseline_vals + params_vals['A2'].value * np.exp(-0.5*((x_data - params_vals['x02'].value)/params_vals['sig2'].value)**2)
    ax1.plot(x_data, comp1, 'g--', linewidth=1, alpha=0.7, label='Peak 1')
    ax1.plot(x_data, comp2, 'm--', linewidth=1, alpha=0.7, label='Peak 2')

ax1.set_xlabel('X')
ax1.set_ylabel('Intensity')
ax1.set_title(f'Spectroscopic Fit: {model_desc} (R²={r2_final:.6f})')
ax1.legend()

ax2.plot(x_data, residuals, 'b.', markersize=2, alpha=0.5)
ax2.axhline(y=0, color='r', linestyle='--')
ax2.set_xlabel('X')
ax2.set_ylabel('Residuals')
ax2.set_title('Residuals')

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

# Summary
if r2_final >= 0.995:
    summary_msg = f"Successfully fitted spectrum 4/5 (T=450) with {model_desc}. R²={r2_final:.6f} exceeds threshold 0.995. The Voigt/alternative profile provides better shape flexibility than pure Lorentzian."
else:
    summary_msg = f"Best fit achieved with {model_desc}. R²={r2_final:.6f} {'meets' if r2_final >= 0.995 else 'below'} threshold 0.995. Locked Lorentzian model gave R²=0.7469."

results = {
    "model_type": model_desc,
    "parameters": results_params,
    "fit_quality": {
        "r_squared": round(float(r2_final), 6),
        "rmse": round(float(rmse_final), 6)
    },
    "summary": summary_msg
}

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