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

# =============================================================================
# STEP 0 — PIPELINE VERIFICATION
# =============================================================================
data = np.load("/Users/maxim.ziatdinov/Code/SciLink/test_series_annealing_data/o3/temp_spectrum_1.npy")

print('Raw data shape:', data.shape)
print('Raw data dtype:', data.dtype)

# Handle different array shapes
if data.ndim == 2:
    if data.shape[1] == 2:
        x = data[:, 0]
        y = data[:, 1]
    elif data.shape[0] == 2:
        x = data[0, :]
        y = data[1, :]
    else:
        # Assume first row is x, second is y or treat as y only
        x = np.linspace(0, 100, data.shape[-1])
        y = data.flatten()
elif data.ndim == 1:
    # Only y values provided; construct x from known range
    y = data
    x = np.linspace(0, 100, len(y))
else:
    raise ValueError(f'Unexpected data shape: {data.shape}')

print(f'x range: [{x.min():.6f}, {x.max():.6f}], n_points={len(x)}')
print(f'y range: [{y.min():.6f}, {y.max():.6f}]')
print(f'x.shape={x.shape}, y.shape={y.shape}')

# CRITICAL: Validate x-range
if x.max() < 1.0:
    print('ERROR: x-range appears normalized or scaled; expected raw units with max~100')
    print('Applying pipeline fix: reconstructing x as linspace(0, 100, n)')
    x = np.linspace(0, 100, len(y))
    print(f'Fixed x range: [{x.min():.4f}, {x.max():.4f}]')
elif x.max() < 10.0:
    print(f'WARNING: x.max()={x.max():.4f} is suspiciously small (expected ~100). Rescaling x.')
    x = x * (100.0 / x.max())
    print(f'Rescaled x range: [{x.min():.4f}, {x.max():.4f}]')

assert x.max() > 10, f'x-range validation failed: max(x)={x.max():.4f}, expected ~100'
print(f'PIPELINE VALIDATION PASSED: x range confirmed [{x.min():.4f}, {x.max():.4f}]')
print(f'x_range_validation: max(x) = {x.max():.4f}')

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

# =============================================================================
# STEP 1 — VISUAL SANITY CHECK
# =============================================================================
fig_check, ax_check = plt.subplots(figsize=(10, 4))
ax_check.plot(x, y, 'k.', markersize=2, label='Raw data')
ax_check.set_xlabel('x')
ax_check.set_ylabel('y')
ax_check.set_title('Raw Data Sanity Check')
ax_check.legend()
plt.tight_layout()
plt.savefig('raw_data_check.png', dpi=100)
plt.close(fig_check)
print('Raw data plot saved to raw_data_check.png')

# =============================================================================
# STEP 2 — MODEL DEFINITION
# =============================================================================
def model(x, c, tau, A1, mu1, gamma1, A2, mu2, gamma2, A3, mu3, gamma3):
    """Exponential saturation baseline + 3 Lorentzian peaks."""
    baseline = c * (1.0 - np.exp(-x / tau))
    L1 = A1 * gamma1**2 / ((x - mu1)**2 + gamma1**2)
    L2 = A2 * gamma2**2 / ((x - mu2)**2 + gamma2**2)
    L3 = A3 * gamma3**2 / ((x - mu3)**2 + gamma3**2)
    return baseline + L1 + L2 + L3

# Initial guesses (physically grounded on x in [0, 100])
p0 = [
    0.5,   # c
    15.0,  # tau
    5.0,   # A1
    30.0,  # mu1
    1.5,   # gamma1
    3.3,   # A2
    55.0,  # mu2
    2.0,   # gamma2
    4.3,   # A3
    79.0,  # mu3
    2.0,   # gamma3
]

# Bounds
lower = [0.0, 1.0,  0.0, 24.0, 0.5,  0.0, 49.0, 0.5,  0.0, 73.0, 0.5]
upper = [2.0, 50.0, 20.0, 36.0, 10.0, 20.0, 61.0, 10.0, 20.0, 85.0, 10.0]

# Sigma weighting
sigma = np.full_like(y, 0.05)

# =============================================================================
# STEP 2 — FIT EXECUTION
# =============================================================================
convergence_flag = 'nominal'
popt = None
pcov = None

print('\nAttempting primary fit with TRF method...')
try:
    popt, pcov = curve_fit(
        model, x, y,
        p0=p0,
        bounds=(lower, upper),
        method='trf',
        sigma=sigma,
        absolute_sigma=True,
        maxfev=30000
    )
    perr = np.sqrt(np.diag(pcov))
    # Check for poorly conditioned covariance
    if np.any(np.diag(pcov) > 1e4):
        raise RuntimeError('Covariance matrix poorly conditioned; falling back.')
    print('Primary fit converged.')
except (RuntimeError, Exception) as e:
    print(f'Primary fit failed: {e}')
    print('Attempting warm-start fallback with modified initial guesses...')
    convergence_flag = 'warm_start_fallback'
    
    # Try with slightly perturbed guesses and looser sigma
    sigma_loose = np.full_like(y, 0.1)
    try:
        p0_warm = [
            0.3, 20.0,
            4.0, 30.0, 2.5,
            3.0, 55.0, 3.0,
            3.5, 79.0, 3.0
        ]
        popt, pcov = curve_fit(
            model, x, y,
            p0=p0_warm,
            bounds=(lower, upper),
            method='trf',
            sigma=sigma_loose,
            absolute_sigma=True,
            maxfev=50000
        )
        perr = np.sqrt(np.diag(pcov))
        print('Warm-start fallback converged.')
    except (RuntimeError, Exception) as e2:
        print(f'Warm-start fallback failed: {e2}')
        print('Executing two-stage fallback...')
        convergence_flag = 'two_stage_fallback'
        
        # Stage 1: Fit baseline only on non-peak regions
        peak_mask = np.ones(len(x), dtype=bool)
        for mu_i, g_i in [(30, 4.5), (55, 6.0), (79, 6.0)]:
            peak_mask &= np.abs(x - mu_i) > 3 * g_i
        x_base = x[peak_mask]
        y_base = y[peak_mask]
        
        def baseline_only(x, c, tau):
            return c * (1.0 - np.exp(-x / tau))
        
        try:
            p_base, _ = curve_fit(
                baseline_only, x_base, y_base,
                p0=[0.5, 15.0],
                bounds=([0.0, 1.0], [2.0, 50.0]),
                method='trf',
                maxfev=10000
            )
            c_fixed, tau_fixed = p_base
            print(f'Stage 1 baseline fit: c={c_fixed:.4f}, tau={tau_fixed:.4f}')
        except Exception:
            c_fixed, tau_fixed = 0.5, 15.0
            print('Stage 1 failed; using default baseline params.')
        
        # Stage 2: Fit peaks only with fixed baseline
        y_no_base = y - baseline_only(x, c_fixed, tau_fixed)
        
        def peaks_only(x, A1, mu1, gamma1, A2, mu2, gamma2, A3, mu3, gamma3):
            L1 = A1 * gamma1**2 / ((x - mu1)**2 + gamma1**2)
            L2 = A2 * gamma2**2 / ((x - mu2)**2 + gamma2**2)
            L3 = A3 * gamma3**2 / ((x - mu3)**2 + gamma3**2)
            return L1 + L2 + L3
        
        p0_peaks = [5.0, 30.0, 1.5, 3.3, 55.0, 2.0, 4.3, 79.0, 2.0]
        lb_peaks = [0.0, 24.0, 0.5, 0.0, 49.0, 0.5, 0.0, 73.0, 0.5]
        ub_peaks = [20.0, 36.0, 10.0, 20.0, 61.0, 10.0, 20.0, 85.0, 10.0]
        
        try:
            p_peaks, pcov_peaks = curve_fit(
                peaks_only, x, y_no_base,
                p0=p0_peaks,
                bounds=(lb_peaks, ub_peaks),
                method='trf',
                maxfev=20000
            )
            # Assemble full parameter vector
            popt = np.array([c_fixed, tau_fixed] + list(p_peaks))
            # Assemble approximate covariance (baseline params have large variance)
            n_full = 11
            pcov = np.zeros((n_full, n_full))
            pcov[0, 0] = 0.01
            pcov[1, 1] = 1.0
            pcov[2:, 2:] = pcov_peaks
            print('Two-stage fallback converged.')
        except Exception as e3:
            print(f'Two-stage fallback also failed: {e3}')
            print('Using initial guesses as final parameters (fit failed).')
            popt = np.array(p0)
            pcov = np.eye(11) * 1e6
            convergence_flag = 'FAILED_using_initial_guesses'

perr = np.sqrt(np.diag(pcov))

# =============================================================================
# STEP 3 — POST-FIT VALIDATION
# =============================================================================
c_fit, tau_fit, A1_fit, mu1_fit, g1_fit, A2_fit, mu2_fit, g2_fit, A3_fit, mu3_fit, g3_fit = popt

print('\n--- POST-FIT VALIDATION ---')
print(f'mu1={mu1_fit:.3f} (bounds: [24, 36]): {"PASS" if 24 <= mu1_fit <= 36 else "FAIL"}')
print(f'mu2={mu2_fit:.3f} (bounds: [49, 61]): {"PASS" if 49 <= mu2_fit <= 61 else "FAIL"}')
print(f'mu3={mu3_fit:.3f} (bounds: [73, 85]): {"PASS" if 73 <= mu3_fit <= 85 else "FAIL"}')
print(f'A1={A1_fit:.3f} > 0.1: {"PASS" if A1_fit > 0.1 else "FAIL"}')
print(f'A2={A2_fit:.3f} > 0.1: {"PASS" if A2_fit > 0.1 else "FAIL"}')
print(f'A3={A3_fit:.3f} > 0.1: {"PASS" if A3_fit > 0.1 else "FAIL"}')
print(f'tau={tau_fit:.3f} > 1.0: {"PASS" if tau_fit > 1.0 else "FAIL"}')

# =============================================================================
# STEP 4 — RESIDUAL ANALYSIS
# =============================================================================
y_fit = model(x, *popt)
residuals = y - y_fit

# R2
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r2 = 1.0 - ss_res / ss_tot
rmse_total = np.sqrt(np.mean(residuals**2))

# Region masks
peak_mask = (
    ((x >= 25) & (x <= 35)) |
    ((x >= 50) & (x <= 60)) |
    ((x >= 74) & (x <= 84))
)
baseline_mask = ((x >= 0) & (x <= 24)) | ((x >= 85) & (x <= 100))

rms_peak = np.sqrt(np.mean(residuals[peak_mask]**2)) if np.any(peak_mask) else float('nan')
rms_baseline = np.sqrt(np.mean(residuals[baseline_mask]**2)) if np.any(baseline_mask) else float('nan')

# R2 in peak regions
if np.any(peak_mask):
    ss_res_peak = np.sum(residuals[peak_mask]**2)
    ss_tot_peak = np.sum((y[peak_mask] - np.mean(y[peak_mask]))**2)
    r2_peak = 1.0 - ss_res_peak / ss_tot_peak if ss_tot_peak > 0 else float('nan')
else:
    r2_peak = float('nan')

print(f'\n--- FIT QUALITY ---')
print(f'R2_total = {r2:.6f}')
print(f'R2_peak_regions = {r2_peak:.6f}')
print(f'RMSE_total = {rmse_total:.6f}')
print(f'RMS_peak_regions = {rms_peak:.6f} (threshold: 0.25)')
print(f'RMS_baseline_regions = {rms_baseline:.6f} (threshold: 0.15)')

if rms_peak > 0.25:
    print('WARNING: Peak-region RMS exceeds threshold (0.25). Consider Voigt profiles.')
if rms_baseline > 0.15:
    print('WARNING: Baseline-region RMS exceeds threshold (0.15). Consider polynomial baseline.')

# =============================================================================
# STEP 5 — PARAMETER REPORTING
# =============================================================================
# Derived quantities
FWHM1 = 2.0 * g1_fit
FWHM2 = 2.0 * g2_fit
FWHM3 = 2.0 * g3_fit
area1 = np.pi * A1_fit * g1_fit
area2 = np.pi * A2_fit * g2_fit
area3 = np.pi * A3_fit * g3_fit

param_names = ['c', 'tau', 'A1', 'mu1', 'gamma1', 'A2', 'mu2', 'gamma2', 'A3', 'mu3', 'gamma3']
print('\n--- FITTED PARAMETERS (95% CI = 1.96 * stderr) ---')
for name, val, err in zip(param_names, popt, perr):
    ci95 = 1.96 * err
    flag = '*** POORLY CONSTRAINED ***' if (val != 0 and ci95 > 0.5 * abs(val)) else ''
    print(f'  {name:8s} = {val:10.4f} +/- {ci95:.4f} {flag}')

print(f'\nDerived quantities:')
print(f'  FWHM1 = {FWHM1:.4f}, FWHM2 = {FWHM2:.4f}, FWHM3 = {FWHM3:.4f}')
print(f'  Area1 = {area1:.4f}, Area2 = {area2:.4f}, Area3 = {area3:.4f}')
print(f'  convergence_flag = {convergence_flag}')

# =============================================================================
# STEP 3 — FIT VISUALIZATION
# =============================================================================
x_fine = np.linspace(x.min(), x.max(), 2000)
y_fine = model(x_fine, *popt)

# Individual components on fine grid
baseline_fine = c_fit * (1.0 - np.exp(-x_fine / tau_fit))
L1_fine = A1_fit * g1_fit**2 / ((x_fine - mu1_fit)**2 + g1_fit**2)
L2_fine = A2_fit * g2_fit**2 / ((x_fine - mu2_fit)**2 + g2_fit**2)
L3_fine = A3_fit * g3_fit**2 / ((x_fine - mu3_fit)**2 + g3_fit**2)

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

# Panel 1: Data + fit + components
ax1 = axes[0]
ax1.plot(x, y, 'k.', markersize=2, alpha=0.6, label='Data')
ax1.plot(x_fine, y_fine, 'r-', linewidth=2, label='Total fit')
ax1.plot(x_fine, baseline_fine, 'b--', linewidth=1.5, alpha=0.7, label='Baseline')
ax1.plot(x_fine, baseline_fine + L1_fine, 'g-', linewidth=1.5, alpha=0.8,
         label=f'Peak 1 (mu={mu1_fit:.1f})')
ax1.plot(x_fine, baseline_fine + L2_fine, 'm-', linewidth=1.5, alpha=0.8,
         label=f'Peak 2 (mu={mu2_fit:.1f})')
ax1.plot(x_fine, baseline_fine + L3_fine, 'c-', linewidth=1.5, alpha=0.8,
         label=f'Peak 3 (mu={mu3_fit:.1f})')
ax1.set_ylabel('Intensity')
ax1.set_title(f'Spectroscopic Fit: 3 Lorentzian + Exponential Baseline\nR² = {r2:.4f}, RMSE = {rmse_total:.4f}')
ax1.legend(loc='upper right', fontsize=8)
ax1.set_xlim(x.min(), x.max())

# Add vertical lines at peak centers
for mu_i, color in [(mu1_fit, 'g'), (mu2_fit, 'm'), (mu3_fit, 'c')]:
    ax1.axvline(mu_i, color=color, linestyle=':', alpha=0.5, linewidth=1)

# Panel 2: Residuals
ax2 = axes[1]
ax2.plot(x, residuals, 'k.', markersize=1.5, alpha=0.5)
ax2.axhline(0, color='r', linewidth=1)
ax2.axhline(0.25, color='orange', linewidth=1, linestyle='--', alpha=0.7, label='Peak RMS threshold')
ax2.axhline(-0.25, color='orange', linewidth=1, linestyle='--', alpha=0.7)
# Shade peak regions
for x_lo, x_hi in [(25, 35), (50, 60), (74, 84)]:
    ax2.axvspan(x_lo, x_hi, alpha=0.1, color='yellow')
ax2.set_ylabel('Residuals')
ax2.set_title(f'Residuals  |  Peak RMS={rms_peak:.4f}  |  Baseline RMS={rms_baseline:.4f}')
ax2.legend(fontsize=7)
ax2.set_xlim(x.min(), x.max())

# Panel 3: Normalized residuals
ax3 = axes[2]
norm_res = residuals / sigma
ax3.plot(x, norm_res, 'b.', markersize=1.5, alpha=0.5)
ax3.axhline(0, color='r', linewidth=1)
ax3.axhline(2, color='r', linewidth=1, linestyle='--', alpha=0.5)
ax3.axhline(-2, color='r', linewidth=1, linestyle='--', alpha=0.5)
ax3.set_xlabel('x')
ax3.set_ylabel('(y - f(x)) / sigma')
ax3.set_title('Normalized Residuals')
ax3.set_xlim(x.min(), x.max())

plt.tight_layout()
plt.savefig('spectrum_0001_fit.png', dpi=150, bbox_inches='tight')
plt.close(fig)
print('\nVisualization saved to spectrum_0001_fit.png')

# =============================================================================
# STEP 5 — RESULTS JSON
# =============================================================================
def safe_float(v):
    if np.isfinite(v):
        return float(v)
    return None

results = {
    'model_type': 'Exponential saturation baseline + 3 Lorentzian peaks: f(x) = c*(1-exp(-x/tau)) + sum_i[A_i*gamma_i^2/((x-mu_i)^2+gamma_i^2)]',
    'parameters': {
        'baseline': {
            'c': safe_float(c_fit),
            'c_err': safe_float(1.96 * perr[0]),
            'tau': safe_float(tau_fit),
            'tau_err': safe_float(1.96 * perr[1])
        },
        'peak_1': {
            'center': safe_float(mu1_fit),
            'center_err': safe_float(1.96 * perr[3]),
            'amplitude': safe_float(A1_fit),
            'amplitude_err': safe_float(1.96 * perr[2]),
            'gamma_hwhm': safe_float(g1_fit),
            'gamma_err': safe_float(1.96 * perr[4]),
            'fwhm': safe_float(FWHM1),
            'integrated_area': safe_float(area1)
        },
        'peak_2': {
            'center': safe_float(mu2_fit),
            'center_err': safe_float(1.96 * perr[6]),
            'amplitude': safe_float(A2_fit),
            'amplitude_err': safe_float(1.96 * perr[5]),
            'gamma_hwhm': safe_float(g2_fit),
            'gamma_err': safe_float(1.96 * perr[7]),
            'fwhm': safe_float(FWHM2),
            'integrated_area': safe_float(area2)
        },
        'peak_3': {
            'center': safe_float(mu3_fit),
            'center_err': safe_float(1.96 * perr[9]),
            'amplitude': safe_float(A3_fit),
            'amplitude_err': safe_float(1.96 * perr[8]),
            'gamma_hwhm': safe_float(g3_fit),
            'gamma_err': safe_float(1.96 * perr[10]),
            'fwhm': safe_float(FWHM3),
            'integrated_area': safe_float(area3)
        }
    },
    'fit_quality': {
        'r_squared': safe_float(r2),
        'r_squared_peak_regions': safe_float(r2_peak),
        'rmse': safe_float(rmse_total),
        'rms_peak_regions': safe_float(rms_peak),
        'rms_baseline_regions': safe_float(rms_baseline),
        'peak_rms_pass': bool(rms_peak <= 0.25),
        'baseline_rms_pass': bool(rms_baseline <= 0.15)
    },
    'diagnostics': {
        'convergence_flag': convergence_flag,
        'x_range_validation': f'max(x)={x.max():.4f}, confirmed correct domain',
        'n_points': int(len(x)),
        'x_min': safe_float(x.min()),
        'x_max': safe_float(x.max()),
        'y_min': safe_float(y.min()),
        'y_max': safe_float(y.max()),
        'pcov_diagonal': [safe_float(v) for v in np.diag(pcov)]
    },
    'summary': (
        f'Three Lorentzian peaks fitted at mu1={mu1_fit:.2f}, mu2={mu2_fit:.2f}, mu3={mu3_fit:.2f} '
        f'with FWHM={FWHM1:.2f}, {FWHM2:.2f}, {FWHM3:.2f} x-units respectively. '
        f'Exponential baseline: c={c_fit:.3f}, tau={tau_fit:.2f}. '
        f'Overall R2={r2:.4f}, RMSE={rmse_total:.4f}. '
        f'Peak-region RMS={rms_peak:.4f} ({"PASS" if rms_peak <= 0.25 else "FAIL"}), '
        f'Baseline-region RMS={rms_baseline:.4f} ({"PASS" if rms_baseline <= 0.15 else "FAIL"}). '
        f'Convergence: {convergence_flag}.'
    )
}

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