import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.special import voigt_profile, wofz
from scipy.optimize import differential_evolution, least_squares
from lmfit import Parameters, minimize, Model
import warnings
warnings.filterwarnings('ignore')

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

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

# Define trapz compatibility
if hasattr(np, 'trapezoid'):
    _trapz = np.trapezoid
else:
    _trapz = np.trapz

# Define model components
def logistic_bg(x, L, k, x0):
    """Sigmoidal/logistic background"""
    return L / (1.0 + np.exp(-k * (x - x0)))

def linear_baseline(x, intercept, slope):
    return intercept + slope * x

def voigt_peak(x, amplitude, center, sigma, gamma):
    """Voigt profile using scipy's voigt_profile"""
    # voigt_profile expects sigma and gamma as half-widths
    v = voigt_profile(x - center, sigma, gamma)
    # Normalize: voigt_profile at peak is voigt_profile(0, sigma, gamma)
    v_max = voigt_profile(0, sigma, gamma)
    if v_max > 0:
        return amplitude * v / v_max
    return np.zeros_like(x)

def skewed_voigt(x, amplitude, center, sigma, gamma, skewness):
    """Skewed Voigt: Voigt multiplied by a skew factor"""
    v = voigt_peak(x, amplitude, center, sigma, gamma)
    # Apply skewness using an error-function-based asymmetry
    skew_factor = 1.0 + scipy_erf(skewness * (x - center) / (sigma * np.sqrt(2)))
    return v * skew_factor

def gaussian_peak(x, amplitude, center, sigma):
    """Simple Gaussian"""
    return amplitude * np.exp(-0.5 * ((x - center) / sigma) ** 2)

from scipy.special import erf as scipy_erf

def full_model(params, x):
    """Complete model: linear baseline + logistic + 3 peaks (2 skewed Voigt + 1 narrow Gaussian + 1 Voigt)"""
    # Linear baseline
    bl_int = params['baseline_intercept']
    bl_slope = params['baseline_slope']
    bg_lin = linear_baseline(x, bl_int, bl_slope)
    
    # Logistic background
    log_L = params['logistic_amplitude']
    log_k = params['logistic_rate']
    log_x0 = params['logistic_center']
    bg_log = logistic_bg(x, log_L, log_k, log_x0)
    
    # Peak 1 broad: Skewed Voigt
    p1b_amp = params['peak1_broad_amplitude']
    p1b_cen = params['peak1_broad_center']
    p1b_sig = params['peak1_broad_sigma']
    p1b_gam = params['peak1_broad_gamma']
    p1b_skew = params['peak1_broad_skewness']
    pk1_broad = skewed_voigt(x, p1b_amp, p1b_cen, p1b_sig, p1b_gam, p1b_skew)
    
    # Peak 1 narrow: Gaussian spike
    p1n_amp = params['peak1_narrow_amplitude']
    p1n_cen = params['peak1_narrow_center']
    p1n_sig = params['peak1_narrow_sigma']
    pk1_narrow = gaussian_peak(x, p1n_amp, p1n_cen, p1n_sig)
    
    # Peak 2: Skewed Voigt
    p2_amp = params['peak2_amplitude']
    p2_cen = params['peak2_center']
    p2_sig = params['peak2_sigma']
    p2_gam = params['peak2_gamma']
    p2_skew = params['peak2_skewness']
    pk2 = skewed_voigt(x, p2_amp, p2_cen, p2_sig, p2_gam, p2_skew)
    
    # Peak 3: Voigt
    p3_amp = params['peak3_amplitude']
    p3_cen = params['peak3_center']
    p3_sig = params['peak3_sigma']
    p3_gam = params['peak3_gamma']
    pk3 = voigt_peak(x, p3_amp, p3_cen, p3_sig, p3_gam)
    
    return bg_lin + bg_log + pk1_broad + pk1_narrow + pk2 + pk3

def residual(params, x, y):
    return y - full_model(params, x)

# Set up lmfit Parameters with bounds
params = Parameters()

# Linear baseline
params.add('baseline_intercept', value=-0.1, min=-1.0, max=1.0)
params.add('baseline_slope', value=0.005, min=-0.05, max=0.05)

# Logistic background
params.add('logistic_amplitude', value=1.0, min=0.3, max=2.5)
params.add('logistic_rate', value=0.2, min=0.05, max=1.0)
params.add('logistic_center', value=10.0, min=3.0, max=22.0)

# Peak 1 broad (Skewed Voigt) - main feature around x~17
params.add('peak1_broad_amplitude', value=3.0, min=0.5, max=8.0)
params.add('peak1_broad_center', value=17.0, min=14.0, max=20.0)
params.add('peak1_broad_sigma', value=4.0, min=1.0, max=10.0)
params.add('peak1_broad_gamma', value=1.0, min=0.01, max=8.0)
params.add('peak1_broad_skewness', value=-0.5, min=-5.0, max=5.0)

# Peak 1 narrow (Gaussian spike) at x~20
params.add('peak1_narrow_amplitude', value=1.0, min=0.01, max=5.0)
params.add('peak1_narrow_center', value=20.0, min=19.0, max=21.0)
params.add('peak1_narrow_sigma', value=0.5, min=0.2, max=1.5)

# Peak 2 (Skewed Voigt) at x~35
params.add('peak2_amplitude', value=1.5, min=0.1, max=5.0)
params.add('peak2_center', value=35.0, min=32.0, max=38.0)
params.add('peak2_sigma', value=1.0, min=0.2, max=5.0)
params.add('peak2_gamma', value=0.5, min=0.01, max=5.0)
params.add('peak2_skewness', value=-0.5, min=-5.0, max=5.0)

# Peak 3 (Voigt) at x~47
params.add('peak3_amplitude', value=0.5, min=0.01, max=3.0)
params.add('peak3_center', value=47.0, min=44.0, max=50.0)
params.add('peak3_sigma', value=3.0, min=0.5, max=8.0)
params.add('peak3_gamma', value=1.0, min=0.01, max=8.0)

# Step 1: Try a preliminary fit to get good starting parameters
# First, fit just the background to the data in regions away from peaks
mask_bg = ((x < 5) | ((x > 25) & (x < 30)) | ((x > 38) & (x < 42)))
if np.sum(mask_bg) > 10:
    params_bg = Parameters()
    params_bg.add('baseline_intercept', value=-0.1, min=-1.0, max=1.0)
    params_bg.add('baseline_slope', value=0.005, min=-0.05, max=0.05)
    params_bg.add('logistic_amplitude', value=1.0, min=0.3, max=2.5)
    params_bg.add('logistic_rate', value=0.2, min=0.05, max=1.0)
    params_bg.add('logistic_center', value=10.0, min=3.0, max=22.0)
    
    def bg_residual(p, xx, yy):
        bg = linear_baseline(xx, p['baseline_intercept'], p['baseline_slope']) + \
             logistic_bg(xx, p['logistic_amplitude'], p['logistic_rate'], p['logistic_center'])
        return yy - bg
    
    try:
        bg_result = minimize(bg_residual, params_bg, args=(x[mask_bg], y[mask_bg]), method='leastsq', max_nfev=5000)
        for pname in ['baseline_intercept', 'baseline_slope', 'logistic_amplitude', 'logistic_rate', 'logistic_center']:
            params[pname].value = bg_result.params[pname].value
    except:
        pass

# Step 2: Full fit using multiple methods for robustness
# First try with Nelder-Mead (simplex) to get into the right basin
try:
    result_nm = minimize(residual, params, args=(x, y), method='nelder', max_nfev=50000)
    params_after_nm = result_nm.params
except:
    params_after_nm = params

# Then refine with Levenberg-Marquardt
try:
    result = minimize(residual, params_after_nm, args=(x, y), method='leastsq', max_nfev=50000)
except:
    result = minimize(residual, params, args=(x, y), method='leastsq', max_nfev=50000)

# Try another round with different initial conditions if fit is poor
y_fit = full_model(result.params, x)
ss_res = np.sum((y - y_fit) ** 2)
ss_tot = np.sum((y - np.mean(y)) ** 2)
r2 = 1 - ss_res / ss_tot

if r2 < 0.95:
    # Try with differential evolution via scipy
    from scipy.optimize import differential_evolution as de
    
    bounds_de = [
        (-1.0, 1.0),   # baseline_intercept
        (-0.05, 0.05), # baseline_slope
        (0.3, 2.5),    # logistic_amplitude
        (0.05, 1.0),   # logistic_rate
        (3.0, 22.0),   # logistic_center
        (0.5, 8.0),    # peak1_broad_amplitude
        (14.0, 20.0),  # peak1_broad_center
        (1.0, 10.0),   # peak1_broad_sigma
        (0.01, 8.0),   # peak1_broad_gamma
        (-5.0, 5.0),   # peak1_broad_skewness
        (0.01, 5.0),   # peak1_narrow_amplitude
        (19.0, 21.0),  # peak1_narrow_center
        (0.2, 1.5),    # peak1_narrow_sigma
        (0.1, 5.0),    # peak2_amplitude
        (32.0, 38.0),  # peak2_center
        (0.2, 5.0),    # peak2_sigma
        (0.01, 5.0),   # peak2_gamma
        (-5.0, 5.0),   # peak2_skewness
        (0.01, 3.0),   # peak3_amplitude
        (44.0, 50.0),  # peak3_center
        (0.5, 8.0),    # peak3_sigma
        (0.01, 8.0),   # peak3_gamma
    ]
    
    pnames_order = ['baseline_intercept', 'baseline_slope', 'logistic_amplitude', 'logistic_rate', 'logistic_center',
                    'peak1_broad_amplitude', 'peak1_broad_center', 'peak1_broad_sigma', 'peak1_broad_gamma', 'peak1_broad_skewness',
                    'peak1_narrow_amplitude', 'peak1_narrow_center', 'peak1_narrow_sigma',
                    'peak2_amplitude', 'peak2_center', 'peak2_sigma', 'peak2_gamma', 'peak2_skewness',
                    'peak3_amplitude', 'peak3_center', 'peak3_sigma', 'peak3_gamma']
    
    def de_objective(pvec):
        for i, pn in enumerate(pnames_order):
            params[pn].value = pvec[i]
        res = residual(params, x, y)
        return np.sum(res**2)
    
    try:
        de_result = de(de_objective, bounds_de, seed=42, maxiter=500, tol=1e-8, 
                       popsize=20, mutation=(0.5, 1.5), recombination=0.9)
        for i, pn in enumerate(pnames_order):
            params[pn].value = de_result.x[i]
        result = minimize(residual, params, args=(x, y), method='leastsq', max_nfev=50000)
    except Exception as e:
        print(f"DE optimization failed: {e}")

# Extract final parameters
fp = result.params

# Compute fit and components
y_fit = full_model(fp, x)
residuals = y - y_fit

# Compute individual components for visualization
bg_linear = linear_baseline(x, fp['baseline_intercept'].value, fp['baseline_slope'].value)
bg_logistic = logistic_bg(x, fp['logistic_amplitude'].value, fp['logistic_rate'].value, fp['logistic_center'].value)
total_bg = bg_linear + bg_logistic

pk1_broad = skewed_voigt(x, fp['peak1_broad_amplitude'].value, fp['peak1_broad_center'].value, 
                         fp['peak1_broad_sigma'].value, fp['peak1_broad_gamma'].value, fp['peak1_broad_skewness'].value)
pk1_narrow = gaussian_peak(x, fp['peak1_narrow_amplitude'].value, fp['peak1_narrow_center'].value, fp['peak1_narrow_sigma'].value)
pk2 = skewed_voigt(x, fp['peak2_amplitude'].value, fp['peak2_center'].value,
                   fp['peak2_sigma'].value, fp['peak2_gamma'].value, fp['peak2_skewness'].value)
pk3 = voigt_peak(x, fp['peak3_amplitude'].value, fp['peak3_center'].value,
                 fp['peak3_sigma'].value, fp['peak3_gamma'].value)

# Compute fit quality
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))

# Compute FWHM and areas for each peak
def compute_fwhm_numerical(x_arr, y_arr):
    """Compute FWHM numerically from a peak profile"""
    if np.max(y_arr) <= 0:
        return 0.0
    half_max = np.max(y_arr) / 2.0
    above = np.where(y_arr >= half_max)[0]
    if len(above) < 2:
        return 0.0
    return x_arr[above[-1]] - x_arr[above[0]]

def compute_area(x_arr, y_arr):
    """Compute area by numerical integration"""
    return _trapz(y_arr, x_arr)

# Dense x for FWHM calculation
x_dense = np.linspace(x.min(), x.max(), 10000)

pk1b_dense = skewed_voigt(x_dense, fp['peak1_broad_amplitude'].value, fp['peak1_broad_center'].value,
                          fp['peak1_broad_sigma'].value, fp['peak1_broad_gamma'].value, fp['peak1_broad_skewness'].value)
pk1n_dense = gaussian_peak(x_dense, fp['peak1_narrow_amplitude'].value, fp['peak1_narrow_center'].value, fp['peak1_narrow_sigma'].value)
pk2_dense = skewed_voigt(x_dense, fp['peak2_amplitude'].value, fp['peak2_center'].value,
                         fp['peak2_sigma'].value, fp['peak2_gamma'].value, fp['peak2_skewness'].value)
pk3_dense = voigt_peak(x_dense, fp['peak3_amplitude'].value, fp['peak3_center'].value,
                       fp['peak3_sigma'].value, fp['peak3_gamma'].value)

fwhm_pk1_broad = compute_fwhm_numerical(x_dense, pk1b_dense)
fwhm_pk1_narrow = 2.355 * fp['peak1_narrow_sigma'].value  # Gaussian FWHM
fwhm_pk2 = compute_fwhm_numerical(x_dense, pk2_dense)
fwhm_pk3 = compute_fwhm_numerical(x_dense, pk3_dense)

area_pk1_broad = compute_area(x_dense, pk1b_dense)
area_pk1_narrow = compute_area(x_dense, pk1n_dense)
area_pk2 = compute_area(x_dense, pk2_dense)
area_pk3 = compute_area(x_dense, pk3_dense)

# Get uncertainties
def get_err(pname):
    try:
        err = fp[pname].stderr
        if err is None:
            return 0.0
        return float(err)
    except:
        return 0.0

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

ax1.plot(x, y, 'k.', markersize=2, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Total Fit')
ax1.plot(x, total_bg, 'b--', linewidth=1.5, alpha=0.7, label='Background (Linear + Logistic)')
ax1.fill_between(x, total_bg, total_bg + pk1_broad, alpha=0.3, color='green', label=f'Peak 1 Broad (SkewedVoigt, c={fp["peak1_broad_center"].value:.1f})')
ax1.fill_between(x, total_bg, total_bg + pk1_narrow, alpha=0.3, color='cyan', label=f'Peak 1 Narrow (Gaussian, c={fp["peak1_narrow_center"].value:.1f})')
ax1.fill_between(x, total_bg, total_bg + pk2, alpha=0.3, color='orange', label=f'Peak 2 (SkewedVoigt, c={fp["peak2_center"].value:.1f})')
ax1.fill_between(x, total_bg, total_bg + pk3, alpha=0.3, color='purple', label=f'Peak 3 (Voigt, c={fp["peak3_center"].value:.1f})')
ax1.set_ylabel('Intensity')
ax1.set_title(f'Spectroscopic Data Curve Fit (R\u00b2={r_squared:.5f}, RMSE={rmse:.4f})')
ax1.legend(loc='upper right', fontsize=8)
ax1.grid(True, alpha=0.3)

ax2.plot(x, residuals, 'k.', markersize=2, alpha=0.5)
ax2.axhline(y=0, color='r', linestyle='--', linewidth=1)
ax2.set_xlabel('x')
ax2.set_ylabel('Residuals')
ax2.set_title('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": "Linear baseline + Logistic background + SkewedVoigt (peak1 broad) + Gaussian (peak1 narrow) + SkewedVoigt (peak2) + Voigt (peak3)",
    "parameters": {
        "baseline": {
            "intercept": round(float(fp['baseline_intercept'].value), 6),
            "intercept_err": round(get_err('baseline_intercept'), 6),
            "slope": round(float(fp['baseline_slope'].value), 6),
            "slope_err": round(get_err('baseline_slope'), 6)
        },
        "logistic_background": {
            "amplitude": round(float(fp['logistic_amplitude'].value), 6),
            "amplitude_err": round(get_err('logistic_amplitude'), 6),
            "rate": round(float(fp['logistic_rate'].value), 6),
            "rate_err": round(get_err('logistic_rate'), 6),
            "center": round(float(fp['logistic_center'].value), 6),
            "center_err": round(get_err('logistic_center'), 6)
        },
        "peak_1_broad": {
            "type": "SkewedVoigt",
            "center": round(float(fp['peak1_broad_center'].value), 4),
            "center_err": round(get_err('peak1_broad_center'), 4),
            "amplitude": round(float(fp['peak1_broad_amplitude'].value), 4),
            "amplitude_err": round(get_err('peak1_broad_amplitude'), 4),
            "sigma": round(float(fp['peak1_broad_sigma'].value), 4),
            "sigma_err": round(get_err('peak1_broad_sigma'), 4),
            "gamma": round(float(fp['peak1_broad_gamma'].value), 4),
            "gamma_err": round(get_err('peak1_broad_gamma'), 4),
            "skewness": round(float(fp['peak1_broad_skewness'].value), 4),
            "skewness_err": round(get_err('peak1_broad_skewness'), 4),
            "fwhm": round(float(fwhm_pk1_broad), 4),
            "area": round(float(area_pk1_broad), 4)
        },
        "peak_1_narrow": {
            "type": "Gaussian",
            "center": round(float(fp['peak1_narrow_center'].value), 4),
            "center_err": round(get_err('peak1_narrow_center'), 4),
            "amplitude": round(float(fp['peak1_narrow_amplitude'].value), 4),
            "amplitude_err": round(get_err('peak1_narrow_amplitude'), 4),
            "sigma": round(float(fp['peak1_narrow_sigma'].value), 4),
            "sigma_err": round(get_err('peak1_narrow_sigma'), 4),
            "fwhm": round(float(fwhm_pk1_narrow), 4),
            "area": round(float(area_pk1_narrow), 4)
        },
        "peak_2": {
            "type": "SkewedVoigt",
            "center": round(float(fp['peak2_center'].value), 4),
            "center_err": round(get_err('peak2_center'), 4),
            "amplitude": round(float(fp['peak2_amplitude'].value), 4),
            "amplitude_err": round(get_err('peak2_amplitude'), 4),
            "sigma": round(float(fp['peak2_sigma'].value), 4),
            "sigma_err": round(get_err('peak2_sigma'), 4),
            "gamma": round(float(fp['peak2_gamma'].value), 4),
            "gamma_err": round(get_err('peak2_gamma'), 4),
            "skewness": round(float(fp['peak2_skewness'].value), 4),
            "skewness_err": round(get_err('peak2_skewness'), 4),
            "fwhm": round(float(fwhm_pk2), 4),
            "area": round(float(area_pk2), 4)
        },
        "peak_3": {
            "type": "Voigt",
            "center": round(float(fp['peak3_center'].value), 4),
            "center_err": round(get_err('peak3_center'), 4),
            "amplitude": round(float(fp['peak3_amplitude'].value), 4),
            "amplitude_err": round(get_err('peak3_amplitude'), 4),
            "sigma": round(float(fp['peak3_sigma'].value), 4),
            "sigma_err": round(get_err('peak3_sigma'), 4),
            "gamma": round(float(fp['peak3_gamma'].value), 4),
            "gamma_err": round(get_err('peak3_gamma'), 4),
            "fwhm": round(float(fwhm_pk3), 4),
            "area": round(float(area_pk3), 4)
        }
    },
    "fit_quality": {
        "r_squared": round(float(r_squared), 6),
        "rmse": round(float(rmse), 6)
    },
    "summary": f"Spectroscopic data fitted with linear baseline + logistic sigmoid background + 4 peak components (2 skewed Voigt, 1 narrow Gaussian, 1 symmetric Voigt). Main broad feature (SkewedVoigt) centered at x={fp['peak1_broad_center'].value:.2f} with sharp Gaussian spike at x={fp['peak1_narrow_center'].value:.2f}. Secondary skewed Voigt peak at x={fp['peak2_center'].value:.2f}. Broad Voigt bump at x={fp['peak3_center'].value:.2f}. R\u00b2={r_squared:.5f}, RMSE={rmse:.4f}. Total 22 free parameters fitted to {len(x)} data points."
}

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