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

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

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

flagged_indices = set()  # populated from metadata if present
for r in results:
    if r.get('flagged', False):
        flagged_indices.add(r['index'])

# ── Extract parameters ─────────────────────────────────────────────────────────
indices = np.array([r['index'] for r in results], dtype=float)

def get(results, peak, field, err=False):
    key = field + ('_err' if err else '')
    vals = []
    for r in results:
        p = r['parameters'].get(peak, {})
        vals.append(p.get(key, np.nan))
    return np.array(vals, dtype=float)

# Peak centers
p1_center      = get(results, 'peak_1', 'center_mu')
p1_center_err  = get(results, 'peak_1', 'center_mu', err=True)
p2_center      = get(results, 'peak_2', 'center_mu')
p2_center_err  = get(results, 'peak_2', 'center_mu', err=True)
p3_center      = get(results, 'peak_3', 'center_mu')
p3_center_err  = get(results, 'peak_3', 'center_mu', err=True)

# Amplitudes
p1_amp     = get(results, 'peak_1', 'amplitude_A')
p1_amp_err = get(results, 'peak_1', 'amplitude_A', err=True)
p2_amp     = get(results, 'peak_2', 'amplitude_A')
p2_amp_err = get(results, 'peak_2', 'amplitude_A', err=True)
p3_amp     = get(results, 'peak_3', 'amplitude_A')
p3_amp_err = get(results, 'peak_3', 'amplitude_A', err=True)

# FWHMs
p1_fwhm     = get(results, 'peak_1', 'fwhm')
p2_fwhm     = get(results, 'peak_2', 'fwhm')
p3_fwhm     = get(results, 'peak_3', 'fwhm')

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

# ── Colour / style palette ─────────────────────────────────────────────────────
C1, C2, C3 = '#2196F3', '#FF9800', '#4CAF50'   # blue, orange, green
FLAG_COLOR  = 'red'

# ── Helper: linear regression overlay ─────────────────────────────────────────
def add_trendline(ax, x, y, color='grey', lw=1.2):
    mask = np.isfinite(x) & np.isfinite(y)
    if mask.sum() < 2:
        return
    m, b = np.polyfit(x[mask], y[mask], 1)
    xfit = np.linspace(x[mask].min(), x[mask].max(), 200)
    ax.plot(xfit, m * xfit + b, '--', color=color, lw=lw, alpha=0.7,
            label=f'trend (slope={m:.4f})')

# ── Helper: scatter with error bars ───────────────────────────────────────────
def plot_series(ax, x, y, yerr, color, label, marker='o', ms=7):
    # normal points
    normal_mask = np.array([i not in flagged_indices for i in range(len(x))])
    flag_mask   = ~normal_mask

    ax.errorbar(x[normal_mask], y[normal_mask],
                yerr=yerr[normal_mask] if yerr is not None else None,
                fmt=marker, color=color, markersize=ms,
                ecolor=color, elinewidth=1, capsize=3,
                label=label, zorder=3)

    if flag_mask.any():
        ax.scatter(x[flag_mask], y[flag_mask],
                   marker='X', s=160, color=FLAG_COLOR, zorder=5,
                   label='flagged')

# ── Build figure ───────────────────────────────────────────────────────────────
fig = plt.figure(figsize=(16, 12))
fig.patch.set_facecolor('#F7F9FC')

gs = GridSpec(3, 3, figure=fig, hspace=0.48, wspace=0.38)

# Axes layout:
#  [0,0] Peak centres   [0,1] Amplitudes     [0,2] FWHMs
#  [1,0] P1 center zoom [1,1] P2 center zoom [1,2] P3 center zoom
#  [2,0..2] R² (spans full width)
ax_cen  = fig.add_subplot(gs[0, 0])
ax_amp  = fig.add_subplot(gs[0, 1])
ax_fwhm = fig.add_subplot(gs[0, 2])
ax_c1   = fig.add_subplot(gs[1, 0])
ax_c2   = fig.add_subplot(gs[1, 1])
ax_c3   = fig.add_subplot(gs[1, 2])
ax_r2   = fig.add_subplot(gs[2, :])

for ax in [ax_cen, ax_amp, ax_fwhm, ax_c1, ax_c2, ax_c3, ax_r2]:
    ax.set_facecolor('#FFFFFF')
    ax.grid(True, linestyle='--', linewidth=0.5, color='#CCCCCC', alpha=0.8)
    for spine in ax.spines.values():
        spine.set_linewidth(0.8)
        spine.set_color('#AAAAAA')

# ── Panel 1 : All peak centres overlaid ───────────────────────────────────────
plot_series(ax_cen, indices, p1_center, p1_center_err, C1, 'Peak 1')
plot_series(ax_cen, indices, p2_center, p2_center_err, C2, 'Peak 2')
plot_series(ax_cen, indices, p3_center, p3_center_err, C3, 'Peak 3')
add_trendline(ax_cen, indices, p1_center, color=C1)
add_trendline(ax_cen, indices, p2_center, color=C2)
add_trendline(ax_cen, indices, p3_center, color=C3)
ax_cen.set_title('Peak Centres', fontsize=11, fontweight='bold')
ax_cen.set_xlabel('Spectrum Index')
ax_cen.set_ylabel('Center (a.u.)')
ax_cen.legend(fontsize=8, framealpha=0.9)

# ── Panel 2 : Amplitudes ───────────────────────────────────────────────────────
plot_series(ax_amp, indices, p1_amp, p1_amp_err, C1, 'Peak 1')
plot_series(ax_amp, indices, p2_amp, p2_amp_err, C2, 'Peak 2')
plot_series(ax_amp, indices, p3_amp, p3_amp_err, C3, 'Peak 3')
add_trendline(ax_amp, indices, p1_amp, color=C1)
add_trendline(ax_amp, indices, p2_amp, color=C2)
add_trendline(ax_amp, indices, p3_amp, color=C3)
ax_amp.set_title('Peak Amplitudes', fontsize=11, fontweight='bold')
ax_amp.set_xlabel('Spectrum Index')
ax_amp.set_ylabel('Amplitude (a.u.)')
ax_amp.legend(fontsize=8, framealpha=0.9)

# ── Panel 3 : FWHMs ────────────────────────────────────────────────────────────
no_err = np.zeros_like(p1_fwhm)   # FWHM errors not stored separately
plot_series(ax_fwhm, indices, p1_fwhm, no_err, C1, 'Peak 1')
plot_series(ax_fwhm, indices, p2_fwhm, no_err, C2, 'Peak 2')
plot_series(ax_fwhm, indices, p3_fwhm, no_err, C3, 'Peak 3')
add_trendline(ax_fwhm, indices, p1_fwhm, color=C1)
add_trendline(ax_fwhm, indices, p2_fwhm, color=C2)
add_trendline(ax_fwhm, indices, p3_fwhm, color=C3)
ax_fwhm.set_title('Peak FWHMs', fontsize=11, fontweight='bold')
ax_fwhm.set_xlabel('Spectrum Index')
ax_fwhm.set_ylabel('FWHM (a.u.)')
ax_fwhm.legend(fontsize=8, framealpha=0.9)

# ── Panels 4-6 : Individual centre drift (zoomed) ─────────────────────────────
for ax, y, yerr, color, label, title in [
        (ax_c1, p1_center, p1_center_err, C1, 'Peak 1', 'Peak 1 Centre Drift'),
        (ax_c2, p2_center, p2_center_err, C2, 'Peak 2', 'Peak 2 Centre Drift'),
        (ax_c3, p3_center, p3_center_err, C3, 'Peak 3', 'Peak 3 Centre Drift'),
]:
    plot_series(ax, indices, y, yerr, color, label)
    add_trendline(ax, indices, y, color=color)

    # annotate slope
    mask = np.isfinite(y)
    if mask.sum() >= 2:
        m, b = np.polyfit(indices[mask], y[mask], 1)
        ax.annotate(f'slope = {m:+.4f}/idx',
                    xy=(0.05, 0.90), xycoords='axes fraction',
                    fontsize=8, color='#333333',
                    bbox=dict(boxstyle='round,pad=0.3', fc='#F0F0F0', ec='grey', lw=0.6))

    ax.set_title(title, fontsize=11, fontweight='bold')
    ax.set_xlabel('Spectrum Index')
    ax.set_ylabel('Center (a.u.)')

    # tight y-range with a small padding
    ypad = max(0.05, 3 * np.nanstd(y))
    ax.set_ylim(np.nanmean(y) - ypad, np.nanmean(y) + ypad)

# ── Panel 7 : R² ──────────────────────────────────────────────────────────────
normal_mask = np.array([i not in flagged_indices for i in range(len(indices))])
flag_mask   = ~normal_mask

ax_r2.plot(indices[normal_mask], r2[normal_mask],
           'o-', color='#9C27B0', markersize=7, linewidth=1.5,
           label='R²', zorder=3)
if flag_mask.any():
    ax_r2.scatter(indices[flag_mask], r2[flag_mask],
                  marker='X', s=160, color=FLAG_COLOR, zorder=5, label='flagged')

add_trendline(ax_r2, indices, r2, color='#9C27B0')

# annotate each point with its R² value
for idx, val in zip(indices, r2):
    ax_r2.annotate(f'{val:.6f}',
                   xy=(idx, val), xytext=(0, 8),
                   textcoords='offset points', ha='center',
                   fontsize=7.5, color='#555555')

ax_r2.set_title('Fit Quality (R²) Across Series', fontsize=11, fontweight='bold')
ax_r2.set_xlabel('Spectrum Index')
ax_r2.set_ylabel('R²')
y_lo = max(0, np.nanmin(r2) - 0.0002)
y_hi = min(1.0, np.nanmax(r2) + 0.0002)
ax_r2.set_ylim(y_lo, y_hi)
ax_r2.legend(fontsize=9, framealpha=0.9)

# ── Suptitle ───────────────────────────────────────────────────────────────────
fig.suptitle('Fitted Parameter Trends — 3 Voigt Peaks + Exponential Baseline',
             fontsize=13, fontweight='bold', y=1.005, color='#222222')

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