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', {})
temperatures = np.array(series_metadata.get('values', [i for i in range(len(results))]))
var_name = series_metadata.get('variable', 'index')
var_unit = series_metadata.get('unit', '')
xlabel = f"{var_name.capitalize()} ({var_unit})" if var_unit else var_name.capitalize()

# Flagged spectra indices
flagged_indices = set()  # None flagged in this series

# Storage for peak near 70 (present in all spectra)
peak70_center = []
peak70_center_err = []
peak70_amp = []
peak70_amp_err = []
peak70_fwhm = []
peak70_fwhm_err = []

# Storage for low-frequency region
peak_low_centers = []  # list of lists (1 or 2 centers per spectrum)
peak_low_amps = []
peak_low_fwhms = []

# For simplified trend
eff_center = []
eff_center_err = []
total_low_amp = []
total_low_amp_err = []
eff_fwhm = []

baseline_slope = []
baseline_slope_err = []
baseline_intercept = []
baseline_intercept_err = []
r_squared = []
rmse_vals = []

def get_peak_fwhm_and_err(peak_params):
    """Extract FWHM and its error from peak parameters, handling different peak types."""
    if 'fwhm' in peak_params:
        return peak_params['fwhm'], peak_params.get('fwhm_err', 0)
    elif 'gamma' in peak_params:
        # Voigt/Lorentzian parameterized with gamma
        sigma = peak_params.get('sigma', 0)
        gamma = peak_params['gamma']
        gamma_err = peak_params.get('gamma_err', 0)
        # Approximate FWHM for Voigt; if sigma is small, FWHM ~ 2*gamma
        fwhm = 2 * gamma
        fwhm_err = 2 * gamma_err
        return fwhm, fwhm_err
    elif 'width' in peak_params:
        # Lorentzian with width parameter
        w = peak_params['width']
        w_err = peak_params.get('width_err', 0)
        return 2 * w, 2 * w_err
    elif 'sigma' in peak_params:
        # Gaussian
        sigma = peak_params['sigma']
        sigma_err = peak_params.get('sigma_err', 0)
        fwhm = 2.3548 * sigma
        fwhm_err = 2.3548 * sigma_err
        return fwhm, fwhm_err
    else:
        return 0, 0

for res in results:
    params = res['parameters']
    fq = res['fit_quality']

    r_squared.append(fq['r_squared'])
    rmse_vals.append(fq['rmse'])

    # Baseline
    bl = params.get('baseline', {})
    baseline_slope.append(bl.get('slope', 0))
    baseline_slope_err.append(bl.get('slope_err', 0))
    baseline_intercept.append(bl.get('intercept', 0))
    baseline_intercept_err.append(bl.get('intercept_err', 0))

    # Dynamically discover peak keys (anything that's not 'baseline')
    peak_keys = [k for k in params if k != 'baseline']

    # Collect all peaks with their center values
    peaks_info = []
    for pk in peak_keys:
        p = params[pk]
        if isinstance(p, dict) and 'center' in p:
            fwhm_val, fwhm_err_val = get_peak_fwhm_and_err(p)
            peaks_info.append({
                'key': pk,
                'center': p['center'],
                'center_err': p.get('center_err', 0),
                'amplitude': p.get('amplitude', 0),
                'amplitude_err': p.get('amplitude_err', 0),
                'fwhm': fwhm_val,
                'fwhm_err': fwhm_err_val,
            })

    # Classify peaks: high-freq (~70) vs low-freq
    # Use threshold of 55 to separate
    threshold = 55
    high_peaks = [p for p in peaks_info if p['center'] > threshold]
    low_peaks = [p for p in peaks_info if p['center'] <= threshold]

    # High-freq peak (pick the one closest to 70 if multiple)
    if high_peaks:
        hp = min(high_peaks, key=lambda p: abs(p['center'] - 70))
        peak70_center.append(hp['center'])
        peak70_center_err.append(hp['center_err'])
        peak70_amp.append(hp['amplitude'])
        peak70_amp_err.append(hp['amplitude_err'])
        peak70_fwhm.append(hp['fwhm'])
        peak70_fwhm_err.append(hp['fwhm_err'])
    else:
        # No high-freq peak found; use NaN
        peak70_center.append(np.nan)
        peak70_center_err.append(0)
        peak70_amp.append(np.nan)
        peak70_amp_err.append(0)
        peak70_fwhm.append(np.nan)
        peak70_fwhm_err.append(0)

    # Low-freq peaks
    if len(low_peaks) == 0:
        eff_center.append(np.nan)
        eff_center_err.append(0)
        total_low_amp.append(np.nan)
        total_low_amp_err.append(0)
        eff_fwhm.append(np.nan)
        peak_low_centers.append([])
        peak_low_amps.append([])
        peak_low_fwhms.append([])
    elif len(low_peaks) == 1:
        lp = low_peaks[0]
        eff_center.append(lp['center'])
        eff_center_err.append(lp['center_err'])
        total_low_amp.append(lp['amplitude'])
        total_low_amp_err.append(lp['amplitude_err'])
        eff_fwhm.append(lp['fwhm'])
        peak_low_centers.append([lp['center']])
        peak_low_amps.append([lp['amplitude']])
        peak_low_fwhms.append([lp['fwhm']])
    else:
        # Multiple low-freq peaks
        centers = [p['center'] for p in low_peaks]
        amps = [p['amplitude'] for p in low_peaks]
        fwhms = [p['fwhm'] for p in low_peaks]
        center_errs = [p['center_err'] for p in low_peaks]
        amp_errs = [p['amplitude_err'] for p in low_peaks]

        total_a = sum(amps)
        if total_a > 0:
            wt_center = sum(a * c for a, c in zip(amps, centers)) / total_a
        else:
            wt_center = np.mean(centers)
        eff_center.append(wt_center)
        eff_center_err.append(np.sqrt(sum(e**2 for e in center_errs)) / len(center_errs))
        total_low_amp.append(total_a)
        total_low_amp_err.append(np.sqrt(sum(e**2 for e in amp_errs)))
        # Effective width: span of peaks plus average fwhm
        eff_fwhm.append(max(centers) - min(centers) + np.mean(fwhms))
        peak_low_centers.append(centers)
        peak_low_amps.append(amps)
        peak_low_fwhms.append(fwhms)

# Convert to arrays
temps = np.array(temperatures)
peak70_center = np.array(peak70_center)
peak70_center_err = np.array(peak70_center_err)
peak70_amp = np.array(peak70_amp)
peak70_amp_err = np.array(peak70_amp_err)
peak70_fwhm = np.array(peak70_fwhm)
peak70_fwhm_err = np.array(peak70_fwhm_err)
eff_center = np.array(eff_center)
eff_center_err = np.array(eff_center_err)
total_low_amp = np.array(total_low_amp)
total_low_amp_err = np.array(total_low_amp_err)
eff_fwhm = np.array(eff_fwhm)
baseline_slope = np.array(baseline_slope)
baseline_slope_err = np.array(baseline_slope_err)
r_squared = np.array(r_squared)
rmse_vals = np.array(rmse_vals)

# ---- Create dashboard figure ----
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
fig.suptitle('Parameter Trends vs Temperature', fontsize=16, fontweight='bold', y=0.98)

def add_trendline(ax, x, y, color='gray'):
    """Add linear regression trendline."""
    mask = np.isfinite(x) & np.isfinite(y)
    xm, ym = x[mask], y[mask]
    if len(xm) >= 3:
        slope_t, intercept_t, r_val, p_val, std_err = stats.linregress(xm, ym)
        x_fit = np.linspace(xm.min(), xm.max(), 100)
        y_fit = slope_t * x_fit + intercept_t
        ax.plot(x_fit, y_fit, '--', color=color, alpha=0.5, linewidth=1,
                label=f'trend: {slope_t:.2e}/K')

def mark_flagged(ax, x, y, flagged_set):
    """Mark flagged points with red X."""
    for i in flagged_set:
        if i < len(x) and np.isfinite(y[i]):
            ax.plot(x[i], y[i], 'X', color='red', markersize=14, markeredgewidth=2,
                    zorder=10, label='flagged' if i == min(flagged_set) else '')

# ---- Panel 1: Low-frequency peak centers ----
ax = axes[0, 0]
for i, T in enumerate(temps):
    centers = peak_low_centers[i]
    amps = peak_low_amps[i]
    for j, (c, a) in enumerate(zip(centers, amps)):
        if len(centers) == 1:
            ax.plot(T, c, 'o', color='#2c7bb6', markersize=8, zorder=5)
        else:
            color_split = '#d7191c' if j == 0 else '#fdae61'
            ax.plot(T, c, 's', color=color_split, markersize=8, zorder=5)

mask_eff = np.isfinite(eff_center)
ax.errorbar(temps[mask_eff], eff_center[mask_eff], yerr=eff_center_err[mask_eff],
            fmt='D--', color='#7b3294', markersize=6, capsize=3, alpha=0.7, label='Weighted center')
ax.axvline(x=425, color='gray', linestyle=':', alpha=0.5, label='Split onset')
ax.set_ylabel('Peak Center', fontsize=11)
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('Low-Freq Peak Centers', fontsize=12, fontweight='bold')
ax.legend(fontsize=8, loc='best')
mark_flagged(ax, temps, eff_center, flagged_indices)
ax.grid(True, alpha=0.3)

# ---- Panel 2: Peak near 70 - center ----
ax = axes[0, 1]
mask70 = np.isfinite(peak70_center)
ax.errorbar(temps[mask70], peak70_center[mask70], yerr=peak70_center_err[mask70],
            fmt='o-', color='#2c7bb6', markersize=7, capsize=3, linewidth=1.5, label='Peak ~70 center')
add_trendline(ax, temps, peak70_center, color='gray')
ax.set_ylabel('Peak Center', fontsize=11)
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('High-Freq Peak Center (~70)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
mark_flagged(ax, temps, peak70_center, flagged_indices)
ax.grid(True, alpha=0.3)

# ---- Panel 3: Amplitudes ----
ax = axes[0, 2]
mask_la = np.isfinite(total_low_amp)
ax.errorbar(temps[mask_la], total_low_amp[mask_la], yerr=total_low_amp_err[mask_la],
            fmt='s-', color='#fdae61', markersize=7, capsize=3, linewidth=1.5, label='Low-freq total amp')
ax.errorbar(temps[mask70], peak70_amp[mask70], yerr=peak70_amp_err[mask70],
            fmt='o-', color='#2c7bb6', markersize=7, capsize=3, linewidth=1.5, label='Peak ~70 amp')
ax.axvline(x=425, color='gray', linestyle=':', alpha=0.5)
add_trendline(ax, temps, peak70_amp, color='#2c7bb6')
add_trendline(ax, temps, total_low_amp, color='#fdae61')
ax.set_ylabel('Amplitude', fontsize=11)
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('Peak Amplitudes', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
mark_flagged(ax, temps, total_low_amp, flagged_indices)
ax.grid(True, alpha=0.3)

# ---- Panel 4: FWHM / Effective Width ----
ax = axes[1, 0]
mask_ef = np.isfinite(eff_fwhm)
ax.plot(temps[mask_ef], eff_fwhm[mask_ef], 's-', color='#fdae61', markersize=7, linewidth=1.5,
        label='Low-freq eff. width')
ax.errorbar(temps[mask70], peak70_fwhm[mask70], yerr=peak70_fwhm_err[mask70],
            fmt='o-', color='#2c7bb6', markersize=7, capsize=3, linewidth=1.5, label='Peak ~70 FWHM')
ax.axvline(x=425, color='gray', linestyle=':', alpha=0.5, label='Split onset')
ax.set_ylabel('Width', fontsize=11)
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('Peak Widths (FWHM / Effective)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
mark_flagged(ax, temps, eff_fwhm, flagged_indices)
ax.grid(True, alpha=0.3)

# ---- Panel 5: Baseline slope ----
ax = axes[1, 1]
ax.errorbar(temps, baseline_slope, yerr=baseline_slope_err, fmt='o-', color='#2c7bb6',
            markersize=7, capsize=3, linewidth=1.5)
add_trendline(ax, temps, baseline_slope, color='gray')
ax.set_ylabel('Baseline Slope', fontsize=11)
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('Baseline Slope', fontsize=12, fontweight='bold')
mark_flagged(ax, temps, baseline_slope, flagged_indices)
ax.grid(True, alpha=0.3)
slope_std = np.std(baseline_slope)
ax.text(0.05, 0.95, f'\u03c3 = {slope_std:.2e}', transform=ax.transAxes,
        fontsize=9, verticalalignment='top',
        bbox=dict(boxstyle='round,pad=0.3', facecolor='wheat', alpha=0.5))

# ---- Panel 6: Fit Quality ----
ax = axes[1, 2]
ax2 = ax.twinx()
ln1 = ax.plot(temps, r_squared, 'o-', color='#2c7bb6', markersize=7, linewidth=1.5,
              label='R\u00b2')
ln2 = ax2.plot(temps, rmse_vals * 1000, 's-', color='#d7191c', markersize=7, linewidth=1.5,
               label='RMSE (\u00d710\u00b3)')
ax.set_ylabel('R\u00b2', fontsize=11, color='#2c7bb6')
ax2.set_ylabel('RMSE (\u00d710\u00b3)', fontsize=11, color='#d7191c')
ax.set_xlabel(xlabel, fontsize=11)
ax.set_title('Fit Quality', fontsize=12, fontweight='bold')
ax.tick_params(axis='y', labelcolor='#2c7bb6')
ax2.tick_params(axis='y', labelcolor='#d7191c')
lns = ln1 + ln2
labs = [l.get_label() for l in lns]
ax.legend(lns, labs, fontsize=8, loc='center left')
mark_flagged(ax, temps, r_squared, flagged_indices)
ax.grid(True, alpha=0.3)

# Add annotation about model transition
fig.text(0.5, 0.005,
         'Note: Model changes from 2-peak (300\u2013400K) to 3-peak (450\u2013500K) due to peak splitting near 35 \u2192 30 + 40',
         ha='center', fontsize=10, style='italic',
         bbox=dict(boxstyle='round,pad=0.4', facecolor='lightyellow', edgecolor='gray', alpha=0.8))

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

print('Dashboard saved as parameter_trends.png')
