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', {})

# Extract series variable
temps = np.array(series_metadata.get('values', list(range(len(results)))))
temp_label = series_metadata.get('variable', 'index').capitalize()
temp_unit = series_metadata.get('unit', '')
x_label = f'{temp_label} ({temp_unit})' if temp_unit else temp_label

# Identify flagged spectra
flagged_indices = set()
for r in results:
    if r.get('flagged', False):
        flagged_indices.add(r['index'])

# Extract parameters
peak1_center = np.array([r['parameters']['peak_1']['center'] for r in results])
peak1_center_err = np.array([r['parameters']['peak_1']['center_err'] for r in results])
peak2_center = np.array([r['parameters']['peak_2']['center'] for r in results])
peak2_center_err = np.array([r['parameters']['peak_2']['center_err'] for r in results])

peak1_amp = np.array([r['parameters']['peak_1']['amplitude'] for r in results])
peak1_amp_err = np.array([r['parameters']['peak_1']['amplitude_err'] for r in results])
peak2_amp = np.array([r['parameters']['peak_2']['amplitude'] for r in results])
peak2_amp_err = np.array([r['parameters']['peak_2']['amplitude_err'] for r in results])

peak1_fwhm = np.array([r['parameters']['peak_1']['fwhm'] for r in results])
peak1_fwhm_err = np.array([r['parameters']['peak_1']['fwhm_err'] for r in results])
peak2_fwhm = np.array([r['parameters']['peak_2']['fwhm'] for r in results])
peak2_fwhm_err = np.array([r['parameters']['peak_2']['fwhm_err'] for r in results])

baseline_slope = np.array([r['parameters']['baseline']['slope_m'] for r in results])
baseline_slope_err = np.array([r['parameters']['baseline']['slope_m_err'] for r in results])

r_squared = np.array([r['fit_quality']['r_squared'] for r in results])
rmse = np.array([r['fit_quality']['rmse'] for r in results])

# Masks for flagged / unflagged
flagged_mask = np.array([i in flagged_indices for i in range(len(results))])
unflagged_mask = ~flagged_mask

# Helper: add linear trend line
def add_trend(ax, x, y, color='gray'):
    slope, intercept, r_val, p_val, std_err = stats.linregress(x, y)
    x_fit = np.linspace(x.min(), x.max(), 100)
    y_fit = slope * x_fit + intercept
    ax.plot(x_fit, y_fit, '--', color=color, alpha=0.6, linewidth=1.2,
            label=f'slope={slope:.4g}, R\u00b2={r_val**2:.4f}')
    ax.legend(fontsize=7, loc='best')

# Color scheme
c1 = '#2176AE'  # blue for peak 1
c2 = '#D7263D'  # red for peak 2
cb = '#57A773'  # green for baseline
cq = '#6B4C9A'  # purple for quality

fig, axes = plt.subplots(2, 3, figsize=(15, 9))
fig.suptitle('Parameter Trends vs Temperature', fontsize=16, fontweight='bold', y=0.98)

# ---- Panel 1: Peak Centers ----
ax = axes[0, 0]
ax.errorbar(temps[unflagged_mask], peak1_center[unflagged_mask],
            yerr=peak1_center_err[unflagged_mask], fmt='o-', color=c1,
            capsize=4, markersize=6, label='Peak 1 center', linewidth=1.5)
ax.errorbar(temps[unflagged_mask], peak2_center[unflagged_mask],
            yerr=peak2_center_err[unflagged_mask], fmt='s-', color=c2,
            capsize=4, markersize=6, label='Peak 2 center', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], peak1_center[flagged_mask], marker='X', s=120, c='red', zorder=5)
    ax.scatter(temps[flagged_mask], peak2_center[flagged_mask], marker='X', s=120, c='red', zorder=5)
add_trend(ax, temps, peak1_center, color=c1)
add_trend(ax, temps, peak2_center, color=c2)
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('Peak Center (a.u.)', fontsize=10)
ax.set_title('Peak Centers', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 2: Peak Amplitudes ----
ax = axes[0, 1]
ax.errorbar(temps[unflagged_mask], peak1_amp[unflagged_mask],
            yerr=peak1_amp_err[unflagged_mask], fmt='o-', color=c1,
            capsize=4, markersize=6, label='Peak 1 amplitude', linewidth=1.5)
ax.errorbar(temps[unflagged_mask], peak2_amp[unflagged_mask],
            yerr=peak2_amp_err[unflagged_mask], fmt='s-', color=c2,
            capsize=4, markersize=6, label='Peak 2 amplitude', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], peak1_amp[flagged_mask], marker='X', s=120, c='red', zorder=5)
    ax.scatter(temps[flagged_mask], peak2_amp[flagged_mask], marker='X', s=120, c='red', zorder=5)
add_trend(ax, temps, peak1_amp, color=c1)
add_trend(ax, temps, peak2_amp, color=c2)
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('Amplitude (a.u.)', fontsize=10)
ax.set_title('Peak Amplitudes', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 3: Peak FWHM ----
ax = axes[0, 2]
ax.errorbar(temps[unflagged_mask], peak1_fwhm[unflagged_mask],
            yerr=peak1_fwhm_err[unflagged_mask], fmt='o-', color=c1,
            capsize=4, markersize=6, label='Peak 1 FWHM', linewidth=1.5)
ax.errorbar(temps[unflagged_mask], peak2_fwhm[unflagged_mask],
            yerr=peak2_fwhm_err[unflagged_mask], fmt='s-', color=c2,
            capsize=4, markersize=6, label='Peak 2 FWHM', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], peak1_fwhm[flagged_mask], marker='X', s=120, c='red', zorder=5)
    ax.scatter(temps[flagged_mask], peak2_fwhm[flagged_mask], marker='X', s=120, c='red', zorder=5)
add_trend(ax, temps, peak1_fwhm, color=c1)
add_trend(ax, temps, peak2_fwhm, color=c2)
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('FWHM (a.u.)', fontsize=10)
ax.set_title('Peak Widths (FWHM)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 4: Peak Separation ----
ax = axes[1, 0]
separation = peak2_center - peak1_center
sep_err = np.sqrt(peak1_center_err**2 + peak2_center_err**2)
ax.errorbar(temps[unflagged_mask], separation[unflagged_mask],
            yerr=sep_err[unflagged_mask], fmt='D-', color='#E8871E',
            capsize=4, markersize=6, label='Peak separation', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], separation[flagged_mask], marker='X', s=120, c='red', zorder=5)
add_trend(ax, temps, separation, color='#E8871E')
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('Peak Separation (a.u.)', fontsize=10)
ax.set_title('Peak Separation (P2 \u2013 P1)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 5: Baseline Slope ----
ax = axes[1, 1]
ax.errorbar(temps[unflagged_mask], baseline_slope[unflagged_mask],
            yerr=baseline_slope_err[unflagged_mask], fmt='^-', color=cb,
            capsize=4, markersize=6, label='Baseline slope', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], baseline_slope[flagged_mask], marker='X', s=120, c='red', zorder=5)
add_trend(ax, temps, baseline_slope, color=cb)
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('Slope m', fontsize=10)
ax.set_title('Baseline Slope', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)
ax.ticklabel_format(style='sci', axis='y', scilimits=(-3,-3))

# ---- Panel 6: Fit Quality (R² and RMSE) ----
ax = axes[1, 2]
color_r2 = cq
ax.plot(temps[unflagged_mask], r_squared[unflagged_mask], 'o-', color=color_r2,
        markersize=6, label='R\u00b2', linewidth=1.5)
if flagged_mask.any():
    ax.scatter(temps[flagged_mask], r_squared[flagged_mask], marker='X', s=120, c='red', zorder=5)
ax.set_xlabel(x_label, fontsize=10)
ax.set_ylabel('R\u00b2', color=color_r2, fontsize=10)
ax.tick_params(axis='y', labelcolor=color_r2)
ax.set_title('Fit Quality', fontsize=12, fontweight='bold')
ax.grid(True, alpha=0.3)

# Twin axis for RMSE
ax2 = ax.twinx()
color_rmse = '#FF8C00'
ax2.plot(temps[unflagged_mask], rmse[unflagged_mask] * 1000, 's--', color=color_rmse,
         markersize=6, label='RMSE (\u00d710\u207B\u00b3)', linewidth=1.5)
if flagged_mask.any():
    ax2.scatter(temps[flagged_mask], rmse[flagged_mask] * 1000, marker='X', s=120, c='red', zorder=5)
ax2.set_ylabel('RMSE (\u00d710\u207B\u00b3)', color=color_rmse, fontsize=10)
ax2.tick_params(axis='y', labelcolor=color_rmse)

# Combined legend
lines1, labels1 = ax.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax.legend(lines1 + lines2, labels1 + labels2, fontsize=8, loc='center left')

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

print('Dashboard saved: parameter_trends.png')
