#!/usr/bin/env python3
"""
3D+3D THEORY UTILITIES
======================

Common utility functions for data analysis, fitting, and visualization.

Author: Simone Calzighetti & Lucy (Claude AI)
Date: December 2025
Version: 1.0
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit, minimize, brentq
from scipy.interpolate import interp1d, UnivariateSpline
from scipy.integrate import quad, simpson
from scipy.stats import chi2 as chi2_dist
import json
import os

# ============================================================================
# STATISTICAL FUNCTIONS
# ============================================================================

def chi_squared(observed, predicted, errors):
    """
    Compute chi-squared statistic.
    
    χ² = Σ [(O - P) / σ]²
    """
    residuals = (observed - predicted) / errors
    return np.sum(residuals**2)

def reduced_chi_squared(observed, predicted, errors, n_params=0):
    """
    Compute reduced chi-squared.
    
    χ²_red = χ² / (N - n_params)
    """
    chi2 = chi_squared(observed, predicted, errors)
    dof = len(observed) - n_params
    return chi2 / dof

def rms(residuals):
    """Root mean square of residuals."""
    return np.sqrt(np.mean(residuals**2))

def weighted_mean(values, errors):
    """Weighted mean and error."""
    weights = 1 / errors**2
    mean = np.sum(weights * values) / np.sum(weights)
    error = np.sqrt(1 / np.sum(weights))
    return mean, error

def bootstrap_error(data, func, n_bootstrap=1000, seed=None):
    """
    Estimate error using bootstrap resampling.
    
    Parameters:
    -----------
    data : array
        Input data
    func : callable
        Function to compute statistic
    n_bootstrap : int
        Number of bootstrap samples
    
    Returns:
    --------
    mean, std : float
        Bootstrap mean and standard deviation
    """
    if seed is not None:
        np.random.seed(seed)
    
    n = len(data)
    results = []
    
    for _ in range(n_bootstrap):
        sample = np.random.choice(data, size=n, replace=True)
        results.append(func(sample))
    
    return np.mean(results), np.std(results)

def jackknife_error(data, func):
    """
    Estimate error using jackknife resampling.
    """
    n = len(data)
    results = np.zeros(n)
    
    for i in range(n):
        sample = np.delete(data, i)
        results[i] = func(sample)
    
    mean = np.mean(results)
    variance = (n - 1) / n * np.sum((results - mean)**2)
    
    return mean, np.sqrt(variance)

# ============================================================================
# FITTING FUNCTIONS
# ============================================================================

def fit_power_law(x, y, y_err=None):
    """
    Fit power law: y = A * x^n
    
    Returns:
    --------
    A, n, cov : amplitude, power, covariance matrix
    """
    def power_law(x, A, n):
        return A * x**n
    
    # Initial guess from log-log fit
    mask = (x > 0) & (y > 0)
    log_x, log_y = np.log(x[mask]), np.log(y[mask])
    n_init = np.polyfit(log_x, log_y, 1)[0]
    A_init = np.exp(np.mean(log_y - n_init * log_x))
    
    try:
        popt, pcov = curve_fit(power_law, x, y, p0=[A_init, n_init], 
                              sigma=y_err, absolute_sigma=True)
        return popt[0], popt[1], pcov
    except:
        return A_init, n_init, None

def fit_exponential(x, y, y_err=None):
    """
    Fit exponential: y = A * exp(-x/λ)
    
    Returns:
    --------
    A, lambda, cov
    """
    def exponential(x, A, lam):
        return A * np.exp(-x / lam)
    
    # Initial guess
    A_init = y[0]
    lam_init = x[-1] / 3
    
    try:
        popt, pcov = curve_fit(exponential, x, y, p0=[A_init, lam_init],
                              sigma=y_err, absolute_sigma=True)
        return popt[0], popt[1], pcov
    except:
        return A_init, lam_init, None

def fit_gaussian(x, y, y_err=None):
    """
    Fit Gaussian: y = A * exp(-(x-μ)²/(2σ²))
    
    Returns:
    --------
    A, mu, sigma, cov
    """
    def gaussian(x, A, mu, sigma):
        return A * np.exp(-(x - mu)**2 / (2 * sigma**2))
    
    # Initial guess
    A_init = y.max()
    mu_init = x[np.argmax(y)]
    sigma_init = (x.max() - x.min()) / 5
    
    try:
        popt, pcov = curve_fit(gaussian, x, y, p0=[A_init, mu_init, sigma_init],
                              sigma=y_err, absolute_sigma=True)
        return popt[0], popt[1], popt[2], pcov
    except:
        return A_init, mu_init, sigma_init, None

# ============================================================================
# FORM FACTORS (3D+3D specific)
# ============================================================================

def yukawa_form_factor(r, lambda_b):
    """
    Yukawa form factor F(r/λ).
    
    F(x) = 1 - (1 + x) * exp(-x)
    
    F → 0 for r << λ
    F → 1 for r >> λ
    """
    x = r / lambda_b
    return 1 - (1 + x) * np.exp(-x)

def screening_factor(r, r_V):
    """
    Vainshtein screening factor S(r).
    
    S = r³ / (r³ + r_V³)
    
    S → 0 for r << r_V (screened)
    S → 1 for r >> r_V (unscreened)
    """
    return r**3 / (r**3 + r_V**3)

def q_field_enhancement(r, lambda_b, beta):
    """
    Q-field gravitational enhancement factor.
    
    g_eff/g_N = 1 + β² * F(r/λ)
    """
    F = yukawa_form_factor(r, lambda_b)
    return 1 + beta**2 * F

# ============================================================================
# INTERPOLATION AND SMOOTHING
# ============================================================================

def smooth_spline(x, y, s=None, k=3):
    """
    Smooth data using cubic spline.
    
    Parameters:
    -----------
    x, y : arrays
        Data points
    s : float
        Smoothing factor (None = automatic)
    k : int
        Spline degree (default 3)
    
    Returns:
    --------
    interp : callable
        Interpolation function
    """
    spl = UnivariateSpline(x, y, s=s, k=k)
    return spl

def log_interpolate(x, y, kind='linear'):
    """
    Interpolate in log-log space.
    """
    mask = (x > 0) & (y > 0)
    log_interp = interp1d(np.log(x[mask]), np.log(y[mask]), kind=kind,
                         fill_value='extrapolate')
    return lambda xi: np.exp(log_interp(np.log(xi)))

# ============================================================================
# NUMERICAL INTEGRATION
# ============================================================================

def integrate_radial(f, r_min, r_max, n_points=1000):
    """
    Integrate f(r) × 4πr² from r_min to r_max.
    """
    r = np.linspace(r_min, r_max, n_points)
    integrand = 4 * np.pi * r**2 * f(r)
    return simpson(integrand, x=r)

def integrate_spherical_shell(f, r1, r2, n_points=100):
    """
    Integrate f(r) over spherical shell from r1 to r2.
    """
    def integrand(r):
        return 4 * np.pi * r**2 * f(r)
    
    result, _ = quad(integrand, r1, r2)
    return result

# ============================================================================
# DATA I/O
# ============================================================================

def save_results(filename, results, format='json'):
    """
    Save results to file.
    
    Parameters:
    -----------
    filename : str
        Output filename
    results : dict
        Results dictionary
    format : str
        'json' or 'npz'
    """
    if format == 'json':
        # Convert numpy arrays to lists
        def convert(obj):
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, dict):
                return {k: convert(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [convert(i) for i in obj]
            else:
                return obj
        
        with open(filename, 'w') as f:
            json.dump(convert(results), f, indent=2)
    
    elif format == 'npz':
        np.savez(filename, **results)
    
    print(f"✓ Saved: {filename}")

def load_results(filename):
    """
    Load results from file.
    """
    if filename.endswith('.json'):
        with open(filename, 'r') as f:
            return json.load(f)
    elif filename.endswith('.npz'):
        return dict(np.load(filename))
    else:
        raise ValueError(f"Unknown format: {filename}")

# ============================================================================
# PLOTTING UTILITIES
# ============================================================================

def setup_plot_style():
    """Set up publication-quality plot style."""
    plt.rcParams.update({
        'font.size': 12,
        'axes.labelsize': 14,
        'axes.titlesize': 16,
        'xtick.labelsize': 12,
        'ytick.labelsize': 12,
        'legend.fontsize': 11,
        'figure.figsize': (10, 6),
        'figure.dpi': 100,
        'savefig.dpi': 150,
        'lines.linewidth': 1.5,
        'axes.grid': True,
        'grid.alpha': 0.3,
    })

def add_golden_ratio_line(ax, x_range, y0, label=True):
    """Add golden ratio reference line to plot."""
    phi = (1 + np.sqrt(5)) / 2
    ax.axhline(phi, color='gold', linestyle='--', alpha=0.7,
              label=f'φ = {phi:.3f}' if label else None)

def add_breathing_scale_marker(ax, lambda_n, label=None, color='gray'):
    """Add vertical line at breathing scale."""
    ax.axvline(lambda_n, color=color, linestyle=':', alpha=0.5,
              label=label or f'λ = {lambda_n:.1f}')

def create_comparison_plot(x, y_obs, y_pred, y_err=None, 
                          xlabel='x', ylabel='y', title=''):
    """
    Create standard comparison plot with residuals.
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8),
                                   gridspec_kw={'height_ratios': [3, 1]})
    
    # Main plot
    if y_err is not None:
        ax1.errorbar(x, y_obs, yerr=y_err, fmt='ko', markersize=4,
                    capsize=2, label='Observed')
    else:
        ax1.scatter(x, y_obs, c='black', s=20, label='Observed')
    
    ax1.plot(x, y_pred, 'r-', linewidth=2, label='3D+3D')
    
    ax1.set_ylabel(ylabel)
    ax1.set_title(title)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Residuals
    residuals = y_obs - y_pred
    
    if y_err is not None:
        ax2.errorbar(x, residuals, yerr=y_err, fmt='ko', markersize=4, capsize=2)
    else:
        ax2.scatter(x, residuals, c='black', s=20)
    
    ax2.axhline(0, color='r', linestyle='-')
    rms_val = rms(residuals)
    ax2.axhline(rms_val, color='gray', linestyle='--', alpha=0.5)
    ax2.axhline(-rms_val, color='gray', linestyle='--', alpha=0.5)
    
    ax2.set_xlabel(xlabel)
    ax2.set_ylabel('Residual')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig, (ax1, ax2)

# ============================================================================
# VALIDATION FUNCTIONS
# ============================================================================

def check_parameter_free(params_used, params_allowed):
    """
    Verify that only allowed (universal) parameters were used.
    
    Returns True if no free parameters per object.
    """
    for p in params_used:
        if p not in params_allowed:
            print(f"⚠ Warning: Non-universal parameter used: {p}")
            return False
    return True

def compute_significance(observed, predicted, error):
    """
    Compute sigma deviation.
    
    σ = |observed - predicted| / error
    """
    return abs(observed - predicted) / error

def validate_prediction(name, predicted, observed, error, threshold=2.0):
    """
    Validate a prediction against observation.
    
    Returns True if agreement within threshold σ.
    """
    sigma = compute_significance(observed, predicted, error)
    passed = sigma < threshold
    
    status = "✓ PASS" if passed else "✗ FAIL"
    print(f"{name}: pred={predicted:.4f}, obs={observed:.4f}±{error:.4f}, "
          f"σ={sigma:.2f} {status}")
    
    return passed

# ============================================================================
# MAIN
# ============================================================================

if __name__ == '__main__':
    # Quick test of utilities
    print("3D+3D Utilities Test")
    print("="*50)
    
    # Test form factors
    r = np.linspace(0.1, 50, 100)
    lambda_b = 4.30
    
    F = yukawa_form_factor(r, lambda_b)
    
    print(f"\nYukawa form factor at r/λ = 0.1: F = {yukawa_form_factor(0.1*lambda_b, lambda_b):.4f}")
    print(f"Yukawa form factor at r/λ = 1.0: F = {yukawa_form_factor(1.0*lambda_b, lambda_b):.4f}")
    print(f"Yukawa form factor at r/λ = 10:  F = {yukawa_form_factor(10*lambda_b, lambda_b):.4f}")
    
    # Test statistics
    obs = np.array([100, 105, 98, 102, 101])
    pred = np.array([100, 100, 100, 100, 100])
    err = np.array([3, 3, 3, 3, 3])
    
    print(f"\nχ² test:")
    print(f"  χ² = {chi_squared(obs, pred, err):.2f}")
    print(f"  χ²_red = {reduced_chi_squared(obs, pred, err, n_params=0):.2f}")
    print(f"  RMS = {rms(obs - pred):.2f}")
    
    print("\n✓ All utilities working!")
