import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from lmfit.models import LorentzianModel
from scipy import sparse
from scipy.sparse.linalg import spsolve

# ---- Step 1-2: Load and sort by ascending Raman shift ----
DATA_PATH = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260310_190310/results/analysis_raman_silicon_CurveFit_20260310_190743_001/temp_spectrum_0.npy'
data = np.load(DATA_PATH)
x_raw, y_raw = data[:, 0], data[:, 1]
sort_idx = np.argsort(x_raw)
x = x_raw[sort_idx]
y = y_raw[sort_idx]

# ---- Step 3: ALS Baseline Correction ----
def als_baseline(y, lam=1e6, p=0.001, niter=50):
    """Asymmetric Least Squares baseline estimation."""
    L = len(y)
    D = sparse.diags([1, -2, 1], [0, -1, -2], shape=(L, L - 2))
    D = lam * D.dot(D.transpose())
    w = np.ones(L)
    for i in range(niter):
        W = sparse.spdiags(w, 0, L, L)
        Z = W + D
        z = spsolve(Z, w * y)
        w_new = p * (y > z) + (1 - p) * (y <= z)
        if np.allclose(w, w_new):
            break
        w = w_new
    return z

baseline = als_baseline(y, lam=1e6, p=0.001, niter=50)
y_corrected = y - baseline

# Check negative points
neg_frac = np.sum(y_corrected < 0) / len(y_corrected)
print(f"Fraction of negative corrected points: {neg_frac:.4f}")

# ---- Step 4: Build 3-peak Lorentzian model ----
# Peak 1: First-order Si optical phonon at ~520 cm-1
# Plan: center=520, FWHM=3.5 -> sigma (HWHM) = 1.75, amplitude guess
# lmfit Lorentzian: amplitude/(pi*sigma) at peak, so amplitude = peak_height * pi * sigma
# Plan amplitude=800 is peak intensity. lmfit amplitude = 800 * pi * 1.75 ~ 4398
p1 = LorentzianModel(prefix='p1_')
p1_params = p1.make_params()
p1_params['p1_center'].set(value=520.0, min=515, max=525)
p1_params['p1_sigma'].set(value=1.75, min=1.25, max=4.0)  # HWHM; FWHM=2*sigma, constrain FWHM to [2.5, 8.0]
p1_params['p1_amplitude'].set(value=800.0 * np.pi * 1.75, min=0)

# Peak 2: 2TA overtone at ~300 cm-1
# Plan: center=302, FWHM=40 -> sigma=20, amplitude=200 (peak height)
p2 = LorentzianModel(prefix='p2_')
p2_params = p2.make_params()
p2_params['p2_center'].set(value=302.0, min=280, max=320)
p2_params['p2_sigma'].set(value=20.0, min=10.0, max=40.0)  # HWHM; FWHM=2*sigma, constrain FWHM to [20, 80]
p2_params['p2_amplitude'].set(value=200.0 * np.pi * 20.0, min=0)

# Peak 3: 2TO overtone at ~960 cm-1
# Plan: center=950, FWHM=40 -> sigma=20, amplitude=160 (peak height)
p3 = LorentzianModel(prefix='p3_')
p3_params = p3.make_params()
p3_params['p3_center'].set(value=950.0, min=930, max=970)
p3_params['p3_sigma'].set(value=20.0, min=10.0, max=40.0)  # HWHM; FWHM=2*sigma, constrain FWHM to [20, 80]
p3_params['p3_amplitude'].set(value=160.0 * np.pi * 20.0, min=0)

# Composite model (no explicit background model since ALS already removed baseline)
composite = p1 + p2 + p3
all_params = p1_params
all_params.update(p2_params)
all_params.update(p3_params)

# Fit to baseline-corrected data using Levenberg-Marquardt
result = composite.fit(y_corrected, all_params, x=x, method='leastsq')

# ---- Step 5: Compute R-squared on full spectrum (baseline + peaks vs original data) ----
# As per mandatory rule: R-squared on baseline + peaks vs original data
y_model_full = baseline + result.best_fit
ss_res_full = np.sum((y - y_model_full) ** 2)
ss_tot_full = np.sum((y - np.mean(y)) ** 2)
r_squared = 1.0 - ss_res_full / ss_tot_full

# RMSE
rmse = np.sqrt(np.mean((y - y_model_full) ** 2))

# Also compute R-squared on corrected data for reference
ss_res_corr = np.sum((y_corrected - result.best_fit) ** 2)
ss_tot_corr = np.sum((y_corrected - np.mean(y_corrected)) ** 2)
r_squared_corrected = 1.0 - ss_res_corr / ss_tot_corr

print(f"R-squared (full spectrum): {r_squared:.6f}")
print(f"R-squared (corrected): {r_squared_corrected:.6f}")
print(f"RMSE: {rmse:.4f}")

# Extract parameters
def get_peak_params(result, prefix, x):
    center = result.params[f'{prefix}center'].value
    center_err = result.params[f'{prefix}center'].stderr if result.params[f'{prefix}center'].stderr else 0.0
    sigma = result.params[f'{prefix}sigma'].value
    sigma_err = result.params[f'{prefix}sigma'].stderr if result.params[f'{prefix}sigma'].stderr else 0.0
    amplitude = result.params[f'{prefix}amplitude'].value
    amplitude_err = result.params[f'{prefix}amplitude'].stderr if result.params[f'{prefix}amplitude'].stderr else 0.0
    fwhm = 2.0 * sigma
    fwhm_err = 2.0 * sigma_err
    # Peak intensity (height) = amplitude / (pi * sigma)
    peak_intensity = amplitude / (np.pi * sigma) if sigma > 0 else 0.0
    # Integrated area = amplitude (for lmfit Lorentzian, amplitude IS the integrated area)
    integrated_area = amplitude
    return {
        'center': round(center, 4),
        'center_err': round(center_err, 4),
        'sigma': round(sigma, 4),
        'sigma_err': round(sigma_err, 4),
        'fwhm': round(fwhm, 4),
        'fwhm_err': round(fwhm_err, 4),
        'amplitude': round(amplitude, 4),
        'amplitude_err': round(amplitude_err, 4),
        'peak_intensity': round(peak_intensity, 4),
        'integrated_area': round(integrated_area, 4)
    }

p1_results = get_peak_params(result, 'p1_', x)
p2_results = get_peak_params(result, 'p2_', x)
p3_results = get_peak_params(result, 'p3_', x)

# Step 3 validation: Check residuals near 520 cm-1 for systematic asymmetry
mask_520 = (x > 515) & (x < 525)
residuals_520 = y_corrected[mask_520] - result.best_fit[mask_520]
# Check if left side residuals are systematically different from right side
center_fit = p1_results['center']
left_mask = (x > 515) & (x < center_fit)
right_mask = (x > center_fit) & (x < 525)
resid_left_mean = np.mean(y_corrected[left_mask] - result.best_fit[left_mask]) if np.sum(left_mask) > 0 else 0
resid_right_mean = np.mean(y_corrected[right_mask] - result.best_fit[right_mask]) if np.sum(right_mask) > 0 else 0
asymmetry_detected = abs(resid_left_mean - resid_right_mean) > 5.0  # threshold

# Step 4 validation
fwhm_valid = 2.5 <= p1_results['fwhm'] <= 4.0
position_valid = abs(p1_results['center'] - 520.5) <= 2.0

print(f"First-order Si peak: center={p1_results['center']} cm-1, FWHM={p1_results['fwhm']} cm-1")
print(f"FWHM in valid range [2.5, 4.0]: {fwhm_valid}")
print(f"Position within 2 cm-1 of 520.5: {position_valid}")
print(f"Asymmetry near 520 cm-1 detected: {asymmetry_detected}")
print(f"Residual left mean: {resid_left_mean:.3f}, right mean: {resid_right_mean:.3f}")

# ---- Visualization ----
fig, axes = plt.subplots(3, 1, figsize=(12, 10), gridspec_kw={'height_ratios': [3, 1, 1]})

# Top panel: Original data + full model
ax1 = axes[0]
ax1.plot(x, y, 'k-', linewidth=0.5, alpha=0.7, label='Raw Data')
ax1.plot(x, y_model_full, 'r-', linewidth=1.5, label=f'Full Model (R²={r_squared:.4f})')
ax1.plot(x, baseline, 'g--', linewidth=1.0, alpha=0.7, label='ALS Baseline')

# Plot individual peak components (on top of baseline)
components = result.eval_components(x=x)
for prefix, label, color in [('p1_', 'Si 1st-order (~520)', 'blue'),
                               ('p2_', '2TA overtone (~300)', 'orange'),
                               ('p3_', '2TO overtone (~960)', 'purple')]:
    if prefix in components:
        ax1.plot(x, components[prefix] + baseline, '--', color=color, linewidth=1.2, alpha=0.8, label=label)

ax1.set_ylabel('Intensity (counts)', fontsize=12)
ax1.set_title('Raman Spectrum of Silicon: 3-Lorentzian Fit with ALS Baseline', fontsize=14)
ax1.legend(fontsize=9, loc='upper right')
ax1.set_xlim(x[0], x[-1])

# Middle panel: Baseline-corrected data + peak fits
ax2 = axes[1]
ax2.plot(x, y_corrected, 'k-', linewidth=0.5, alpha=0.7, label='Corrected Data')
ax2.plot(x, result.best_fit, 'r-', linewidth=1.5, label='Peak Model')
for prefix, label, color in [('p1_', 'Si 1st-order', 'blue'),
                               ('p2_', '2TA', 'orange'),
                               ('p3_', '2TO', 'purple')]:
    if prefix in components:
        ax2.fill_between(x, 0, components[prefix], alpha=0.2, color=color, label=label)

ax2.set_ylabel('Corrected Intensity', fontsize=12)
ax2.legend(fontsize=9, loc='upper right')
ax2.set_xlim(x[0], x[-1])

# Bottom panel: Residuals
ax3 = axes[2]
residuals = y - y_model_full
ax3.plot(x, residuals, 'k-', linewidth=0.5, alpha=0.7)
ax3.axhline(y=0, color='r', linestyle='--', linewidth=0.8)
ax3.set_xlabel('Raman Shift (cm⁻¹)', fontsize=12)
ax3.set_ylabel('Residuals', fontsize=12)
ax3.set_xlim(x[0], x[-1])

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

# ---- Build summary ----
validation_notes = []
if fwhm_valid:
    validation_notes.append(f"Si 1st-order FWHM ({p1_results['fwhm']:.2f} cm-1) is within expected range [2.5, 4.0] cm-1 for crystalline Si.")
else:
    validation_notes.append(f"WARNING: Si 1st-order FWHM ({p1_results['fwhm']:.2f} cm-1) is outside expected range [2.5, 4.0] cm-1.")
if position_valid:
    validation_notes.append(f"Si 1st-order position ({p1_results['center']:.2f} cm-1) is within 2 cm-1 of 520.5 cm-1 reference.")
else:
    validation_notes.append(f"WARNING: Si 1st-order position ({p1_results['center']:.2f} cm-1) deviates >2 cm-1 from 520.5 cm-1 reference.")
if asymmetry_detected:
    validation_notes.append("Residual asymmetry detected near 520 cm-1. Consider Voigt or Fano profile for improved fit.")
else:
    validation_notes.append("No significant residual asymmetry near 520 cm-1. Lorentzian profile is appropriate.")
validation_notes.append(f"Negative points after baseline correction: {neg_frac*100:.1f}% (threshold: <5%).")

summary_text = " ".join(validation_notes)

results = {
    "model_type": "ALS_baseline_plus_3_Lorentzian_peaks",
    "parameters": {
        "baseline": {
            "method": "ALS",
            "lambda": 1e6,
            "p": 0.001
        },
        "Si_first_order": {
            "center": p1_results['center'],
            "center_err": p1_results['center_err'],
            "fwhm": p1_results['fwhm'],
            "fwhm_err": p1_results['fwhm_err'],
            "peak_intensity": p1_results['peak_intensity'],
            "integrated_area": p1_results['integrated_area'],
            "amplitude": p1_results['amplitude'],
            "amplitude_err": p1_results['amplitude_err']
        },
        "2TA_overtone": {
            "center": p2_results['center'],
            "center_err": p2_results['center_err'],
            "fwhm": p2_results['fwhm'],
            "fwhm_err": p2_results['fwhm_err'],
            "peak_intensity": p2_results['peak_intensity'],
            "integrated_area": p2_results['integrated_area'],
            "amplitude": p2_results['amplitude'],
            "amplitude_err": p2_results['amplitude_err']
        },
        "2TO_overtone": {
            "center": p3_results['center'],
            "center_err": p3_results['center_err'],
            "fwhm": p3_results['fwhm'],
            "fwhm_err": p3_results['fwhm_err'],
            "peak_intensity": p3_results['peak_intensity'],
            "integrated_area": p3_results['integrated_area'],
            "amplitude": p3_results['amplitude'],
            "amplitude_err": p3_results['amplitude_err']
        }
    },
    "fit_quality": {
        "r_squared": round(r_squared, 6),
        "rmse": round(rmse, 4),
        "r_squared_corrected_spectrum": round(r_squared_corrected, 6)
    },
    "summary": summary_text
}

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