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/analysis_session_20260401_224117/results/analysis_raman_silicon_CurveFit_20260401_224208_001/temp_spectrum_0.npy')
x = data[:, 0]
y = data[:, 1]

# Define model components
def fano(x, center, gamma, amplitude, q):
    """Fano line shape: A * (q + epsilon)^2 / (1 + epsilon^2) where epsilon = (x - center)/gamma"""
    epsilon = (x - center) / gamma
    return amplitude * (q + epsilon)**2 / (1.0 + epsilon**2)

def lorentzian(x, center, gamma, amplitude):
    """Lorentzian: A * gamma^2 / ((x - center)^2 + gamma^2)"""
    return amplitude * gamma**2 / ((x - center)**2 + gamma**2)

def full_model(x, y0, 
               si_center, si_gamma, si_amplitude, si_q,
               ta_center, ta_gamma, ta_amplitude,
               to_center, to_gamma, to_amplitude):
    """Full model: constant baseline + Fano (Si 520) + Lorentzian (2TA ~300) + Lorentzian (2TO ~950)"""
    return (y0 
            + fano(x, si_center, si_gamma, si_amplitude, si_q)
            + lorentzian(x, ta_center, ta_gamma, ta_amplitude)
            + lorentzian(x, to_center, to_gamma, to_amplitude))

# Create lmfit model
model = Model(full_model, independent_vars=['x'])

# --- Stage 1: Fit dominant Si peak alone to get robust Fano parameters ---
def fano_plus_baseline(x, y0, si_center, si_gamma, si_amplitude, si_q):
    return y0 + fano(x, si_center, si_gamma, si_amplitude, si_q)

stage1_model = Model(fano_plus_baseline, independent_vars=['x'])
stage1_params = Parameters()
stage1_params.add('y0', value=105, min=80, max=130)
stage1_params.add('si_center', value=520, min=505, max=535)
stage1_params.add('si_gamma', value=3, min=0.5, max=30)
stage1_params.add('si_amplitude', value=800, min=0)
stage1_params.add('si_q', value=-10, min=-100, max=-1)

# Use region around 520 for stage 1
mask_si = (x > 450) & (x < 600)
try:
    stage1_result = stage1_model.fit(y[mask_si], x=x[mask_si], params=stage1_params, method='leastsq')
    si_center_init = stage1_result.params['si_center'].value
    si_gamma_init = stage1_result.params['si_gamma'].value
    si_amplitude_init = stage1_result.params['si_amplitude'].value
    si_q_init = stage1_result.params['si_q'].value
    y0_init = stage1_result.params['y0'].value
    print(f"Stage 1 converged: center={si_center_init:.2f}, gamma={si_gamma_init:.2f}, q={si_q_init:.2f}, amp={si_amplitude_init:.2f}")
except Exception as e:
    print(f"Stage 1 failed ({e}), using initial guesses")
    si_center_init = 520
    si_gamma_init = 3
    si_amplitude_init = 800
    si_q_init = -10
    y0_init = 105

# --- Stage 2: Full fit with all 3 peaks + baseline ---
params = Parameters()
params.add('y0', value=y0_init, min=80, max=130)
# Fano peak (first-order Si phonon ~520 cm-1)
params.add('si_center', value=si_center_init, min=505, max=535)
params.add('si_gamma', value=si_gamma_init, min=0.5, max=30)
params.add('si_amplitude', value=si_amplitude_init, min=0)
params.add('si_q', value=si_q_init, min=-100, max=-1)
# 2TA overtone Lorentzian (~300 cm-1)
params.add('ta_center', value=300, min=285, max=315)
params.add('ta_gamma', value=25, min=1, max=80)
params.add('ta_amplitude', value=200, min=0)
# 2TO overtone Lorentzian (~950 cm-1)
params.add('to_center', value=950, min=935, max=965)
params.add('to_gamma', value=20, min=1, max=80)
params.add('to_amplitude', value=160, min=0)

result = model.fit(y, x=x, params=params, method='leastsq', max_nfev=10000)

# Extract fit and components
y_fit = result.best_fit
residuals = y - y_fit

p = result.params
y0_val = p['y0'].value
baseline_component = np.full_like(x, y0_val)
fano_component = fano(x, p['si_center'].value, p['si_gamma'].value, 
                      p['si_amplitude'].value, p['si_q'].value)
ta_component = lorentzian(x, p['ta_center'].value, p['ta_gamma'].value, p['ta_amplitude'].value)
to_component = lorentzian(x, p['to_center'].value, p['to_gamma'].value, p['to_amplitude'].value)

# Compute fit quality
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1.0 - ss_res / ss_tot
rmse = np.sqrt(np.mean(residuals**2))

print(f"R² = {r_squared:.6f}")
print(f"RMSE = {rmse:.4f}")
print(result.fit_report())

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

# Main plot
ax1.plot(x, y, 'k.', markersize=1.5, alpha=0.5, label='Data')
ax1.plot(x, y_fit, 'r-', linewidth=2, label='Best Fit')
ax1.plot(x, baseline_component + fano_component, 'b--', linewidth=1.2, alpha=0.7, 
         label=f'Fano (Si 1st order, {p["si_center"].value:.1f} cm$^{{-1}}$)')
ax1.plot(x, baseline_component + ta_component, 'g--', linewidth=1.2, alpha=0.7,
         label=f'Lorentzian (2TA, {p["ta_center"].value:.1f} cm$^{{-1}}$)')
ax1.plot(x, baseline_component + to_component, 'm--', linewidth=1.2, alpha=0.7,
         label=f'Lorentzian (2TO, {p["to_center"].value:.1f} cm$^{{-1}}$)')
ax1.axhline(y=y0_val, color='gray', linestyle=':', alpha=0.5, label=f'Baseline (y₀={y0_val:.1f})')
ax1.set_ylabel('Intensity (a.u.)', fontsize=12)
ax1.set_title('Raman Spectrum of p-type Silicon: Fano + 2×Lorentzian Fit', fontsize=13)
ax1.legend(fontsize=9, loc='upper right')
ax1.set_xlim(x.min(), x.max())

# Residuals plot
ax2.plot(x, residuals, 'k-', linewidth=0.8, alpha=0.7)
ax2.axhline(y=0, color='r', linestyle='--', alpha=0.5)
ax2.set_xlabel('Raman Shift (cm$^{-1}$)', fontsize=12)
ax2.set_ylabel('Residuals', fontsize=12)
ax2.set_title(f'Residuals (R² = {r_squared:.6f}, RMSE = {rmse:.4f})', fontsize=11)

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

# --- Compute FWHM for Fano peak ---
# For Fano, the FWHM is approximately 2*gamma but depends on q.
# We'll compute it numerically from the component.
fano_only = fano_component
fano_max = np.max(fano_only)
fano_hm = fano_max / 2.0
above_hm = np.where(fano_only >= fano_hm)[0]
if len(above_hm) > 1:
    fano_fwhm_numerical = x[above_hm[-1]] - x[above_hm[0]]
else:
    fano_fwhm_numerical = 2.0 * p['si_gamma'].value

# Compute FWHM for Lorentzians: FWHM = 2*gamma
ta_fwhm = 2.0 * p['ta_gamma'].value
to_fwhm = 2.0 * p['to_gamma'].value

# Helper to get value and stderr
def pval(name):
    return p[name].value

def perr(name):
    return p[name].stderr if p[name].stderr is not None else 0.0

# Build results
results = {
    "model_type": "Constant baseline + Fano (1st-order Si phonon ~520 cm-1) + Lorentzian (2TA overtone ~300 cm-1) + Lorentzian (2TO overtone ~950 cm-1); 11 parameters total",
    "parameters": {
        "baseline": {
            "y0": round(pval('y0'), 4),
            "y0_err": round(perr('y0'), 4)
        },
        "peak_1_Fano_Si_first_order": {
            "center": round(pval('si_center'), 4),
            "center_err": round(perr('si_center'), 4),
            "gamma": round(pval('si_gamma'), 4),
            "gamma_err": round(perr('si_gamma'), 4),
            "FWHM_numerical": round(fano_fwhm_numerical, 4),
            "amplitude": round(pval('si_amplitude'), 4),
            "amplitude_err": round(perr('si_amplitude'), 4),
            "q_asymmetry": round(pval('si_q'), 4),
            "q_err": round(perr('si_q'), 4),
            "lineshape": "Fano"
        },
        "peak_2_Lorentzian_2TA_overtone": {
            "center": round(pval('ta_center'), 4),
            "center_err": round(perr('ta_center'), 4),
            "gamma": round(pval('ta_gamma'), 4),
            "gamma_err": round(perr('ta_gamma'), 4),
            "FWHM": round(ta_fwhm, 4),
            "amplitude": round(pval('ta_amplitude'), 4),
            "amplitude_err": round(perr('ta_amplitude'), 4),
            "lineshape": "Lorentzian"
        },
        "peak_3_Lorentzian_2TO_overtone": {
            "center": round(pval('to_center'), 4),
            "center_err": round(perr('to_center'), 4),
            "gamma": round(pval('to_gamma'), 4),
            "gamma_err": round(perr('to_gamma'), 4),
            "FWHM": round(to_fwhm, 4),
            "amplitude": round(pval('to_amplitude'), 4),
            "amplitude_err": round(perr('to_amplitude'), 4),
            "lineshape": "Lorentzian"
        }
    },
    "fit_quality": {
        "r_squared": round(r_squared, 6),
        "rmse": round(rmse, 4),
        "reduced_chi_squared": round(result.redchi, 4) if result.redchi is not None else None,
        "n_data_points": len(x),
        "n_parameters": 11,
        "degrees_of_freedom": len(x) - 11
    },
    "summary": (
        f"Raman spectrum of p-type silicon fitted with Fano (1st-order Si phonon) + 2 Lorentzians (2TA, 2TO overtones) on a constant baseline. "
        f"The first-order Si phonon is at {pval('si_center'):.2f} cm-1 with Fano asymmetry parameter q={pval('si_q'):.2f} "
        f"(negative q consistent with p-type doping and Fano interference with free-hole continuum). "
        f"Fano half-width gamma={pval('si_gamma'):.2f} cm-1 (numerical FWHM ~ {fano_fwhm_numerical:.1f} cm-1). "
        f"2TA overtone at {pval('ta_center'):.1f} cm-1 (FWHM={ta_fwhm:.1f} cm-1), "
        f"2TO overtone at {pval('to_center'):.1f} cm-1 (FWHM={to_fwhm:.1f} cm-1). "
        f"R²={r_squared:.6f}, RMSE={rmse:.4f}. "
        f"Two-stage fitting strategy used: Stage 1 fitted Fano peak in 450-600 cm-1 region, "
        f"Stage 2 refined all 11 parameters simultaneously."
    )
}

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