import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy.special import voigt_profile
import warnings
warnings.filterwarnings('ignore')

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

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

# Define Lorentzian profile
def lorentzian(x, A, c, gamma):
    return A / (1.0 + ((x - c) / gamma)**2)

# Define 3-Lorentzian + linear baseline model
def model_lorentzian(x, m, b, A1, c1, gamma1, A2, c2, gamma2, A3, c3, gamma3):
    baseline = m * x + b
    peak1 = lorentzian(x, A1, c1, gamma1)
    peak2 = lorentzian(x, A2, c2, gamma2)
    peak3 = lorentzian(x, A3, c3, gamma3)
    return baseline + peak1 + peak2 + peak3

# Define Gaussian profile for pseudo-Voigt
def gaussian(x, A, c, sigma):
    return A * np.exp(-0.5 * ((x - c) / sigma)**2)

# Define pseudo-Voigt: eta * Lorentzian + (1-eta) * Gaussian (per-peak eta)
def pseudo_voigt_peak(x, A, c, gamma, sigma, eta):
    L = A / (1.0 + ((x - c) / gamma)**2)
    G = A * np.exp(-0.5 * ((x - c) / sigma)**2)
    return eta * L + (1.0 - eta) * G

def model_pseudo_voigt(x, m, b, A1, c1, gamma1, sigma1, eta1,
                       A2, c2, gamma2, sigma2, eta2,
                       A3, c3, gamma3, sigma3, eta3):
    baseline = m * x + b
    peak1 = pseudo_voigt_peak(x, A1, c1, gamma1, sigma1, eta1)
    peak2 = pseudo_voigt_peak(x, A2, c2, gamma2, sigma2, eta2)
    peak3 = pseudo_voigt_peak(x, A3, c3, gamma3, sigma3, eta3)
    return baseline + peak1 + peak2 + peak3

# ========== PASS 1: Unweighted Lorentzian fit ==========
p0_lor = [0.005, 0.15,   # m, b
          4.88, 30.0, 1.0,   # A1, c1, gamma1
          2.88, 55.0, 1.3,   # A2, c2, gamma2
          3.87, 78.0, 0.8]   # A3, c3, gamma3

bounds_lower = [-0.05, -1.0,
                0.01, 28.0, 0.1,
                0.01, 53.0, 0.1,
                0.01, 76.0, 0.1]
bounds_upper = [0.05, 2.0,
                20.0, 32.0, 5.0,
                20.0, 57.0, 5.0,
                20.0, 80.0, 5.0]

try:
    popt_lor, pcov_lor = curve_fit(model_lorentzian, x, y, p0=p0_lor,
                                    bounds=(bounds_lower, bounds_upper),
                                    maxfev=50000, method='trf')
    y_fit_lor = model_lorentzian(x, *popt_lor)
    residuals_lor = y - y_fit_lor
    ss_res_lor = np.sum(residuals_lor**2)
    ss_tot = np.sum((y - np.mean(y))**2)
    r2_lor = 1.0 - ss_res_lor / ss_tot
    print(f"Pass 1 (Lorentzian) R^2 = {r2_lor:.6f}")
    lorentzian_success = True
except Exception as e:
    print(f"Lorentzian fit failed: {e}")
    lorentzian_success = False
    r2_lor = 0.0

# ========== PASS 2: Weighted Lorentzian fit ==========
if lorentzian_success:
    # Estimate weights from pass 1 residuals using rolling std
    window = max(20, len(x) // 50)
    local_std = np.array([np.std(residuals_lor[max(0, i-window):min(len(residuals_lor), i+window+1)])
                          for i in range(len(residuals_lor))])
    local_std = np.maximum(local_std, 1e-6)
    weights = 1.0 / local_std

    try:
        popt_lor_w, pcov_lor_w = curve_fit(model_lorentzian, x, y, p0=popt_lor,
                                             sigma=1.0/weights, absolute_sigma=True,
                                             bounds=(bounds_lower, bounds_upper),
                                             maxfev=50000, method='trf')
        y_fit_lor_w = model_lorentzian(x, *popt_lor_w)
        residuals_lor_w = y - y_fit_lor_w
        ss_res_lor_w = np.sum(residuals_lor_w**2)
        r2_lor_w = 1.0 - ss_res_lor_w / ss_tot
        print(f"Pass 2 (Weighted Lorentzian) R^2 = {r2_lor_w:.6f}")
    except Exception as e:
        print(f"Weighted Lorentzian fit failed: {e}, using unweighted")
        popt_lor_w = popt_lor
        pcov_lor_w = pcov_lor
        y_fit_lor_w = y_fit_lor
        residuals_lor_w = residuals_lor
        r2_lor_w = r2_lor
else:
    r2_lor_w = 0.0

# ========== Check if pseudo-Voigt needed ==========
use_voigt = False
if r2_lor_w < 0.995:
    print(f"Lorentzian R^2 = {r2_lor_w:.6f} < 0.995, attempting pseudo-Voigt...")
    use_voigt = True

    # Initialize pseudo-Voigt from Lorentzian results
    if lorentzian_success:
        m0, b0 = popt_lor_w[0], popt_lor_w[1]
        A1_0, c1_0, g1_0 = popt_lor_w[2], popt_lor_w[3], popt_lor_w[4]
        A2_0, c2_0, g2_0 = popt_lor_w[5], popt_lor_w[6], popt_lor_w[7]
        A3_0, c3_0, g3_0 = popt_lor_w[8], popt_lor_w[9], popt_lor_w[10]
    else:
        m0, b0 = 0.005, 0.15
        A1_0, c1_0, g1_0 = 4.88, 30.0, 1.0
        A2_0, c2_0, g2_0 = 2.88, 55.0, 1.3
        A3_0, c3_0, g3_0 = 3.87, 78.0, 0.8

    # p0: m, b, A1, c1, gamma1, sigma1, eta1, A2, c2, gamma2, sigma2, eta2, A3, c3, gamma3, sigma3, eta3
    p0_pv = [m0, b0,
             A1_0, c1_0, g1_0, g1_0 * 1.2, 0.7,
             A2_0, c2_0, g2_0, g2_0 * 1.2, 0.7,
             A3_0, c3_0, g3_0, g3_0 * 1.2, 0.7]

    bounds_lower_pv = [-0.05, -1.0,
                       0.01, 28.0, 0.05, 0.05, 0.0,
                       0.01, 53.0, 0.05, 0.05, 0.0,
                       0.01, 76.0, 0.05, 0.05, 0.0]
    bounds_upper_pv = [0.05, 2.0,
                       20.0, 32.0, 5.0, 5.0, 1.0,
                       20.0, 57.0, 5.0, 5.0, 1.0,
                       20.0, 80.0, 5.0, 5.0, 1.0]

    # Unweighted pseudo-Voigt
    try:
        popt_pv, pcov_pv = curve_fit(model_pseudo_voigt, x, y, p0=p0_pv,
                                      bounds=(bounds_lower_pv, bounds_upper_pv),
                                      maxfev=100000, method='trf')
        y_fit_pv = model_pseudo_voigt(x, *popt_pv)
        residuals_pv = y - y_fit_pv
        ss_res_pv = np.sum(residuals_pv**2)
        r2_pv = 1.0 - ss_res_pv / ss_tot
        print(f"Pseudo-Voigt unweighted R^2 = {r2_pv:.6f}")

        # Weighted pseudo-Voigt
        window = max(20, len(x) // 50)
        local_std_pv = np.array([np.std(residuals_pv[max(0, i-window):min(len(residuals_pv), i+window+1)])
                                  for i in range(len(residuals_pv))])
        local_std_pv = np.maximum(local_std_pv, 1e-6)
        weights_pv = 1.0 / local_std_pv

        popt_pv_w, pcov_pv_w = curve_fit(model_pseudo_voigt, x, y, p0=popt_pv,
                                           sigma=1.0/weights_pv, absolute_sigma=True,
                                           bounds=(bounds_lower_pv, bounds_upper_pv),
                                           maxfev=100000, method='trf')
        y_fit_pv_w = model_pseudo_voigt(x, *popt_pv_w)
        residuals_pv_w = y - y_fit_pv_w
        ss_res_pv_w = np.sum(residuals_pv_w**2)
        r2_pv_w = 1.0 - ss_res_pv_w / ss_tot
        print(f"Pseudo-Voigt weighted R^2 = {r2_pv_w:.6f}")
        voigt_success = True
    except Exception as e:
        print(f"Pseudo-Voigt fit failed: {e}")
        voigt_success = False
        r2_pv_w = 0.0
else:
    voigt_success = False
    r2_pv_w = 0.0

# ========== Select best model ==========
if use_voigt and voigt_success and r2_pv_w > r2_lor_w:
    model_name = "3-Peak Pseudo-Voigt + Linear Baseline"
    popt_final = popt_pv_w
    pcov_final = pcov_pv_w
    y_fit_final = y_fit_pv_w
    residuals_final = residuals_pv_w
    r2_final = r2_pv_w
    is_voigt = True
else:
    model_name = "3-Peak Lorentzian + Linear Baseline"
    popt_final = popt_lor_w if lorentzian_success else popt_lor
    pcov_final = pcov_lor_w if lorentzian_success else pcov_lor
    y_fit_final = y_fit_lor_w if lorentzian_success else y_fit_lor
    residuals_final = residuals_lor_w if lorentzian_success else residuals_lor
    r2_final = r2_lor_w if lorentzian_success else r2_lor
    is_voigt = False

rmse_final = np.sqrt(np.mean(residuals_final**2))
reduced_chi2 = np.sum(residuals_final**2) / (len(y) - len(popt_final))
residual_std = np.std(residuals_final)

# Parameter uncertainties
perr = np.sqrt(np.diag(pcov_final)) if pcov_final is not None else np.zeros(len(popt_final))

# ========== Extract parameters and compute derived quantities ==========
if is_voigt:
    # Pseudo-Voigt: m, b, A1, c1, gamma1, sigma1, eta1, A2, c2, gamma2, sigma2, eta2, A3, c3, gamma3, sigma3, eta3
    param_names = ['m', 'b', 'A1', 'c1', 'gamma1', 'sigma1', 'eta1',
                   'A2', 'c2', 'gamma2', 'sigma2', 'eta2',
                   'A3', 'c3', 'gamma3', 'sigma3', 'eta3']
    m_val, b_val = popt_final[0], popt_final[1]
    m_err, b_err = perr[0], perr[1]

    peaks_info = []
    for i in range(3):
        idx_base = 2 + i * 5
        A = popt_final[idx_base]
        c = popt_final[idx_base + 1]
        gamma = popt_final[idx_base + 2]
        sigma = popt_final[idx_base + 3]
        eta = popt_final[idx_base + 4]
        A_err = perr[idx_base]
        c_err = perr[idx_base + 1]
        gamma_err = perr[idx_base + 2]
        sigma_err = perr[idx_base + 3]
        eta_err = perr[idx_base + 4]

        # FWHM for Lorentzian component: 2*gamma
        fwhm_lor = 2.0 * gamma
        # FWHM for Gaussian component: 2*sigma*sqrt(2*ln(2))
        fwhm_gau = 2.0 * sigma * np.sqrt(2.0 * np.log(2.0))
        # Effective FWHM (approximate for pseudo-Voigt)
        fwhm_eff = eta * fwhm_lor + (1.0 - eta) * fwhm_gau

        # Integrated area: Lorentzian integral = A * pi * gamma, Gaussian integral = A * sigma * sqrt(2*pi)
        area_lor = A * np.pi * gamma
        area_gau = A * sigma * np.sqrt(2.0 * np.pi)
        area = eta * area_lor + (1.0 - eta) * area_gau

        # Compute component for plotting
        comp_y = pseudo_voigt_peak(x, A, c, gamma, sigma, eta)

        peaks_info.append({
            'A': A, 'A_err': A_err,
            'c': c, 'c_err': c_err,
            'gamma': gamma, 'gamma_err': gamma_err,
            'sigma': sigma, 'sigma_err': sigma_err,
            'eta': eta, 'eta_err': eta_err,
            'fwhm': fwhm_eff,
            'fwhm_lor': fwhm_lor, 'fwhm_gau': fwhm_gau,
            'area': area,
            'component_y': comp_y
        })
else:
    # Lorentzian: m, b, A1, c1, gamma1, A2, c2, gamma2, A3, c3, gamma3
    m_val, b_val = popt_final[0], popt_final[1]
    m_err, b_err = perr[0], perr[1]

    peaks_info = []
    for i in range(3):
        idx_base = 2 + i * 3
        A = popt_final[idx_base]
        c = popt_final[idx_base + 1]
        gamma = popt_final[idx_base + 2]
        A_err = perr[idx_base]
        c_err = perr[idx_base + 1]
        gamma_err = perr[idx_base + 2]

        fwhm = 2.0 * gamma
        fwhm_err = 2.0 * gamma_err
        area = A * np.pi * gamma  # Lorentzian integrated area
        # Error propagation for area: area = A * pi * gamma
        area_err = np.pi * np.sqrt((gamma * A_err)**2 + (A * gamma_err)**2)

        comp_y = lorentzian(x, A, c, gamma)

        peaks_info.append({
            'A': A, 'A_err': A_err,
            'c': c, 'c_err': c_err,
            'gamma': gamma, 'gamma_err': gamma_err,
            'fwhm': fwhm, 'fwhm_err': fwhm_err,
            'area': area, 'area_err': area_err,
            'component_y': comp_y
        })

# Baseline component
baseline_y = m_val * x + b_val

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

# Main plot
ax1 = axes[0]
ax1.scatter(x, y, s=1, alpha=0.4, color='gray', label='Data', zorder=1)
ax1.plot(x, y_fit_final, 'r-', linewidth=2, label=f'Fit (R²={r2_final:.6f})', zorder=3)
ax1.plot(x, baseline_y, 'k--', linewidth=1, alpha=0.5, label='Baseline', zorder=2)

colors = ['blue', 'green', 'orange']
for i, pinfo in enumerate(peaks_info):
    ax1.fill_between(x, baseline_y, baseline_y + pinfo['component_y'], alpha=0.2, color=colors[i],
                     label=f'Peak {i+1} (c={pinfo["c"]:.2f})')
    ax1.plot(x, baseline_y + pinfo['component_y'], color=colors[i], linewidth=1, alpha=0.7)

ax1.set_ylabel('Intensity', fontsize=12)
ax1.set_title(f'{model_name}', fontsize=14)
ax1.legend(fontsize=9, loc='upper right')
ax1.grid(True, alpha=0.3)

# Residuals plot
ax2 = axes[1]
ax2.scatter(x, residuals_final, s=1, alpha=0.4, color='gray')
ax2.axhline(y=0, color='r', linestyle='--', linewidth=1)
ax2.set_xlabel('x (spectral position)', fontsize=12)
ax2.set_ylabel('Residual', fontsize=12)
ax2.set_title(f'Residuals (RMSE={rmse_final:.6f}, std={residual_std:.6f})', fontsize=11)
ax2.grid(True, alpha=0.3)

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

# ========== Build results JSON ==========
results = {
    "model_type": model_name,
    "parameters": {
        "baseline": {
            "slope": float(m_val),
            "slope_err": float(m_err),
            "intercept": float(b_val),
            "intercept_err": float(b_err)
        }
    },
    "fit_quality": {
        "r_squared": float(r2_final),
        "rmse": float(rmse_final),
        "reduced_chi_squared": float(reduced_chi2),
        "residual_std": float(residual_std)
    },
    "summary": ""
}

for i, pinfo in enumerate(peaks_info):
    peak_key = f"peak_{i+1}"
    if is_voigt:
        results["parameters"][peak_key] = {
            "center": float(pinfo['c']),
            "center_err": float(pinfo['c_err']),
            "amplitude": float(pinfo['A']),
            "amplitude_err": float(pinfo['A_err']),
            "gamma_hwhm": float(pinfo['gamma']),
            "gamma_hwhm_err": float(pinfo['gamma_err']),
            "sigma": float(pinfo['sigma']),
            "sigma_err": float(pinfo['sigma_err']),
            "eta_mixing": float(pinfo['eta']),
            "eta_mixing_err": float(pinfo['eta_err']),
            "fwhm_effective": float(pinfo['fwhm']),
            "fwhm_lorentzian": float(pinfo['fwhm_lor']),
            "fwhm_gaussian": float(pinfo['fwhm_gau']),
            "integrated_area": float(pinfo['area'])
        }
    else:
        results["parameters"][peak_key] = {
            "center": float(pinfo['c']),
            "center_err": float(pinfo['c_err']),
            "amplitude": float(pinfo['A']),
            "amplitude_err": float(pinfo['A_err']),
            "gamma_hwhm": float(pinfo['gamma']),
            "gamma_hwhm_err": float(pinfo['gamma_err']),
            "fwhm": float(pinfo['fwhm']),
            "fwhm_err": float(pinfo['fwhm_err']),
            "integrated_area": float(pinfo['area']),
            "integrated_area_err": float(pinfo['area_err'])
        }

# Build summary
if is_voigt:
    eta_vals = [peaks_info[i]['eta'] for i in range(3)]
    summary = (f"Pseudo-Voigt model with 3 peaks on linear baseline. "
               f"R²={r2_final:.6f}, RMSE={rmse_final:.6f}. "
               f"Peak centers: {peaks_info[0]['c']:.2f}, {peaks_info[1]['c']:.2f}, {peaks_info[2]['c']:.2f}. "
               f"Lorentzian mixing fractions (eta): {eta_vals[0]:.3f}, {eta_vals[1]:.3f}, {eta_vals[2]:.3f}. ")
    if r2_final < 0.995:
        summary += (f"R² < 0.995 indicates residual model misspecification. "
                    f"Structured residuals around peaks suggest additional complexity "
                    f"(e.g., asymmetric lineshapes or additional minor peaks). ")
    else:
        summary += "Excellent fit quality achieved with pseudo-Voigt profiles. "
    if r2_lor_w > 0:
        summary += f"Pure Lorentzian R²={r2_lor_w:.6f} for comparison. "
else:
    summary = (f"Lorentzian model with 3 peaks on linear baseline. "
               f"R²={r2_final:.6f}, RMSE={rmse_final:.6f}. "
               f"Peak centers: {peaks_info[0]['c']:.2f}, {peaks_info[1]['c']:.2f}, {peaks_info[2]['c']:.2f}. ")
    if r2_final >= 0.995:
        summary += "Excellent fit quality; pure Lorentzian profiles sufficient. "
    else:
        summary += ("R² < 0.995; pseudo-Voigt was attempted but did not improve or failed. "
                    "Structured residuals may indicate lineshape complexity beyond 3-peak model. ")

summary += (f"Peak centers and integrated areas are high-confidence quantities. "
            f"Parameter uncertainties from covariance matrix may underestimate true uncertainty "
            f"due to systematic model misspecification if residuals show structure.")

results["summary"] = summary

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