import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import os
import glob

# Define the model: single Gaussian on constant baseline
def gaussian_model(x, a, A, mu, sigma):
    return a + A * np.exp(-0.5 * ((x - mu) / sigma) ** 2)

# Find all spectrum files
base_dir = '/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_series/'
spectrum_files = sorted(glob.glob(os.path.join(base_dir, 'temp_spectrum_*.npy')))

if not spectrum_files:
    # Try just the single file
    spectrum_files = [os.path.join(base_dir, "/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_series/temp_spectrum_4.npy")]

all_results = {}
all_fit_data = []

for fpath in spectrum_files:
    fname = os.path.basename(fpath)
    spec_id = fname.replace('temp_spectrum_', '').replace('.npy', '')
    
    # Load data
    data = np.load(fpath)
    
    # Handle different data shapes
    if data.ndim == 2:
        if data.shape[0] == 2:
            x = data[0]
            y = data[1]
        elif data.shape[1] == 2:
            x = data[:, 0]
            y = data[:, 1]
        else:
            x = np.linspace(0, 50, data.shape[-1])
            y = data[0] if data.shape[0] < data.shape[1] else data[:, 0]
    elif data.ndim == 1:
        x = np.linspace(0, 50, len(data))
        y = data
    else:
        raise ValueError(f"Unexpected data shape: {data.shape}")
    
    # Initial guesses
    # Baseline: estimate from edges
    edge_mask = (x < 5) | (x > 40)
    if np.sum(edge_mask) > 0:
        a_init = np.mean(y[edge_mask])
    else:
        a_init = 0.1
    
    # Peak amplitude and center
    idx_max = np.argmax(y)
    A_init = y[idx_max] - a_init
    if A_init < 1.0:
        A_init = 4.0
    mu_init = x[idx_max]
    sigma_init = 2.5
    
    p0 = [a_init, A_init, mu_init, sigma_init]
    bounds_lower = [0.0, 1.0, 5.0, 0.5]
    bounds_upper = [0.5, 6.0, 45.0, 10.0]
    
    # Clamp initial guesses within bounds
    for i in range(len(p0)):
        p0[i] = max(bounds_lower[i], min(bounds_upper[i], p0[i]))
    
    try:
        popt, pcov = curve_fit(gaussian_model, x, y, p0=p0,
                               bounds=(bounds_lower, bounds_upper),
                               method='trf', maxfev=10000)
        perr = np.sqrt(np.diag(pcov))
    except Exception as e:
        print(f"Fitting failed for {fname}: {e}")
        # Try without bounds using LM
        try:
            popt, pcov = curve_fit(gaussian_model, x, y, p0=p0,
                                   method='lm', maxfev=10000)
            perr = np.sqrt(np.diag(pcov))
        except Exception as e2:
            print(f"Fitting also failed with LM for {fname}: {e2}")
            continue
    
    a_fit, A_fit, mu_fit, sigma_fit = popt
    a_err, A_err, mu_err, sigma_err = perr
    
    # Compute FWHM
    fwhm = 2.3548200450309493 * sigma_fit  # 2 * sqrt(2 * ln(2)) * sigma
    fwhm_err = 2.3548200450309493 * sigma_err
    
    # Compute fit quality
    y_fit = gaussian_model(x, *popt)
    residuals = y - y_fit
    ss_res = np.sum(residuals ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0
    rmse = np.sqrt(np.mean(residuals ** 2))
    
    all_fit_data.append({
        'spec_id': spec_id,
        'x': x,
        'y': y,
        'y_fit': y_fit,
        'residuals': residuals,
        'popt': popt,
        'perr': perr,
        'fwhm': fwhm,
        'fwhm_err': fwhm_err,
        'r_squared': r_squared,
        'rmse': rmse
    })
    
    all_results[f'spectrum_{spec_id}'] = {
        'peak_1': {
            'center': float(mu_fit),
            'center_err': float(mu_err),
            'amplitude': float(A_fit),
            'amplitude_err': float(A_err),
            'sigma': float(sigma_fit),
            'sigma_err': float(sigma_err),
            'fwhm': float(fwhm),
            'fwhm_err': float(fwhm_err)
        },
        'baseline': {
            'offset': float(a_fit),
            'offset_err': float(a_err)
        },
        'fit_quality': {
            'r_squared': float(r_squared),
            'rmse': float(rmse)
        }
    }

# Create visualization
num_spectra = len(all_fit_data)
if num_spectra == 0:
    print("No spectra were successfully fitted.")
else:
    fig, axes = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]}, sharex=True)
    
    colors = plt.cm.viridis(np.linspace(0, 1, max(num_spectra, 1)))
    
    for i, fd in enumerate(all_fit_data):
        color = colors[i]
        label_data = f'Spectrum {fd["spec_id"]}'
        label_fit = f'Fit {fd["spec_id"]} (R²={fd["r_squared"]:.5f})'
        
        axes[0].scatter(fd['x'], fd['y'], s=5, alpha=0.5, color=color, label=label_data)
        axes[0].plot(fd['x'], fd['y_fit'], '-', color=color, linewidth=2, label=label_fit)
        
        # Also show baseline and Gaussian component separately
        a_fit = fd['popt'][0]
        axes[0].axhline(y=a_fit, color=color, linestyle=':', alpha=0.3, linewidth=1)
        
        axes[1].plot(fd['x'], fd['residuals'], '.', color=color, markersize=3, alpha=0.5,
                     label=f'Residuals {fd["spec_id"]}')
    
    axes[0].set_ylabel('Intensity', fontsize=12)
    axes[0].set_title('Spectroscopic Data: Single Gaussian + Constant Baseline Fit', fontsize=14)
    axes[0].legend(fontsize=8, ncol=2, loc='upper right')
    axes[0].grid(True, alpha=0.3)
    
    axes[1].axhline(y=0, color='black', linestyle='-', linewidth=0.5)
    axes[1].set_xlabel('X (channel / energy / wavelength)', fontsize=12)
    axes[1].set_ylabel('Residuals', fontsize=12)
    axes[1].legend(fontsize=8, ncol=2)
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('spectrum_0004_fit.png', dpi=150, bbox_inches='tight')
    plt.close()
    print("Saved spectrum_0004_fit.png")

# Prepare JSON output
# If only one spectrum, simplify the output
if num_spectra == 1:
    fd = all_fit_data[0]
    sid = f'spectrum_{fd["spec_id"]}'
    res = all_results[sid]
    
    results = {
        "model_type": "Single Gaussian peak on constant (flat) baseline: y = a + A * exp(-0.5 * ((x - mu) / sigma)^2)",
        "parameters": {
            "peak_1": res['peak_1'],
            "baseline": res['baseline']
        },
        "fit_quality": res['fit_quality'],
        "summary": (
            f"Single Gaussian fit to spectroscopic data. "
            f"Peak center at {res['peak_1']['center']:.3f} +/- {res['peak_1']['center_err']:.4f}, "
            f"FWHM = {res['peak_1']['fwhm']:.3f} +/- {res['peak_1']['fwhm_err']:.4f}, "
            f"amplitude = {res['peak_1']['amplitude']:.3f}, "
            f"baseline offset = {res['baseline']['offset']:.4f}. "
            f"R² = {res['fit_quality']['r_squared']:.6f}, RMSE = {res['fit_quality']['rmse']:.6f}. "
            f"Excellent fit quality."
        )
    }
else:
    # Multiple spectra
    parameters = {}
    avg_r2 = []
    avg_rmse = []
    centers = []
    for sid, res in all_results.items():
        parameters[sid] = {
            'peak_1': res['peak_1'],
            'baseline': res['baseline']
        }
        avg_r2.append(res['fit_quality']['r_squared'])
        avg_rmse.append(res['fit_quality']['rmse'])
        centers.append(res['peak_1']['center'])
    
    results = {
        "model_type": "Single Gaussian peak on constant (flat) baseline: y = a + A * exp(-0.5 * ((x - mu) / sigma)^2)",
        "parameters": parameters,
        "fit_quality": {
            "r_squared": float(np.mean(avg_r2)),
            "rmse": float(np.mean(avg_rmse)),
            "r_squared_per_spectrum": {sid: all_results[sid]['fit_quality']['r_squared'] for sid in all_results},
            "rmse_per_spectrum": {sid: all_results[sid]['fit_quality']['rmse'] for sid in all_results}
        },
        "summary": (
            f"Fitted {num_spectra} spectra with single Gaussian + constant baseline model. "
            f"Peak centers range from {min(centers):.3f} to {max(centers):.3f}. "
            f"Mean R² = {np.mean(avg_r2):.6f}, Mean RMSE = {np.mean(avg_rmse):.6f}."
        )
    }

print(f"FIT_RESULTS_JSON:{json.dumps(results)}")
