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
positions = np.array(series_metadata.get('values', list(range(len(results)))))
var_name = series_metadata.get('variable', 'index')
var_unit = series_metadata.get('unit', '')
xlabel = f"{var_name} ({var_unit})" if var_unit else var_name

# Flagged spectra (none in this case)
flagged_indices = set()

# Extract parameters
n = len(results)
peak1_center = np.full(n, np.nan)
peak1_center_err = np.full(n, np.nan)
peak1_amplitude = np.full(n, np.nan)
peak1_amplitude_err = np.full(n, np.nan)
peak1_sigma = np.full(n, np.nan)
peak1_sigma_err = np.full(n, np.nan)
baseline_offset = np.full(n, np.nan)
baseline_offset_err = np.full(n, np.nan)
r_squared = np.full(n, np.nan)
num_peaks = np.full(n, np.nan)

# Also track peak_2 if present
peak2_center = np.full(n, np.nan)
peak2_center_err = np.full(n, np.nan)
peak2_amplitude = np.full(n, np.nan)
peak2_amplitude_err = np.full(n, np.nan)

for i, res in enumerate(results):
    params = res['parameters']
    fq = res['fit_quality']
    
    # Peak 1
    if 'peak_1' in params:
        p1 = params['peak_1']
        peak1_center[i] = p1.get('center', np.nan)
        peak1_center_err[i] = p1.get('center_err', 0)
        peak1_amplitude[i] = p1.get('amplitude', np.nan)
        peak1_amplitude_err[i] = p1.get('amplitude_err', 0)
        peak1_sigma[i] = p1.get('sigma', np.nan)
        peak1_sigma_err[i] = p1.get('sigma_err', 0)
    
    # Peak 2
    if 'peak_2' in params:
        p2 = params['peak_2']
        peak2_center[i] = p2.get('center', np.nan)
        peak2_center_err[i] = p2.get('center_err', 0)
        peak2_amplitude[i] = p2.get('amplitude', np.nan)
        peak2_amplitude_err[i] = p2.get('amplitude_err', 0)
    
    # Baseline
    if 'baseline' in params:
        bl = params['baseline']
        baseline_offset[i] = bl.get('offset', np.nan)
        baseline_offset_err[i] = bl.get('offset_err', 0)
    
    # Fit quality
    r_squared[i] = fq.get('r_squared', np.nan)
    
    # Count peaks
    peak_keys = [k for k in params.keys() if k.startswith('peak_')]
    num_peaks[i] = len(peak_keys)

# Identify flagged and normal masks
flagged_mask = np.array([i in flagged_indices for i in range(n)])
normal_mask = ~flagged_mask

# Helper: add trend line for normal (non-flagged) points
def add_trendline(ax, x, y, mask):
    valid = mask & ~np.isnan(y)
    if np.sum(valid) >= 2:
        slope, intercept, r_val, p_val, std_err = stats.linregress(x[valid], y[valid])
        x_fit = np.linspace(x[valid].min(), x[valid].max(), 100)
        y_fit = slope * x_fit + intercept
        ax.plot(x_fit, y_fit, 'k--', alpha=0.4, linewidth=1,
                label=f'trend: slope={slope:.4f}')
        ax.legend(fontsize=7, loc='best')

# Helper: plot flagged points
def mark_flagged(ax, x, y, mask):
    if np.any(mask):
        ax.scatter(x[mask], y[mask], marker='X', s=120, c='red',
                   zorder=10, label='flagged', edgecolors='darkred', linewidths=1.2)

# Create figure: 2x3 layout
fig, axes = plt.subplots(2, 3, figsize=(16, 9))
fig.suptitle('Parameter Trends Across Series', fontsize=16, fontweight='bold', y=0.98)

# ---- Panel 1: Peak 1 Center ----
ax = axes[0, 0]
ax.errorbar(positions[normal_mask], peak1_center[normal_mask],
            yerr=peak1_center_err[normal_mask], fmt='o-', color='#1f77b4',
            capsize=4, markersize=6, label='Peak 1')
# Also plot peak_2 centers if any
has_p2 = ~np.isnan(peak2_center)
if np.any(has_p2):
    ax.errorbar(positions[has_p2], peak2_center[has_p2],
                yerr=peak2_center_err[has_p2], fmt='s', color='#ff7f0e',
                capsize=4, markersize=8, label='Peak 2')
mark_flagged(ax, positions, peak1_center, flagged_mask)
add_trendline(ax, positions, peak1_center, normal_mask)
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('Center position', fontsize=10)
ax.set_title('Peak Center(s)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 2: Peak 1 Amplitude ----
ax = axes[0, 1]
ax.errorbar(positions[normal_mask], peak1_amplitude[normal_mask],
            yerr=peak1_amplitude_err[normal_mask], fmt='o-', color='#2ca02c',
            capsize=4, markersize=6, label='Peak 1')
if np.any(has_p2):
    ax.errorbar(positions[has_p2], peak2_amplitude[has_p2],
                yerr=peak2_amplitude_err[has_p2], fmt='s', color='#ff7f0e',
                capsize=4, markersize=8, label='Peak 2')
mark_flagged(ax, positions, peak1_amplitude, flagged_mask)
add_trendline(ax, positions, peak1_amplitude, normal_mask)
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('Amplitude', fontsize=10)
ax.set_title('Peak Amplitude(s)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 3: Peak 1 Sigma (width) ----
ax = axes[0, 2]
ax.errorbar(positions[normal_mask], peak1_sigma[normal_mask],
            yerr=peak1_sigma_err[normal_mask], fmt='o-', color='#d62728',
            capsize=4, markersize=6, label='Peak 1')
mark_flagged(ax, positions, peak1_sigma, flagged_mask)
add_trendline(ax, positions, peak1_sigma, normal_mask)
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('Sigma', fontsize=10)
ax.set_title('Peak 1 Width (\u03c3)', fontsize=12, fontweight='bold')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# ---- Panel 4: Baseline Offset ----
ax = axes[1, 0]
ax.errorbar(positions[normal_mask], baseline_offset[normal_mask],
            yerr=baseline_offset_err[normal_mask], fmt='o-', color='#9467bd',
            capsize=4, markersize=6)
mark_flagged(ax, positions, baseline_offset, flagged_mask)
add_trendline(ax, positions, baseline_offset, normal_mask)
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('Baseline offset', fontsize=10)
ax.set_title('Baseline Offset', fontsize=12, fontweight='bold')
ax.grid(True, alpha=0.3)

# ---- Panel 5: R-squared ----
ax = axes[1, 1]
ax.plot(positions[normal_mask], r_squared[normal_mask], 'o-', color='#8c564b',
        markersize=6)
mark_flagged(ax, positions, r_squared, flagged_mask)
# Shade quality regions
ax.axhspan(0.999, 1.0, color='green', alpha=0.07, label='Excellent')
ax.axhspan(0.99, 0.999, color='yellow', alpha=0.07, label='Good')
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('R\u00b2', fontsize=10)
ax.set_title('Fit Quality (R\u00b2)', fontsize=12, fontweight='bold')
# Set ylim to zoom around the data
r2_min = np.nanmin(r_squared)
r2_range = 1.0 - r2_min
ax.set_ylim(r2_min - 0.2 * r2_range, 1.0 + 0.05 * r2_range)
ax.legend(fontsize=7, loc='lower right')
ax.grid(True, alpha=0.3)

# ---- Panel 6: Number of peaks / Model complexity ----
ax = axes[1, 2]
colors_bar = ['#1f77b4' if not flagged_mask[i] else 'red' for i in range(n)]
ax.bar(positions, num_peaks, color=colors_bar, alpha=0.7, edgecolor='black', width=0.6)
ax.set_xlabel(xlabel, fontsize=10)
ax.set_ylabel('Number of peaks', fontsize=10)
ax.set_title('Model Complexity', fontsize=12, fontweight='bold')
ax.set_yticks([1, 2, 3])
ax.set_ylim(0, 3.5)
ax.grid(True, alpha=0.3, axis='y')

# Add annotations for the two-peak spectrum
for i in range(n):
    if num_peaks[i] > 1:
        ax.annotate(f'{int(num_peaks[i])} peaks', (positions[i], num_peaks[i]),
                    textcoords='offset points', xytext=(0, 10),
                    ha='center', fontsize=9, fontweight='bold', color='#ff7f0e')

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 as parameter_trends.png')
