import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

# ── 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.get('values'):
    x_vals = np.array(series_metadata['values'], dtype=float)
    x_label = f"{series_metadata.get('variable', 'Series variable')} ({series_metadata.get('unit', '')})".strip()
else:
    x_vals = np.array([r['index'] for r in results], dtype=float)
    x_label = 'Spectrum index'

flagged = set()  # no spectra flagged

# ── Parameter extraction ────────────────────────────────────────────────────
def get(r, *keys, default=np.nan):
    obj = r
    for k in keys:
        if isinstance(obj, dict) and k in obj:
            obj = obj[k]
        else:
            return default
    return obj if obj is not None else default

p1_cen  = np.array([get(r,'parameters','peak_1','center')        for r in results])
p2_cen  = np.array([get(r,'parameters','peak_2','center')        for r in results])
p3_cen  = np.array([get(r,'parameters','peak_3','center')        for r in results])
p1_cen_e= np.array([get(r,'parameters','peak_1','center_err')    for r in results])
p2_cen_e= np.array([get(r,'parameters','peak_2','center_err')    for r in results])
p3_cen_e= np.array([get(r,'parameters','peak_3','center_err')    for r in results])

p1_amp  = np.array([get(r,'parameters','peak_1','amplitude')     for r in results])
p2_amp  = np.array([get(r,'parameters','peak_2','amplitude')     for r in results])
p3_amp  = np.array([get(r,'parameters','peak_3','amplitude')     for r in results])
p1_amp_e= np.array([get(r,'parameters','peak_1','amplitude_err') for r in results])
p2_amp_e= np.array([get(r,'parameters','peak_2','amplitude_err') for r in results])
p3_amp_e= np.array([get(r,'parameters','peak_3','amplitude_err') for r in results])

p1_fwhm = np.array([get(r,'parameters','peak_1','fwhm')          for r in results])
p2_fwhm = np.array([get(r,'parameters','peak_2','fwhm')          for r in results])
p3_fwhm = np.array([get(r,'parameters','peak_3','fwhm')          for r in results])
p1_fwhm_e=np.array([get(r,'parameters','peak_1','gamma_err')*2   for r in results])
p2_fwhm_e=np.array([get(r,'parameters','peak_2','gamma_err')*2   for r in results])
p3_fwhm_e=np.array([get(r,'parameters','peak_3','gamma_err')*2   for r in results])

r2      = np.array([get(r,'fit_quality','r_squared')             for r in results])
r2_pk   = np.array([get(r,'fit_quality','r_squared_peak_regions')for r in results])

# ── Helper: linear trend ────────────────────────────────────────────────────
def trend_line(x, y):
    mask = np.isfinite(y)
    if mask.sum() < 2:
        return None, None
    coeffs = np.polyfit(x[mask], y[mask], 1)
    return coeffs, np.poly1d(coeffs)(x)

# ── Color palette ───────────────────────────────────────────────────────────
C = ['#2196F3', '#FF9800', '#4CAF50']   # blue, orange, green (peaks 1-3)
TRENDC = ['#0D47A1', '#E65100', '#1B5E20']

# ── Figure layout ───────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(15, 9))
fig.suptitle('Fitted Parameter Trends Across Series', fontsize=15, fontweight='bold', y=1.01)

ax_cen, ax_amp, ax_fwhm = axes[0]
ax_area, ax_base, ax_r2  = axes[1]

marker_kw  = dict(marker='o', markersize=6, linewidth=1.5, capsize=4)
flag_kw    = dict(marker='x', color='red', markersize=14, linewidth=2.5,
                  linestyle='none', zorder=10, label='Flagged')

# ── 1. Peak centres ─────────────────────────────────────────────────────────
for cen, cerr, label, col, tc in zip(
        [p1_cen, p2_cen, p3_cen],
        [p1_cen_e, p2_cen_e, p3_cen_e],
        ['Peak 1', 'Peak 2', 'Peak 3'], C, TRENDC):
    ax_cen.errorbar(x_vals, cen, yerr=cerr, label=label, color=col, **marker_kw)
    _, ty = trend_line(x_vals, cen)
    if ty is not None:
        ax_cen.plot(x_vals, ty, '--', color=tc, linewidth=1, alpha=0.7)
for fi in flagged:
    ax_cen.plot(x_vals[fi], p1_cen[fi], **flag_kw)
ax_cen.set_title('Peak Centres', fontweight='bold')
ax_cen.set_xlabel(x_label)
ax_cen.set_ylabel('Centre position (a.u.)')
ax_cen.legend(fontsize=8)
ax_cen.grid(True, alpha=0.3)

# ── 2. Peak amplitudes ───────────────────────────────────────────────────────
for amp, aerr, label, col, tc in zip(
        [p1_amp, p2_amp, p3_amp],
        [p1_amp_e, p2_amp_e, p3_amp_e],
        ['Peak 1', 'Peak 2', 'Peak 3'], C, TRENDC):
    ax_amp.errorbar(x_vals, amp, yerr=aerr, label=label, color=col, **marker_kw)
    _, ty = trend_line(x_vals, amp)
    if ty is not None:
        ax_amp.plot(x_vals, ty, '--', color=tc, linewidth=1, alpha=0.7)
for fi in flagged:
    ax_amp.plot(x_vals[fi], p1_amp[fi], **flag_kw)
ax_amp.set_title('Peak Amplitudes', fontweight='bold')
ax_amp.set_xlabel(x_label)
ax_amp.set_ylabel('Amplitude (a.u.)')
ax_amp.legend(fontsize=8)
ax_amp.grid(True, alpha=0.3)

# ── 3. Peak FWHM ─────────────────────────────────────────────────────────────
for fwhm, ferr, label, col, tc in zip(
        [p1_fwhm, p2_fwhm, p3_fwhm],
        [p1_fwhm_e, p2_fwhm_e, p3_fwhm_e],
        ['Peak 1', 'Peak 2', 'Peak 3'], C, TRENDC):
    ax_fwhm.errorbar(x_vals, fwhm, yerr=ferr, label=label, color=col, **marker_kw)
    _, ty = trend_line(x_vals, fwhm)
    if ty is not None:
        ax_fwhm.plot(x_vals, ty, '--', color=tc, linewidth=1, alpha=0.7)
for fi in flagged:
    ax_fwhm.plot(x_vals[fi], p1_fwhm[fi], **flag_kw)
ax_fwhm.set_title('Peak FWHM', fontweight='bold')
ax_fwhm.set_xlabel(x_label)
ax_fwhm.set_ylabel('FWHM (a.u.)')
ax_fwhm.legend(fontsize=8)
ax_fwhm.grid(True, alpha=0.3)

# ── 4. Integrated areas ───────────────────────────────────────────────────────
p1_area = np.array([get(r,'parameters','peak_1','integrated_area') for r in results])
p2_area = np.array([get(r,'parameters','peak_2','integrated_area') for r in results])
p3_area = np.array([get(r,'parameters','peak_3','integrated_area') for r in results])

for area, label, col, tc in zip(
        [p1_area, p2_area, p3_area],
        ['Peak 1', 'Peak 2', 'Peak 3'], C, TRENDC):
    ax_area.plot(x_vals, area, marker='o', markersize=6, linewidth=1.5,
                 label=label, color=col)
    _, ty = trend_line(x_vals, area)
    if ty is not None:
        ax_area.plot(x_vals, ty, '--', color=tc, linewidth=1, alpha=0.7)
for fi in flagged:
    ax_area.plot(x_vals[fi], p1_area[fi], **flag_kw)
ax_area.set_title('Integrated Peak Areas', fontweight='bold')
ax_area.set_xlabel(x_label)
ax_area.set_ylabel('Integrated area (a.u.)')
ax_area.legend(fontsize=8)
ax_area.grid(True, alpha=0.3)

# ── 5. Baseline parameters ────────────────────────────────────────────────────
base_c   = np.array([get(r,'parameters','baseline','c')       for r in results])
base_c_e = np.array([get(r,'parameters','baseline','c_err')   for r in results])
base_tau = np.array([get(r,'parameters','baseline','tau')     for r in results])

ax_base2 = ax_base.twinx()
ax_base.errorbar(x_vals, base_c, yerr=base_c_e, color='#9C27B0', marker='s',
                 markersize=6, linewidth=1.5, capsize=4, label='c (saturation level)')
_, ty = trend_line(x_vals, base_c)
if ty is not None:
    ax_base.plot(x_vals, ty, '--', color='#4A148C', linewidth=1, alpha=0.7)
ax_base2.plot(x_vals, base_tau, color='#F44336', marker='^', markersize=6,
              linewidth=1.5, linestyle='-.', label='tau (decay const)')
for fi in flagged:
    ax_base.plot(x_vals[fi], base_c[fi], **flag_kw)
ax_base.set_title('Baseline Parameters', fontweight='bold')
ax_base.set_xlabel(x_label)
ax_base.set_ylabel('c  (saturation level)', color='#9C27B0')
ax_base2.set_ylabel('tau  (decay constant)', color='#F44336')
lines1, labs1 = ax_base.get_legend_handles_labels()
lines2, labs2 = ax_base2.get_legend_handles_labels()
ax_base.legend(lines1 + lines2, labs1 + labs2, fontsize=8)
ax_base.grid(True, alpha=0.3)

# ── 6. R² quality ────────────────────────────────────────────────────────────
ax_r2.plot(x_vals, r2,    color='#009688', marker='D', markersize=7,
           linewidth=1.5, label='R² (global)')
ax_r2.plot(x_vals, r2_pk, color='#FF5722', marker='s', markersize=7,
           linewidth=1.5, linestyle='--', label='R² (peak regions)')
for fi in flagged:
    ax_r2.plot(x_vals[fi], r2[fi], **flag_kw)
ax_r2.set_title('Fit Quality (R²)', fontweight='bold')
ax_r2.set_xlabel(x_label)
ax_r2.set_ylabel('R²')
ax_r2.set_ylim(bottom=max(0, min(np.nanmin(r2), np.nanmin(r2_pk)) - 0.002))
ax_r2.legend(fontsize=8)
ax_r2.grid(True, alpha=0.3)

# ── Flagged legend entry (only if any flagged) ────────────────────────────────
if flagged:
    flag_handle = Line2D([0], [0], marker='x', color='red', markersize=10,
                         linewidth=0, label='Flagged spectrum')
    fig.legend(handles=[flag_handle], loc='upper right', fontsize=9,
               bbox_to_anchor=(1.0, 1.0))

# ── Trend-line legend note ────────────────────────────────────────────────────
fig.text(0.5, -0.01,
         'Dashed lines = linear regression trend  |  Error bars = 1σ fit uncertainty',
         ha='center', fontsize=9, color='grey')

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