import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import warnings
warnings.filterwarnings('ignore')

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

results = data['results']
series_metadata = data.get('series_metadata', {})

# ── Series x-axis ──────────────────────────────────────────────────────────
if series_metadata and 'values' in series_metadata:
    x_vals  = np.array(series_metadata['values'], dtype=float)
    x_label = series_metadata.get('variable', 'Series Index')
    x_unit  = series_metadata.get('unit', '')
    if x_unit:
        x_label = f"{x_label} ({x_unit})"
else:
    x_vals  = np.array([r['index'] for r in results], dtype=float)
    x_label = 'Spectrum Index'

flagged_indices = set()  # none flagged
n = len(results)

# ── Parameter extraction helpers ───────────────────────────────────────────
def get_peak(r, pk, param, err_param=None):
    val = r['parameters'][pk].get(param, np.nan)
    err = r['parameters'][pk].get(err_param, np.nan) if err_param else np.nan
    # Cap unrealistically huge errors (fitting artefacts) at NaN
    if err_param and abs(err) > 1e4:
        err = np.nan
    return val, err

peaks = ['peak_1', 'peak_2', 'peak_3']
peak_labels = ['Peak 1', 'Peak 2', 'Peak 3']
peak_colors = ['#1f77b4', '#ff7f0e', '#2ca02c']

# Centers
centers     = np.array([[r['parameters'][pk]['center']       for pk in peaks] for r in results])
center_errs = np.array([[r['parameters'][pk]['center_err']   for pk in peaks] for r in results])

# Amplitudes (area)
amps     = np.array([[r['parameters'][pk]['amplitude_area'] for pk in peaks] for r in results])
amp_errs = np.array([[r['parameters'][pk]['amplitude_err']  for pk in peaks] for r in results])

# FWHMs
fwhms     = np.array([[r['parameters'][pk]['fwhm']      for pk in peaks] for r in results])
fwhm_errs = np.array([[r['parameters'][pk]['fwhm_err']  for pk in peaks] for r in results])

# eta_effective
etas = np.array([[r['parameters'][pk]['eta_effective'] for pk in peaks] for r in results])

# Baseline slope
a1      = np.array([r['parameters']['baseline']['a1']     for r in results])
a1_errs = np.array([r['parameters']['baseline']['a1_err'] for r in results])

# R-squared
r2 = np.array([r['fit_quality']['r_squared'] for r in results])

# ── Linear regression helper ────────────────────────────────────────────────
def linreg_line(xv, yv):
    """Return (x_dense, y_pred) for a linear fit, ignoring NaNs."""
    mask = np.isfinite(xv) & np.isfinite(yv)
    if mask.sum() < 2:
        return None, None
    coeffs = np.polyfit(xv[mask], yv[mask], 1)
    xd = np.linspace(xv[mask].min(), xv[mask].max(), 200)
    return xd, np.polyval(coeffs, xd)

def add_trend(ax, xv, yv, color, lw=1.2):
    xd, yd = linreg_line(xv, yv)
    if xd is not None:
        ax.plot(xd, yd, '--', color=color, linewidth=lw, alpha=0.7)

# ── Flag marker overlay helper ──────────────────────────────────────────────
def mark_flagged(ax, x_vals, y_vals, flagged_set):
    for i, xi in enumerate(x_vals):
        if results[i]['index'] in flagged_set:
            ax.plot(xi, y_vals[i], marker='x', color='red',
                    markersize=14, markeredgewidth=2.5,
                    linestyle='none', zorder=10, label='Flagged' if i == 0 else '')

# ── Figure layout ───────────────────────────────────────────────────────────
fig = plt.figure(figsize=(16, 11))
gs  = GridSpec(2, 3, figure=fig, hspace=0.42, wspace=0.35)

ax1 = fig.add_subplot(gs[0, 0])  # Peak centers
ax2 = fig.add_subplot(gs[0, 1])  # Peak amplitudes
ax3 = fig.add_subplot(gs[0, 2])  # Peak FWHMs
ax4 = fig.add_subplot(gs[1, 0])  # eta_effective
ax5 = fig.add_subplot(gs[1, 1])  # Baseline slope a1
ax6 = fig.add_subplot(gs[1, 2])  # R-squared

axes = [ax1, ax2, ax3, ax4, ax5, ax6]

# Shared style
for ax in axes:
    ax.tick_params(labelsize=9)
    ax.set_xlabel(x_label, fontsize=9)
    for spine in ['top', 'right']:
        ax.spines[spine].set_visible(False)

marker_style = dict(marker='o', markersize=6, linewidth=1.5, capsize=4,
                    elinewidth=1.2, capthick=1.2)

# ── Panel 1 — Peak Centers ──────────────────────────────────────────────────
for j, (pk, lbl, col) in enumerate(zip(peaks, peak_labels, peak_colors)):
    yv = centers[:, j]
    ev = center_errs[:, j]
    ax1.errorbar(x_vals, yv, yerr=ev, label=lbl, color=col, **marker_style)
    add_trend(ax1, x_vals, yv, col)
    mark_flagged(ax1, x_vals, yv, flagged_indices)

ax1.set_title('Peak Centers', fontsize=11, fontweight='bold')
ax1.set_ylabel('Center (x units)', fontsize=9)
ax1.legend(fontsize=8, framealpha=0.6)

# ── Panel 2 — Peak Amplitudes ───────────────────────────────────────────────
for j, (pk, lbl, col) in enumerate(zip(peaks, peak_labels, peak_colors)):
    yv = amps[:, j]
    ev = amp_errs[:, j]
    ax2.errorbar(x_vals, yv, yerr=ev, label=lbl, color=col, **marker_style)
    add_trend(ax2, x_vals, yv, col)
    mark_flagged(ax2, x_vals, yv, flagged_indices)

ax2.set_title('Peak Amplitudes (Area)', fontsize=11, fontweight='bold')
ax2.set_ylabel('Amplitude (area)', fontsize=9)
ax2.legend(fontsize=8, framealpha=0.6)

# ── Panel 3 — Peak FWHMs ────────────────────────────────────────────────────
for j, (pk, lbl, col) in enumerate(zip(peaks, peak_labels, peak_colors)):
    yv = fwhms[:, j]
    ev = fwhm_errs[:, j]
    ax3.errorbar(x_vals, yv, yerr=ev, label=lbl, color=col, **marker_style)
    add_trend(ax3, x_vals, yv, col)
    mark_flagged(ax3, x_vals, yv, flagged_indices)

ax3.set_title('Peak FWHM', fontsize=11, fontweight='bold')
ax3.set_ylabel('FWHM (x units)', fontsize=9)
ax3.legend(fontsize=8, framealpha=0.6)

# ── Panel 4 — eta_effective (Lorentzian fraction) ───────────────────────────
for j, (pk, lbl, col) in enumerate(zip(peaks, peak_labels, peak_colors)):
    yv = etas[:, j]
    ax4.plot(x_vals, yv, marker='o', markersize=6, linewidth=1.5,
             label=lbl, color=col)
    add_trend(ax4, x_vals, yv, col)
    mark_flagged(ax4, x_vals, yv, flagged_indices)

ax4.set_title('Effective Lorentzian Fraction (η)', fontsize=11, fontweight='bold')
ax4.set_ylabel('η (0 = Gaussian, 1 = Lorentzian)', fontsize=9)
ax4.set_ylim(0, 1.1)
ax4.legend(fontsize=8, framealpha=0.6)

# ── Panel 5 — Baseline Slope a1 ─────────────────────────────────────────────
ax5.errorbar(x_vals, a1, yerr=a1_errs, color='#9467bd',
             **{**marker_style, 'label': 'a1 (slope)'})
add_trend(ax5, x_vals, a1, '#9467bd')
for i in range(n):
    if results[i]['index'] in flagged_indices:
        ax5.plot(x_vals[i], a1[i], marker='x', color='red',
                 markersize=14, markeredgewidth=2.5, linestyle='none', zorder=10)

ax5.set_title('Baseline Slope (a₁)', fontsize=11, fontweight='bold')
ax5.set_ylabel('a₁', fontsize=9)
ax5.legend(fontsize=8, framealpha=0.6)

# ── Panel 6 — R-squared ─────────────────────────────────────────────────────
ax6.plot(x_vals, r2, marker='s', markersize=7, linewidth=1.8,
         color='#d62728', label='R²')
add_trend(ax6, x_vals, r2, '#d62728')
for i in range(n):
    if results[i]['index'] in flagged_indices:
        ax6.plot(x_vals[i], r2[i], marker='x', color='red',
                 markersize=14, markeredgewidth=2.5, linestyle='none', zorder=10)

# Annotate R² values
for xi, ri in zip(x_vals, r2):
    ax6.annotate(f'{ri:.6f}', xy=(xi, ri),
                 xytext=(0, 8), textcoords='offset points',
                 ha='center', fontsize=7.5, color='#d62728')

ax6.set_title('Fit Quality (R²)', fontsize=11, fontweight='bold')
ax6.set_ylabel('R²', fontsize=9)
y_lo = max(0, r2.min() - 5e-5)
y_hi = min(1, r2.max() + 5e-5)
ax6.set_ylim(y_lo, y_hi)
ax6.legend(fontsize=8, framealpha=0.6)

# ── Global title ────────────────────────────────────────────────────────────
fig.suptitle('Fitted Parameter Trends Across Series\n'
             '(dashed lines = linear regression; error bars = 1σ fit uncertainty)',
             fontsize=13, fontweight='bold', y=1.01)

# ── Legend for flagged marker (if any) ──────────────────────────────────────
if flagged_indices:
    from matplotlib.lines import Line2D
    flag_handle = Line2D([0], [0], marker='x', color='red', linestyle='none',
                         markersize=12, markeredgewidth=2.5, label='Flagged')
    fig.legend(handles=[flag_handle], loc='upper right',
               fontsize=9, framealpha=0.8)

plt.savefig('parameter_trends.png', dpi=150, bbox_inches='tight')
plt.close('all')
