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_193636/results/analysis_raman_silicon_CurveFit_20260310_195008_001/temp_spectrum_0.npy'
data = np.load(DATA_PATH)
if data.ndim == 1:
    n = len(data) // 2
    x_raw, y_raw = data[:n], data[n:]
elif data.shape[1] == 2:
    x_raw, y_raw = data[:, 0], data[:, 1]
elif data.shape[0] == 2:
    x_raw, y_raw = data[0, :], data[1, :]
else:
    x_raw, y_raw = data[:, 0], data[:, 1]

sort_idx = np.argsort(x_raw)
x = x_raw[sort_idx].astype(float)
y = y_raw[sort_idx].astype(float)

# ---- Step 3: ALS baseline correction ----
def als_baseline(y, lam=1e6, p=0.001, niter=50):
    L = len(y)
    D = sparse.diags([1, -2, 1], [0, -1, -2], shape=(L, L - 2))
    D = D.tocsc()
    H = lam * D.dot(D.T)
    w = np.ones(L)
    for i in range(niter):
        W = sparse.spdiags(w, 0, L, L)
        C = W + H
        z = spsolve(C.tocsc(), 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_corr = y - baseline

# Check that fewer than 5% of baseline-corrected points are negative
neg_frac = np.sum(y_corr < 0) / len(y_corr)
print(f"Fraction of negative baseline-corrected points: {neg_frac:.4f}")

# ---- Step 4: Build 3-peak Lorentzian model ----
# Peak 1: Si F2g first-order phonon at ~520 cm-1
# Plan: center=520, FWHM=3.5 -> sigma=FWHM/2=1.75, amplitude from plan=800
# lmfit Lorentzian: amplitude/(pi*sigma) * sigma^2/((x-center)^2 + sigma^2)
# The plan's Lorentzian: A * gamma^2 / ((x-x0)^2 + gamma^2)
# lmfit amplitude = A * pi * sigma, where sigma = gamma = HWHM
# Plan amplitude_520=800 -> lmfit amplitude = 800 * pi * 1.75 ~ 4398
# But let's use a data-driven initial guess for robustness

peak1 = LorentzianModel(prefix='p1_')
peak2 = LorentzianModel(prefix='p2_')
peak3 = LorentzianModel(prefix='p3_')

composite = peak1 + peak2 + peak3

params = composite.make_params()

# Estimate initial amplitude from data for peak 1 (520 cm-1)
idx_520 = np.argmin(np.abs(x - 520))
max_intensity_520 = y_corr[idx_520] if y_corr[idx_520] > 0 else np.max(y_corr)
# lmfit Lorentzian peak height = amplitude / (pi * sigma)
# So amplitude = height * pi * sigma
initial_sigma_p1 = 1.75  # FWHM=3.5 -> sigma=1.75
initial_amp_p1 = max_intensity_520 * np.pi * initial_sigma_p1

# Peak 1: Si F2g at ~520 cm-1
params['p1_center'].set(value=520.0, min=515.0, max=525.0)
params['p1_sigma'].set(value=initial_sigma_p1, min=1.0, max=7.5)  # FWHM=2*sigma, plan FWHM [2.0,15.0] -> sigma [1.0, 7.5]
params['p1_amplitude'].set(value=initial_amp_p1, min=0)

# Peak 2: Si 2TA at ~300 cm-1
idx_300 = np.argmin(np.abs(x - 302))
intensity_300 = max(y_corr[idx_300], 10)
initial_sigma_p2 = 15.0  # FWHM=30 -> sigma=15
initial_amp_p2 = intensity_300 * np.pi * initial_sigma_p2

params['p2_center'].set(value=302.0, min=280.0, max=320.0)
params['p2_sigma'].set(value=initial_sigma_p2, min=7.5, max=30.0)  # FWHM [15,60] -> sigma [7.5, 30]
params['p2_amplitude'].set(value=initial_amp_p2, min=0)

# Peak 3: Si 2TO at ~960 cm-1
idx_960 = np.argmin(np.abs(x - 950))
intensity_960 = max(y_corr[idx_960], 10)
initial_sigma_p3 = 12.5  # FWHM=25 -> sigma=12.5
initial_amp_p3 = intensity_960 * np.pi * initial_sigma_p3

params['p3_center'].set(value=950.0, min=930.0, max=970.0)
params['p3_sigma'].set(value=initial_sigma_p3, min=7.5, max=25.0)  # FWHM [15,50] -> sigma [7.5, 25]
params['p3_amplitude'].set(value=initial_amp_p3, min=0)

# Fit using Levenberg-Marquardt on baseline-corrected data
result = composite.fit(y_corr, params, x=x, method='leastsq')

# ---- Step 5: Compute R-squared on ORIGINAL data (baseline + peaks vs original) ----
# As per mandatory domain rules: Report R-squared on full spectrum (baseline + peaks vs original data)
y_fit_total = baseline + result.best_fit
ss_res = np.sum((y - y_fit_total) ** 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_total) ** 2))

# Also compute R-squared on baseline-corrected data for additional info
ss_res_corr = np.sum((y_corr - result.best_fit) ** 2)
ss_tot_corr = np.sum((y_corr - np.mean(y_corr)) ** 2)
r_squared_corr = 1.0 - ss_res_corr / ss_tot_corr

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

# Extract parameters
def get_peak_params(result, prefix):
    center = result.params[f'{prefix}center'].value
    center_err = result.params[f'{prefix}center'].stderr
    sigma = result.params[f'{prefix}sigma'].value
    sigma_err = result.params[f'{prefix}sigma'].stderr
    amplitude = result.params[f'{prefix}amplitude'].value
    amplitude_err = result.params[f'{prefix}amplitude'].stderr
    fwhm = 2.0 * sigma
    fwhm_err = 2.0 * sigma_err if sigma_err is not None else None
    height = amplitude / (np.pi * sigma) if sigma > 0 else 0
    height_err = None
    if amplitude_err is not None and sigma_err is not None and sigma > 0:
        # Propagate: height = amp/(pi*sigma)
        height_err = height * np.sqrt((amplitude_err/amplitude)**2 + (sigma_err/sigma)**2) if amplitude > 0 else None
    return {
        'center': round(center, 4),
        'center_err': round(center_err, 4) if center_err is not None else None,
        'fwhm': round(fwhm, 4),
        'fwhm_err': round(fwhm_err, 4) if fwhm_err is not None else None,
        'amplitude_lmfit': round(amplitude, 4),
        'amplitude_lmfit_err': round(amplitude_err, 4) if amplitude_err is not None else None,
        'height': round(height, 4),
        'height_err': round(height_err, 4) if height_err is not None else None,
        'sigma': round(sigma, 4),
        'sigma_err': round(sigma_err, 4) if sigma_err is not None else None
    }

peak1_params = get_peak_params(result, 'p1_')
peak2_params = get_peak_params(result, 'p2_')
peak3_params = get_peak_params(result, 'p3_')

# Compute integrated areas using np.trapezoid on individual components
comps = result.eval_components(x=x)
peak1_curve = comps['p1_']
peak2_curve = comps['p2_']
peak3_curve = comps['p3_']

area_p1 = np.trapezoid(peak1_curve, x)
area_p2 = np.trapezoid(peak2_curve, x)
area_p3 = np.trapezoid(peak3_curve, x)

# Intensity ratio: second-order to first-order
intensity_ratio_2nd_to_1st = (area_p2 + area_p3) / area_p1 if area_p1 > 0 else None

# Check residuals near 520 cm-1 for systematic structure (Fano check)
mask_520 = (x > 510) & (x < 530)
residuals_full = y - y_fit_total
residuals_520 = residuals_full[mask_520]
# Check for asymmetry in residuals
if len(residuals_520) > 5:
    x_520_region = x[mask_520]
    center_520 = peak1_params['center']
    low_side = residuals_520[x_520_region < center_520]
    high_side = residuals_520[x_520_region > center_520]
    mean_low = np.mean(low_side) if len(low_side) > 0 else 0
    mean_high = np.mean(high_side) if len(high_side) > 0 else 0
    asymmetry_note = f"Residual asymmetry near 520 cm-1: low-side mean={mean_low:.2f}, high-side mean={mean_high:.2f}"
else:
    asymmetry_note = "Insufficient points near 520 cm-1 to check asymmetry"

print(asymmetry_note)

# Validation summary
fwhm_520 = peak1_params['fwhm']
fwhm_validation = "PASS" if 2.5 <= fwhm_520 <= 4.0 else f"NOTE: FWHM={fwhm_520:.2f} outside expected 2.5-4.0 cm-1 for single-crystal Si"
r2_validation = "PASS" if r_squared > 0.98 else f"NOTE: R-squared={r_squared:.4f} below 0.98 threshold"

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

# Main plot: data + fit + components
ax1 = axes[0]
ax1.plot(x, y, 'k-', linewidth=0.5, alpha=0.7, label='Raw Data')
ax1.plot(x, y_fit_total, 'r-', linewidth=1.5, label=f'Total Fit (R²={r_squared:.5f})')
ax1.plot(x, baseline, 'b--', linewidth=1.0, alpha=0.7, label='ALS Baseline')
ax1.plot(x, baseline + peak1_curve, 'g--', linewidth=1.0, alpha=0.7, label=f'Si F2g (~{peak1_params["center"]:.1f} cm⁻¹)')
ax1.plot(x, baseline + peak2_curve, 'm--', linewidth=1.0, alpha=0.7, label=f'Si 2TA (~{peak2_params["center"]:.1f} cm⁻¹)')
ax1.plot(x, baseline + peak3_curve, 'c--', linewidth=1.0, alpha=0.7, label=f'Si 2TO (~{peak3_params["center"]:.1f} cm⁻¹)')
ax1.set_ylabel('Intensity (counts)', fontsize=12)
ax1.set_title('Raman Spectrum of Crystalline Silicon - 3-Peak Lorentzian Fit', fontsize=14)
ax1.legend(fontsize=9, loc='upper right')
ax1.set_xlim(x[0], x[-1])

# Residuals
ax2 = axes[1]
ax2.plot(x, residuals_full, 'k-', linewidth=0.5)
ax2.axhline(y=0, color='r', linestyle='--', linewidth=0.5)
ax2.set_ylabel('Residuals', fontsize=12)
ax2.set_xlim(x[0], x[-1])
ax2.set_title('Residuals (Data - Fit)', fontsize=11)

# Baseline-corrected with peak components
ax3 = axes[2]
ax3.plot(x, y_corr, 'k-', linewidth=0.5, alpha=0.7, label='Baseline-corrected')
ax3.plot(x, result.best_fit, 'r-', linewidth=1.5, label='Peak Model')
ax3.fill_between(x, peak1_curve, alpha=0.3, color='green', label='Si F2g')
ax3.fill_between(x, peak2_curve, alpha=0.3, color='magenta', label='Si 2TA')
ax3.fill_between(x, peak3_curve, alpha=0.3, color='cyan', label='Si 2TO')
ax3.set_xlabel('Raman Shift (cm⁻¹)', fontsize=12)
ax3.set_ylabel('Intensity', fontsize=12)
ax3.set_title('Baseline-Corrected Spectrum with Peak Components', fontsize=11)
ax3.legend(fontsize=9, loc='upper right')
ax3.set_xlim(x[0], x[-1])

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

# ---- Build summary ----
summary_parts = []
summary_parts.append(f"Three Lorentzian peaks fitted to crystalline Si Raman spectrum after ALS baseline correction (lambda=1e6, p=0.001).")
summary_parts.append(f"Si F2g first-order phonon at {peak1_params['center']:.2f} cm-1 with FWHM={fwhm_520:.2f} cm-1.")
summary_parts.append(f"Si 2TA second-order overtone at {peak2_params['center']:.2f} cm-1 with FWHM={peak2_params['fwhm']:.2f} cm-1.")
summary_parts.append(f"Si 2TO second-order overtone at {peak3_params['center']:.2f} cm-1 with FWHM={peak3_params['fwhm']:.2f} cm-1.")
summary_parts.append(f"Intensity ratio (2nd-order/1st-order): {intensity_ratio_2nd_to_1st:.4f}." if intensity_ratio_2nd_to_1st is not None else "Could not compute intensity ratio.")
summary_parts.append(f"FWHM validation for 520 peak: {fwhm_validation}.")
summary_parts.append(f"R-squared validation: {r2_validation}.")
summary_parts.append(asymmetry_note)
summary_parts.append(f"Negative baseline-corrected fraction: {neg_frac:.4f} ({'PASS' if neg_frac < 0.05 else 'WARN: >5%'}).")

results = {
    "model_type": "ALS baseline (lambda=1e6, p=0.001) + 3 Lorentzian peaks (Si F2g, Si 2TA, Si 2TO)",
    "parameters": {
        "peak_1_Si_F2g": peak1_params,
        "peak_2_Si_2TA": peak2_params,
        "peak_3_Si_2TO": peak3_params,
        "integrated_area_Si_F2g": round(area_p1, 4),
        "integrated_area_Si_2TA": round(area_p2, 4),
        "integrated_area_Si_2TO": round(area_p3, 4),
        "intensity_ratio_2nd_to_1st": round(intensity_ratio_2nd_to_1st, 6) if intensity_ratio_2nd_to_1st is not None else None
    },
    "fit_quality": {
        "r_squared": round(r_squared, 6),
        "r_squared_baseline_corrected": round(r_squared_corr, 6),
        "rmse": round(rmse, 4),
        "reduced_chi_squared": round(result.redchi, 6) if result.redchi is not None else None
    },
    "summary": " ".join(summary_parts)
}

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