import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import stats

# ── Load data ──────────────────────────────────────────────────────────
with open('series_fit_results.json', 'r') as f:
    data = json.load(f)

results = data['results']
series_metadata = data.get('series_metadata', {})
conc = np.array(series_metadata['values'])  # ×10¹⁹ cm⁻³
unit = series_metadata.get('unit', '×10¹⁹ cm⁻³')
n = len(results)

# ── Extract parameters ─────────────────────────────────────────────────
peak_center = np.array([r['parameters']['fano_peak']['center'] for r in results])
peak_center_err = np.array([r['parameters']['fano_peak']['center_err'] for r in results])
fwhm = np.array([r['parameters']['fano_peak']['fwhm'] for r in results])
fwhm_err = np.array([r['parameters']['fano_peak']['fwhm_err'] for r in results])
q_asym = np.array([r['parameters']['fano_peak']['q_asymmetry'] for r in results])
inv_q = np.array([r['parameters']['fano_peak']['inverse_q'] for r in results])
inv_q_err = np.array([r['parameters']['fano_peak']['inverse_q_err'] for r in results])
r_squared = np.array([r['fit_quality']['r_squared'] for r in results])
red_chi2 = np.array([r['fit_quality']['reduced_chi_squared'] for r in results])

# Identify flagged indices
flagged_idx = [r['index'] for r in results if r.get('flagged', False)]
# Also flag index 8 as suspect (reduced_chi² ~1079 vs ~220 average)
suspect_idx = [i for i in range(n) if red_chi2[i] > 500 and i not in flagged_idx]

mask_good = np.array([(i not in flagged_idx) for i in range(n)])
mask_clean = np.array([(i not in flagged_idx and i not in suspect_idx) for i in range(n)])

# ── Weighted linear regression: peak center vs [B] ────────────────────
# Primary: exclude flagged only
w_primary = 1.0 / peak_center_err[mask_good]**2
from numpy.polynomial.polynomial import polyfit as np_polyfit

def weighted_linreg(x, y, w):
    """Weighted least squares: y = slope*x + intercept. Returns slope, intercept, and uncertainties."""
    sw = np.sum(w)
    sx = np.sum(w * x)
    sy = np.sum(w * y)
    sxx = np.sum(w * x**2)
    sxy = np.sum(w * x * y)
    delta = sw * sxx - sx**2
    intercept = (sxx * sy - sx * sxy) / delta
    slope = (sw * sxy - sx * sy) / delta
    # Residuals
    resid = y - (slope * x + intercept)
    # Weighted residual variance
    n_pts = len(x)
    s2 = np.sum(w * resid**2) / (n_pts - 2) if n_pts > 2 else 1.0
    intercept_err = np.sqrt(s2 * sxx / delta)
    slope_err = np.sqrt(s2 * sw / delta)
    # R² (weighted)
    ss_tot = np.sum(w * (y - np.sum(w * y) / sw)**2)
    ss_res = np.sum(w * resid**2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0
    return slope, intercept, slope_err, intercept_err, r2, resid

# Regression excluding flagged
slope1, intercept1, slope1_err, intercept1_err, r2_1, resid1 = weighted_linreg(
    conc[mask_good], peak_center[mask_good], w_primary)

# Regression excluding flagged AND suspect
w_clean = 1.0 / peak_center_err[mask_clean]**2
slope2, intercept2, slope2_err, intercept2_err, r2_2, resid2 = weighted_linreg(
    conc[mask_clean], peak_center[mask_clean], w_clean)

# FWHM linear trend (unweighted for simplicity, exclude flagged)
slope_fwhm, intercept_fwhm, _, _, _ , _ = weighted_linreg(
    conc[mask_good], fwhm[mask_good], 1.0/fwhm_err[mask_good]**2)

# ── Print calibration results ──────────────────────────────────────────
print('='*70)
print('CALIBRATION MODEL: Si Raman Peak Position vs Boron Concentration')
print('='*70)
print(f'\nModel: ω(cm⁻¹) = m × [B] + b   where [B] in {unit}')
print(f'\n--- Primary (excluding flagged index {flagged_idx}) ---')
print(f'  Slope     m = {slope1:.4f} ± {slope1_err:.4f} cm⁻¹ / (10¹⁹ cm⁻³)')
print(f'  Intercept b = {intercept1:.4f} ± {intercept1_err:.4f} cm⁻¹')
print(f'  R²          = {r2_1:.6f}')
print(f'\n--- Clean (excluding flagged + suspect idx {flagged_idx + suspect_idx}) ---')
print(f'  Slope     m = {slope2:.4f} ± {slope2_err:.4f} cm⁻¹ / (10¹⁹ cm⁻³)')
print(f'  Intercept b = {intercept2:.4f} ± {intercept2_err:.4f} cm⁻¹')
print(f'  R²          = {r2_2:.6f}')
print(f'\n--- FWHM broadening trend (excluding flagged) ---')
print(f'  FWHM slope  = {slope_fwhm:.4f} cm⁻¹ / (10¹⁹ cm⁻³)')
print(f'\n--- Peak shift summary ---')
print(f'  Undoped reference (extrapolated): {intercept1:.2f} cm⁻¹')
print(f'  Total shift over range: {peak_center[mask_good][-1] - peak_center[mask_good][0]:.2f} cm⁻¹')
print(f'  Sensitivity: ~{abs(slope1):.3f} cm⁻¹ per 10¹⁹ cm⁻³ boron')
print('='*70)

# ── Dashboard figure ───────────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
fig.suptitle('B-doped Si Raman Calibration — Parameter Trends', fontsize=15, fontweight='bold', y=0.98)

colors_good = '#2166ac'
colors_flag = '#d73027'
colors_suspect = '#fc8d59'
ms = 8

# Helper: plot with flagged markers
def plot_with_flags(ax, x, y, yerr, ylabel, title, flagged=flagged_idx, suspect=suspect_idx):
    good = [i for i in range(n) if i not in flagged and i not in suspect]
    # Good points
    ax.errorbar(x[good], y[good], yerr=yerr[good] if yerr is not None else None,
                fmt='o', color=colors_good, ms=ms, capsize=3, label='Good fits', zorder=3)
    # Suspect points
    if suspect:
        ax.errorbar(x[suspect], y[suspect], yerr=yerr[suspect] if yerr is not None else None,
                    fmt='s', color=colors_suspect, ms=ms, capsize=3, label='Suspect', zorder=3)
    # Flagged points
    if flagged:
        ax.errorbar(x[flagged], y[flagged], yerr=yerr[flagged] if yerr is not None else None,
                    fmt='X', color=colors_flag, ms=ms+3, capsize=3, label='Flagged', zorder=4,
                    markeredgecolor='darkred', markeredgewidth=1.5)
    ax.set_xlabel(f'[B] ({unit})', fontsize=10)
    ax.set_ylabel(ylabel, fontsize=10)
    ax.set_title(title, fontsize=11, fontweight='bold')
    ax.tick_params(labelsize=9)

# ── Panel (0,0): Calibration plot — Peak Center vs [B] ────────────────
ax = axes[0, 0]
plot_with_flags(ax, conc, peak_center, peak_center_err,
                'Fano Peak Center (cm⁻¹)', 'Calibration: Peak Position vs [B]')
# Regression line (primary)
xfit = np.linspace(0, conc.max()*1.05, 200)
ax.plot(xfit, slope1*xfit + intercept1, '-', color='#1a9850', lw=2, alpha=0.85,
        label=f'Fit (excl. flagged)\nω = {slope1:.3f}[B] + {intercept1:.2f}')
# Confidence band
resid_std = np.std(peak_center[mask_good] - (slope1*conc[mask_good] + intercept1))
ax.fill_between(xfit, slope1*xfit + intercept1 - 2*resid_std,
                slope1*xfit + intercept1 + 2*resid_std,
                alpha=0.12, color='#1a9850', label='±2σ band')
ax.legend(fontsize=7.5, loc='lower left')
ax.set_xlim(-0.5, conc.max()+1)

# ── Panel (0,1): Residuals from calibration ────────────────────────────
ax = axes[0, 1]
resid_all = peak_center - (slope1*conc + intercept1)
good_i = [i for i in range(n) if i not in flagged_idx and i not in suspect_idx]
ax.errorbar(conc[good_i], resid_all[good_i], yerr=peak_center_err[good_i],
            fmt='o', color=colors_good, ms=ms, capsize=3, label='Good fits')
if suspect_idx:
    ax.errorbar(conc[suspect_idx], resid_all[suspect_idx], yerr=peak_center_err[suspect_idx],
                fmt='s', color=colors_suspect, ms=ms, capsize=3, label='Suspect')
if flagged_idx:
    ax.errorbar(conc[flagged_idx], resid_all[flagged_idx], yerr=peak_center_err[flagged_idx],
                fmt='X', color=colors_flag, ms=ms+3, capsize=3, label='Flagged',
                markeredgecolor='darkred', markeredgewidth=1.5)
ax.axhline(0, color='gray', ls='--', lw=1)
ax.axhline(2*resid_std, color='#1a9850', ls=':', lw=1, alpha=0.5)
ax.axhline(-2*resid_std, color='#1a9850', ls=':', lw=1, alpha=0.5)
ax.set_xlabel(f'[B] ({unit})', fontsize=10)
ax.set_ylabel('Residual (cm⁻¹)', fontsize=10)
ax.set_title('Calibration Residuals', fontsize=11, fontweight='bold')
ax.legend(fontsize=7.5)
ax.tick_params(labelsize=9)

# ── Panel (0,2): FWHM broadening ──────────────────────────────────────
ax = axes[0, 2]
plot_with_flags(ax, conc, fwhm, fwhm_err,
                'Fano FWHM (cm⁻¹)', 'Linewidth Broadening vs [B]')
ax.plot(xfit, slope_fwhm*xfit + intercept_fwhm, '-', color='#1a9850', lw=2, alpha=0.8,
        label=f'Slope = {slope_fwhm:.4f} cm⁻¹/(10¹⁹cm⁻³)')
ax.legend(fontsize=7.5, loc='upper left')

# ── Panel (1,0): 1/q asymmetry ─────────────────────────────────────────
ax = axes[1, 0]
plot_with_flags(ax, conc, inv_q, inv_q_err,
                '1/q (Fano asymmetry)', 'Fano Asymmetry Parameter 1/q')
ax.axhline(0, color='gray', ls='--', lw=0.8)
ax.legend(fontsize=7.5)

# ── Panel (1,1): Fit quality — R² ─────────────────────────────────────
ax = axes[1, 1]
plot_with_flags(ax, conc, r_squared, None,
                'R²', 'Fit Quality (R²) Across Series')
ax.axhline(np.mean(r_squared[mask_good]), color='gray', ls=':', lw=1, label=f'Mean (good) = {np.mean(r_squared[mask_good]):.6f}')
ax.set_ylim(min(r_squared.min()*0.9999, 0.992), 1.0001)
ax.legend(fontsize=7.5)
ax.ticklabel_format(axis='y', useOffset=False)

# ── Panel (1,2): Reduced chi² ──────────────────────────────────────────
ax = axes[1, 2]
plot_with_flags(ax, conc, red_chi2, None,
                'Reduced χ²', 'Reduced χ² Across Series')
ax.axhline(np.mean(red_chi2[mask_clean]), color='gray', ls=':', lw=1,
           label=f'Mean (clean) = {np.mean(red_chi2[mask_clean]):.0f}')
ax.set_yscale('log')
ax.legend(fontsize=7.5)

# ── Annotation box with calibration summary ────────────────────────────
textstr = (
    'Calibration Summary\n'
    '─────────────────────\n'
    f'ω = ({slope1:.4f} ± {slope1_err:.4f})·[B]\n'
    f'    + ({intercept1:.2f} ± {intercept1_err:.2f}) cm⁻¹\n'
    f'R² = {r2_1:.6f}\n'
    f'Sensitivity: {abs(slope1):.3f} cm⁻¹/(10¹⁹cm⁻³)\n'
    f'N pts used: {int(mask_good.sum())} / {n}'
)
fig.text(0.5, 0.01, textstr, transform=fig.transFigure,
         fontsize=10, verticalalignment='bottom', horizontalalignment='center',
         bbox=dict(boxstyle='round,pad=0.5', facecolor='lightyellow', alpha=0.9, edgecolor='gray'),
         fontfamily='monospace')

plt.tight_layout(rect=[0, 0.09, 1, 0.96])
plt.savefig('parameter_trends.png', dpi=150, bbox_inches='tight')
plt.close('all')

# ── Also output calibration as JSON-like summary ───────────────────────
print('\n--- Machine-readable calibration output ---')
calib = {
    'model': 'linear',
    'equation': 'omega_cm-1 = slope * [B]_1e19 + intercept',
    'slope': round(slope1, 5),
    'slope_err': round(slope1_err, 5),
    'slope_unit': 'cm-1 per 1e19 cm-3',
    'intercept': round(intercept1, 4),
    'intercept_err': round(intercept1_err, 4),
    'intercept_unit': 'cm-1',
    'R_squared': round(r2_1, 6),
    'n_points_used': int(mask_good.sum()),
    'excluded_indices': flagged_idx,
    'valid_range_1e19': [float(conc[mask_good].min()), float(conc[mask_good].max())],
    'fwhm_broadening_slope': round(slope_fwhm, 5),
    'fwhm_broadening_slope_unit': 'cm-1 per 1e19 cm-3'
}
print(json.dumps(calib, indent=2))
