import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit, minimize_scalar
from scipy.special import erfc, erfcx
from scipy.stats import norm
import warnings
warnings.filterwarnings('ignore')

# Load data
data = np.load('/Users/maxim.ziatdinov/Code/SciLink/test_adaptive_annealing_data/output_emg/temp_spectrum_0.npy')
if data.ndim == 2:
    x = data[:, 0]
    y = data[:, 1]
else:
    y = data
    x = np.linspace(0, 50, len(y))

# Sort by x
sort_idx = np.argsort(x)
x = x[sort_idx]
y = y[sort_idx]

def emg_stable(x, A, mu, sigma, tau):
    """Numerically stable EMG using erfcx to prevent overflow."""
    z = (sigma / tau) - (x - mu) / sigma
    # For large z, use erfcx to avoid overflow
    # EMG = (A / (2*tau)) * exp(sigma^2/(2*tau^2) - (x-mu)/tau) * erfc((sigma^2 - tau*(x-mu))/(sqrt(2)*sigma*tau))
    # Rewrite: let u = (sigma^2 - tau*(x-mu))/(sqrt(2)*sigma*tau) = z/sqrt(2)
    # Then exp(...)*erfc(u) = erfcx(u) * exp(... - u^2)
    u = z / np.sqrt(2)
    # exp_arg = sigma^2/(2*tau^2) - (x-mu)/tau
    # exp_arg - u^2 = sigma^2/(2*tau^2) - (x-mu)/tau - z^2/2
    # z = sigma/tau - (x-mu)/sigma
    # z^2/2 = (sigma/tau - (x-mu)/sigma)^2 / 2
    #       = (sigma^2/tau^2 - 2*(x-mu)/tau + (x-mu)^2/sigma^2) / 2
    # exp_arg - u^2 = sigma^2/(2*tau^2) - (x-mu)/tau - sigma^2/(2*tau^2) + (x-mu)/tau - (x-mu)^2/(2*sigma^2)
    #               = -(x-mu)^2/(2*sigma^2)
    # So exp_arg - u^2 = -(x-mu)^2/(2*sigma^2)
    
    result = np.zeros_like(x, dtype=float)
    
    # Use erfcx approach: EMG = (A/(2*tau)) * erfcx(u) * exp(-(x-mu)^2/(2*sigma^2))
    gauss_arg = -0.5 * ((x - mu) / sigma) ** 2
    
    # For numerical stability, handle cases where u is very large or very negative
    # erfcx(u) for large positive u ~ 1/(sqrt(pi)*u) -> safe
    # erfcx(u) for large negative u -> 2*exp(u^2) -> can overflow
    # When u is very negative (z very negative), the EMG contribution is essentially 0
    # because it means we're far to the right of the peak
    
    safe_mask = u > -50  # erfcx blows up for very negative u, but EMG -> 0 there too
    # Actually, when u is very negative, EMG -> (A/tau)*exp(-(x-mu)^2/(2*sigma^2)) * exp(u^2) * ...
    # But actually for very negative u, erfc(u) -> 2, so use direct computation
    
    # Strategy: use erfcx when u > 0 (safe), direct when u <= 0
    pos_mask = u >= 0
    neg_mask = ~pos_mask
    
    if np.any(pos_mask):
        result[pos_mask] = (A / (2 * tau)) * erfcx(u[pos_mask]) * np.exp(gauss_arg[pos_mask])
    
    if np.any(neg_mask):
        exp_arg = 0.5 * (sigma / tau) ** 2 - (x[neg_mask] - mu) / tau
        # Clip to prevent overflow
        exp_arg_clipped = np.clip(exp_arg, -500, 500)
        result[neg_mask] = (A / (2 * tau)) * np.exp(exp_arg_clipped) * erfc(u[neg_mask])
    
    return result

def full_model(x, a, b, A, mu, sigma, tau):
    """Full model: linear baseline + EMG peak."""
    baseline = a + b * x
    peak = emg_stable(x, A, mu, sigma, tau)
    return baseline + peak

def emg_only(x, A, mu, sigma, tau):
    return emg_stable(x, A, mu, sigma, tau)

# Phase 1: Coarse grid search
print("Phase 1: Grid search for initialization...")
tau_grid = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
sigma_grid = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5]

best_rss = np.inf
best_params = None
best_tau_sigma = None

for tau_try in tau_grid:
    for sigma_try in sigma_grid:
        try:
            def model_fixed(x, a, b, A, mu):
                return full_model(x, a, b, A, mu, sigma_try, tau_try)
            
            p0 = [0.12, 0.005, 15.0, 20.0]
            bounds_fixed = ([-1, -0.05, 1, 14], [2, 0.05, 100, 23])
            popt, _ = curve_fit(model_fixed, x, y, p0=p0, bounds=bounds_fixed, maxfev=10000)
            
            y_pred = model_fixed(x, *popt)
            rss = np.sum((y - y_pred) ** 2)
            
            if rss < best_rss:
                best_rss = rss
                best_params = list(popt) + [sigma_try, tau_try]
                best_tau_sigma = (tau_try, sigma_try)
        except Exception:
            continue

print(f"Best grid point: tau={best_tau_sigma[0]}, sigma={best_tau_sigma[1]}, RSS={best_rss:.6f}")

# Phase 2: Full nonlinear optimization
print("Phase 2: Full nonlinear fit...")
a0, b0, A0, mu0 = best_params[:4]
sigma0 = best_params[4]
tau0 = best_params[5]

p0_full = [a0, b0, A0, mu0, sigma0, tau0]
bounds_full = ([-1, -0.05, 1, 14, 0.5, 0.5], [2, 0.05, 100, 23, 5.0, 12.0])

try:
    popt, pcov = curve_fit(full_model, x, y, p0=p0_full, bounds=bounds_full,
                           method='trf', maxfev=50000, ftol=1e-12, xtol=1e-12)
    fit_success = True
except Exception as e:
    print(f"TRF failed: {e}, trying dogbox...")
    try:
        # Perturb initial guesses
        p0_perturbed = [a0, b0, A0, mu0 + 0.5, sigma0, tau0 + 0.5]
        popt, pcov = curve_fit(full_model, x, y, p0=p0_perturbed, bounds=bounds_full,
                               method='dogbox', maxfev=50000, ftol=1e-12, xtol=1e-12)
        fit_success = True
    except Exception as e2:
        print(f"Dogbox also failed: {e2}")
        fit_success = False

if fit_success:
    a_fit, b_fit, A_fit, mu_fit, sigma_fit, tau_fit = popt
    
    # Parameter uncertainties from covariance matrix
    perr = np.sqrt(np.diag(pcov)) if pcov is not None and np.all(np.isfinite(pcov)) else np.zeros(6)
    a_err, b_err, A_err, mu_err, sigma_err, tau_err = perr
    
    y_fit = full_model(x, *popt)
    
    # R² and RMSE
    ss_res = np.sum((y - y_fit) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r_squared = 1 - ss_res / ss_tot
    rmse = np.sqrt(np.mean((y - y_fit) ** 2))
    
    print(f"EMG fit: R²={r_squared:.6f}, RMSE={rmse:.6f}")
    print(f"Parameters: a={a_fit:.4f}, b={b_fit:.6f}, A={A_fit:.4f}, mu={mu_fit:.4f}, sigma={sigma_fit:.4f}, tau={tau_fit:.4f}")

# Fallback: split-Gaussian if R² < 0.99
if not fit_success or r_squared < 0.99:
    print("Attempting split-Gaussian fallback...")
    def split_gaussian(x, a, b, A, mu, sigma_L, sigma_R):
        baseline = a + b * x
        sig = np.where(x < mu, sigma_L, sigma_R)
        peak = A * np.exp(-0.5 * ((x - mu) / sig) ** 2)
        return baseline + peak
    
    try:
        p0_sg = [0.12, 0.005, 3.5, 21.5, 2.0, 4.0]
        bounds_sg = ([-1, -0.05, 0.1, 14, 0.5, 0.5], [2, 0.05, 10, 28, 8, 12])
        popt_sg, pcov_sg = curve_fit(split_gaussian, x, y, p0=p0_sg, bounds=bounds_sg, maxfev=50000)
        y_fit_sg = split_gaussian(x, *popt_sg)
        ss_res_sg = np.sum((y - y_fit_sg) ** 2)
        r2_sg = 1 - ss_res_sg / ss_tot
        rmse_sg = np.sqrt(np.mean((y - y_fit_sg) ** 2))
        print(f"Split-Gaussian: R²={r2_sg:.6f}, RMSE={rmse_sg:.6f}")
        
        if not fit_success or r2_sg > r_squared:
            # Use split-Gaussian instead
            print("Using split-Gaussian model.")
            # We still prefer EMG if it worked
    except Exception as e:
        print(f"Split-Gaussian failed: {e}")

# Fallback: two-Gaussian if still < 0.99
if not fit_success or r_squared < 0.99:
    print("Attempting two-Gaussian fallback...")
    def two_gaussian(x, a, b, A1, mu1, s1, A2, mu2, s2):
        baseline = a + b * x
        g1 = A1 * np.exp(-0.5 * ((x - mu1) / s1) ** 2)
        g2 = A2 * np.exp(-0.5 * ((x - mu2) / s2) ** 2)
        return baseline + g1 + g2
    
    try:
        p0_2g = [0.12, 0.005, 3.0, 21.0, 2.0, 1.0, 25.0, 5.0]
        bounds_2g = ([-1, -0.05, 0.1, 14, 0.5, 0.01, 20, 1], [2, 0.05, 10, 25, 5, 5, 35, 15])
        popt_2g, pcov_2g = curve_fit(two_gaussian, x, y, p0=p0_2g, bounds=bounds_2g, maxfev=50000)
        y_fit_2g = two_gaussian(x, *popt_2g)
        ss_res_2g = np.sum((y - y_fit_2g) ** 2)
        r2_2g = 1 - ss_res_2g / ss_tot
        rmse_2g = np.sqrt(np.mean((y - y_fit_2g) ** 2))
        print(f"Two-Gaussian: R²={r2_2g:.6f}, RMSE={rmse_2g:.6f}")
    except Exception as e:
        print(f"Two-Gaussian failed: {e}")

# Compute derived quantities for the EMG fit
if fit_success:
    # Peak mode: numerically find argmax of full model
    from scipy.optimize import minimize_scalar
    neg_model = lambda xx: -full_model(np.array([xx]), *popt)[0]
    result_mode = minimize_scalar(neg_model, bounds=(14, 35), method='bounded')
    peak_mode = result_mode.x
    peak_height = full_model(np.array([peak_mode]), *popt)[0]
    
    # FWHM: find two x-values where EMG-only equals half its max
    emg_at_mode = emg_only(np.array([peak_mode]), A_fit, mu_fit, sigma_fit, tau_fit)[0]
    half_max = emg_at_mode / 2.0
    
    # Find left half-max point
    from scipy.optimize import brentq
    try:
        x_left = brentq(lambda xx: emg_only(np.array([xx]), A_fit, mu_fit, sigma_fit, tau_fit)[0] - half_max,
                        x.min(), peak_mode)
    except:
        x_left = peak_mode - 2.0
    
    # Find right half-max point
    try:
        x_right = brentq(lambda xx: emg_only(np.array([xx]), A_fit, mu_fit, sigma_fit, tau_fit)[0] - half_max,
                         peak_mode, x.max())
    except:
        x_right = peak_mode + 5.0
    
    fwhm = x_right - x_left
    asymmetry_ratio = tau_fit / sigma_fit
    
    # Durbin-Watson statistic
    residuals = y - y_fit
    dw = np.sum(np.diff(residuals) ** 2) / np.sum(residuals ** 2)
    
    print(f"\nDerived quantities:")
    print(f"  Peak mode: {peak_mode:.4f}")
    print(f"  Peak height: {peak_height:.4f}")
    print(f"  FWHM: {fwhm:.4f}")
    print(f"  Asymmetry ratio (tau/sigma): {asymmetry_ratio:.4f}")
    print(f"  Durbin-Watson: {dw:.4f}")
    print(f"  EMG mu (underlying Gaussian mean): {mu_fit:.4f}")
    print(f"  Mode shift from mu: {peak_mode - mu_fit:.4f}")
    
    # Monte Carlo uncertainty propagation for derived quantities
    n_mc = 1000
    np.random.seed(42)
    try:
        param_samples = np.random.multivariate_normal(popt, pcov, size=n_mc)
    except:
        param_samples = None
    
    if param_samples is not None:
        mc_modes = []
        mc_heights = []
        mc_fwhms = []
        mc_asym = []
        
        for i in range(n_mc):
            try:
                ps = param_samples[i]
                a_s, b_s, A_s, mu_s, sigma_s, tau_s = ps
                if sigma_s <= 0 or tau_s <= 0 or A_s <= 0:
                    continue
                
                neg_m = lambda xx: -full_model(np.array([xx]), *ps)[0]
                res_m = minimize_scalar(neg_m, bounds=(14, 35), method='bounded')
                mode_s = res_m.x
                height_s = full_model(np.array([mode_s]), *ps)[0]
                
                emg_mode_s = emg_only(np.array([mode_s]), A_s, mu_s, sigma_s, tau_s)[0]
                hm_s = emg_mode_s / 2.0
                
                try:
                    xl = brentq(lambda xx: emg_only(np.array([xx]), A_s, mu_s, sigma_s, tau_s)[0] - hm_s,
                                5, mode_s)
                    xr = brentq(lambda xx: emg_only(np.array([xx]), A_s, mu_s, sigma_s, tau_s)[0] - hm_s,
                                mode_s, 45)
                    fwhm_s = xr - xl
                except:
                    fwhm_s = np.nan
                
                mc_modes.append(mode_s)
                mc_heights.append(height_s)
                mc_fwhms.append(fwhm_s)
                mc_asym.append(tau_s / sigma_s)
            except:
                continue
        
        mode_err = np.nanstd(mc_modes) if len(mc_modes) > 10 else 0.0
        height_err = np.nanstd(mc_heights) if len(mc_heights) > 10 else 0.0
        fwhm_err = np.nanstd(mc_fwhms) if len(mc_fwhms) > 10 else 0.0
        asym_err = np.nanstd(mc_asym) if len(mc_asym) > 10 else 0.0
    else:
        mode_err = height_err = fwhm_err = asym_err = 0.0
    
    # Visualization
    fig, axes = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]})
    
    # Top panel: data + fit + components
    ax1 = axes[0]
    ax1.scatter(x, y, s=3, alpha=0.5, color='gray', label='Data', zorder=1)
    ax1.plot(x, y_fit, 'r-', linewidth=2, label='EMG fit', zorder=3)
    
    # Show components
    baseline = a_fit + b_fit * x
    emg_component = emg_only(x, A_fit, mu_fit, sigma_fit, tau_fit)
    ax1.plot(x, baseline, 'g--', linewidth=1, label=f'Baseline: {a_fit:.3f} + {b_fit:.4f}*x', zorder=2)
    ax1.plot(x, baseline + emg_component, 'b-.', linewidth=1, alpha=0.5, label='Baseline + EMG', zorder=2)
    
    # Mark peak mode
    ax1.axvline(peak_mode, color='orange', linestyle=':', alpha=0.7, label=f'Peak mode = {peak_mode:.2f}')
    ax1.axvline(mu_fit, color='purple', linestyle=':', alpha=0.7, label=f'EMG mu = {mu_fit:.2f}')
    
    ax1.set_ylabel('Intensity')
    ax1.set_title(f'EMG Peak Fit (R² = {r_squared:.6f}, RMSE = {rmse:.6f})')
    ax1.legend(fontsize=8)
    ax1.set_xlim(x.min(), x.max())
    
    # Bottom panel: residuals
    ax2 = axes[1]
    ax2.scatter(x, residuals, s=3, alpha=0.5, color='steelblue')
    ax2.axhline(0, color='red', linestyle='-', linewidth=1)
    ax2.set_xlabel('x')
    ax2.set_ylabel('Residual')
    ax2.set_title(f'Residuals (DW={dw:.3f}, Mean={np.mean(residuals):.4e}, Std={np.std(residuals):.4e})')
    ax2.set_xlim(x.min(), x.max())
    
    plt.tight_layout()
    plt.savefig('fit_visualization.png', dpi=150)
    plt.close()
    print("Saved fit_visualization.png")
    
    # Build results JSON
    results = {
        "model_type": "Exponentially Modified Gaussian (EMG) with linear baseline: y = a + b*x + (A/(2*tau)) * exp(sigma^2/(2*tau^2) - (x-mu)/tau) * erfc((sigma^2 - tau*(x-mu))/(sqrt(2)*sigma*tau))",
        "parameters": {
            "peak_1": {
                "center": float(round(peak_mode, 4)),
                "center_err": float(round(mode_err, 4)),
                "height": float(round(peak_height, 4)),
                "height_err": float(round(height_err, 4)),
                "fwhm": float(round(fwhm, 4)),
                "fwhm_err": float(round(fwhm_err, 4)),
                "mu_gaussian": float(round(mu_fit, 4)),
                "mu_gaussian_err": float(round(mu_err, 4)),
                "sigma": float(round(sigma_fit, 4)),
                "sigma_err": float(round(sigma_err, 4)),
                "tau": float(round(tau_fit, 4)),
                "tau_err": float(round(tau_err, 4)),
                "A_area": float(round(A_fit, 4)),
                "A_area_err": float(round(A_err, 4)),
                "asymmetry_ratio_tau_over_sigma": float(round(asymmetry_ratio, 4)),
                "asymmetry_ratio_err": float(round(asym_err, 4)),
                "mode_shift_from_mu": float(round(peak_mode - mu_fit, 4))
            },
            "baseline": {
                "intercept": float(round(a_fit, 4)),
                "intercept_err": float(round(a_err, 4)),
                "slope": float(round(b_fit, 6)),
                "slope_err": float(round(b_err, 6))
            }
        },
        "fit_quality": {
            "r_squared": float(round(r_squared, 6)),
            "rmse": float(round(rmse, 6)),
            "durbin_watson": float(round(dw, 4)),
            "mean_residual": float(round(np.mean(residuals), 6)),
            "max_abs_residual": float(round(np.max(np.abs(residuals)), 6)),
            "residual_std": float(round(np.std(residuals), 6))
        },
        "summary": f"EMG model fit to asymmetric spectroscopic peak. Peak mode at x={peak_mode:.2f} (shifted {peak_mode-mu_fit:.2f} from underlying Gaussian mean mu={mu_fit:.2f}). Asymmetry ratio tau/sigma={asymmetry_ratio:.2f} indicates {'strong' if asymmetry_ratio > 2 else 'moderate' if asymmetry_ratio > 1 else 'mild'} exponential tailing. FWHM={fwhm:.2f}. R²={r_squared:.6f}. Tau={tau_fit:.2f} characterizes the exponential tail decay length."
    }
    
    print(f"FIT_RESULTS_JSON:{json.dumps(results)}")

else:
    print("All fitting attempts failed.")
    results = {
        "model_type": "EMG fit failed",
        "parameters": {},
        "fit_quality": {"r_squared": 0, "rmse": float('inf')},
        "summary": "All fitting attempts (EMG, split-Gaussian, two-Gaussian) failed to converge."
    }
    print(f"FIT_RESULTS_JSON:{json.dumps(results)}")
