"""
ASYMPTOTIC SAFETY EXPLORER - LPA' Implementation
=================================================

Functional Renormalization Group analysis for Q₂-Q₃ scalar sector
in 3D+3D theory, searching for UV fixed point.

Author: Simone Calzighetti, Lucy (Claude AI)
Date: November 21, 2025
Status: ACTIVE RESEARCH - EDISON MODE!

Usage:
    python asymptotic_safety_lpa.py
    
Output:
    - RG flow plots
    - Fixed point search results
    - Phase diagrams
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.optimize import fsolve
import json
from datetime import datetime

# ============================================================================
# PHYSICAL CONSTANTS
# ============================================================================

# Compactification scales (from Paper I)
L4_kpc = 1.9  # kpc
L5_kpc = 4.5  # kpc

# Convert to eV^-1
kpc_to_eV_inv = 6.19e25  # 1 kpc ≈ 6.19×10^25 eV^-1
L4 = L4_kpc * kpc_to_eV_inv
L5 = L5_kpc * kpc_to_eV_inv

# KK scale (cutoff)
Lambda_KK = 1.0 / L4  # ~ 10^-24 eV

# Initial UV scale for RG flow
k_UV = Lambda_KK

print(f"Physical scales:")
print(f"  L₄ = {L4_kpc} kpc = {L4:.2e} eV⁻¹")
print(f"  L₅ = {L5_kpc} kpc = {L5:.2e} eV⁻¹")
print(f"  Λ_KK = {Lambda_KK:.2e} eV")
print()

# ============================================================================
# REGULATOR FUNCTIONS
# ============================================================================

def litim_regulator(q2, k2):
    """
    Optimized Litim regulator: R_k(q²) = (k² - q²)θ(k² - q²)
    
    Args:
        q2: momentum squared
        k2: RG scale squared
    
    Returns:
        R_k(q²)
    """
    return np.maximum(k2 - q2, 0.0)

def regulator_derivative(q2, k2):
    """
    ∂_t R_k = 2k² θ(k² - q²)
    
    Args:
        q2: momentum squared  
        k2: RG scale squared
        
    Returns:
        ∂_t R_k(q²)
    """
    return np.where(q2 < k2, 2*k2, 0.0)

# ============================================================================
# BETA FUNCTIONS - LOCAL POTENTIAL APPROXIMATION (LPA)
# ============================================================================

def beta_functions_LPA(t, y, return_dict=False):
    """
    RG flow equations in Local Potential Approximation.
    
    State vector: y = [m2_2, m2_3, lam2, lam3, lam23]
    
    Args:
        t: RG time, t = log(k/k_UV)
        y: State vector [m²₂, m²₃, λ₂, λ₃, λ₂₃]
        return_dict: If True, return dict with components
        
    Returns:
        dy/dt: Beta functions
    """
    # Unpack state
    m2_2, m2_3, lam2, lam3, lam23 = y
    
    # Current RG scale
    k = k_UV * np.exp(t)
    k2 = k**2
    
    # Dimensionless ratios
    r2 = m2_2 / k2
    r3 = m2_3 / k2
    
    # Regularization: avoid division by zero
    eps = 1e-10
    denom2 = 1.0 + r2 + eps
    denom3 = 1.0 + r3 + eps
    
    # Beta functions for mass parameters
    # ∂_t m² = (k²/16π²) [1/(1 + m²/k²)]
    beta_m2_2 = (k2 / (16*np.pi**2)) * (1.0/denom2)
    beta_m2_3 = (k2 / (16*np.pi**2)) * (1.0/denom3)
    
    # Beta function for self-coupling λ₂
    # ∂_t λ₂ = (1/16π²) [λ₂/(1+r₂) - 2λ₂²/(1+r₂)²]
    beta_lam2 = (1.0 / (16*np.pi**2)) * (
        lam2/denom2 - 2*lam2**2/(denom2**2)
    )
    
    # Beta function for self-coupling λ₃
    beta_lam3 = (1.0 / (16*np.pi**2)) * (
        lam3/denom3 - 2*lam3**2/(denom3**2)
    )
    
    # Beta function for cross-coupling λ₂₃
    # ∂_t λ₂₃ = (1/16π²) [λ₂₃/(1+r₂) + λ₂₃/(1+r₃) - 4λ₂₃²/((1+r₂)(1+r₃))]
    beta_lam23 = (1.0 / (16*np.pi**2)) * (
        lam23/denom2 + lam23/denom3 - 4*lam23**2/(denom2*denom3)
    )
    
    result = np.array([beta_m2_2, beta_m2_3, beta_lam2, beta_lam3, beta_lam23])
    
    if return_dict:
        return {
            'beta_m2_2': beta_m2_2,
            'beta_m2_3': beta_m2_3,
            'beta_lam2': beta_lam2,
            'beta_lam3': beta_lam3,
            'beta_lam23': beta_lam23,
            'k': k,
            'r2': r2,
            'r3': r3
        }
    
    return result

# ============================================================================
# BETA FUNCTIONS - LPA' (with wave function renormalization)
# ============================================================================

def compute_anomalous_dimension(m2, lam, k2):
    """
    Anomalous dimension η in LPA':
    η = (1/16π²k²) [U'''²/(k² + U'')²]
    
    For U = (m²/2)Q² + (λ/4!)Q⁴:
    U'' = m² + (λ/2)Q²
    U''' = λQ
    
    At Q=0 (for simplicity in LPA'):
    η ≈ 0 (no anomalous dimension at origin)
    
    Args:
        m2: mass squared
        lam: quartic coupling
        k2: RG scale squared
        
    Returns:
        η: anomalous dimension
    """
    # Simplified: evaluate at Q=0
    # U''(0) = m²
    # U'''(0) = 0
    # → η = 0 in LPA' at origin
    
    # For non-zero field expectation, use field-dependent η
    # Here we use Q=0 approximation
    return 0.0

def beta_functions_LPA_prime(t, y, return_dict=False):
    """
    RG flow equations in LPA' (LPA + wave function renormalization).
    
    State vector: y = [m2_2, m2_3, lam2, lam3, lam23, Z2, Z3]
    
    Args:
        t: RG time, t = log(k/k_UV)
        y: State vector [m²₂, m²₃, λ₂, λ₃, λ₂₃, Z₂, Z₃]
        return_dict: If True, return dict with components
        
    Returns:
        dy/dt: Beta functions
    """
    # Unpack state
    m2_2, m2_3, lam2, lam3, lam23, Z2, Z3 = y
    
    # Current RG scale
    k = k_UV * np.exp(t)
    k2 = k**2
    
    # Dimensionless ratios
    r2 = m2_2 / k2
    r3 = m2_3 / k2
    
    # Regularization
    eps = 1e-10
    denom2 = 1.0 + r2 + eps
    denom3 = 1.0 + r3 + eps
    
    # Anomalous dimensions
    eta2 = compute_anomalous_dimension(m2_2, lam2, k2)
    eta3 = compute_anomalous_dimension(m2_3, lam3, k2)
    
    # Beta functions for mass (same as LPA but with eta correction)
    beta_m2_2 = (k2 / (16*np.pi**2)) * (1.0/denom2) - 2*eta2*m2_2
    beta_m2_3 = (k2 / (16*np.pi**2)) * (1.0/denom3) - 2*eta3*m2_3
    
    # Beta functions for couplings (with eta correction)
    beta_lam2 = (1.0 / (16*np.pi**2)) * (
        lam2/denom2 - 2*lam2**2/(denom2**2)
    ) - 4*eta2*lam2
    
    beta_lam3 = (1.0 / (16*np.pi**2)) * (
        lam3/denom3 - 2*lam3**2/(denom3**2)
    ) - 4*eta3*lam3
    
    beta_lam23 = (1.0 / (16*np.pi**2)) * (
        lam23/denom2 + lam23/denom3 - 4*lam23**2/(denom2*denom3)
    ) - 2*(eta2 + eta3)*lam23
    
    # Beta functions for wave function renormalization
    beta_Z2 = eta2 * Z2
    beta_Z3 = eta3 * Z3
    
    result = np.array([
        beta_m2_2, beta_m2_3, beta_lam2, beta_lam3, beta_lam23,
        beta_Z2, beta_Z3
    ])
    
    if return_dict:
        return {
            'beta_m2_2': beta_m2_2,
            'beta_m2_3': beta_m2_3,
            'beta_lam2': beta_lam2,
            'beta_lam3': beta_lam3,
            'beta_lam23': beta_lam23,
            'beta_Z2': beta_Z2,
            'beta_Z3': beta_Z3,
            'k': k,
            'eta2': eta2,
            'eta3': eta3
        }
    
    return result

# ============================================================================
# RG FLOW INTEGRATION
# ============================================================================

def run_RG_flow(initial_conditions, t_max=20, method='LPA', rtol=1e-6):
    """
    Integrate RG flow equations from UV to IR.
    
    Args:
        initial_conditions: Initial values at k_UV
        t_max: Maximum RG time (log scale)
        method: 'LPA' or 'LPA_prime'
        rtol: Relative tolerance for ODE solver
        
    Returns:
        solution object from solve_ivp
    """
    if method == 'LPA':
        beta_func = beta_functions_LPA
        n_vars = 5
    elif method == 'LPA_prime':
        beta_func = beta_functions_LPA_prime
        n_vars = 7
    else:
        raise ValueError(f"Unknown method: {method}")
    
    # Check initial conditions size
    if len(initial_conditions) != n_vars:
        raise ValueError(f"Expected {n_vars} initial conditions for {method}")
    
    print(f"Running RG flow ({method}):")
    print(f"  Initial conditions: {initial_conditions}")
    print(f"  t_max = {t_max} (k_IR/k_UV = {np.exp(-t_max):.2e})")
    
    # Integrate
    sol = solve_ivp(
        beta_func,
        t_span=[0, t_max],
        y0=initial_conditions,
        method='RK45',
        dense_output=True,
        rtol=rtol,
        atol=1e-10
    )
    
    if sol.success:
        print(f"  ✓ Integration successful!")
        print(f"  Final values: {sol.y[:, -1]}")
    else:
        print(f"  ✗ Integration failed: {sol.message}")
    
    return sol

# ============================================================================
# FIXED POINT SEARCH
# ============================================================================

def find_fixed_point(initial_guess, method='LPA'):
    """
    Search for fixed point where all beta functions vanish.
    
    Args:
        initial_guess: Starting point for search
        method: 'LPA' or 'LPA_prime'
        
    Returns:
        fixed_point: Solution if found, None otherwise
        info: Dictionary with search info
    """
    if method == 'LPA':
        beta_func = lambda y: beta_functions_LPA(0, y)
    elif method == 'LPA_prime':
        beta_func = lambda y: beta_functions_LPA_prime(0, y)
    else:
        raise ValueError(f"Unknown method: {method}")
    
    print(f"\nSearching for fixed point ({method}):")
    print(f"  Initial guess: {initial_guess}")
    
    # Use fsolve to find zeros of beta functions
    solution, info, ier, msg = fsolve(
        beta_func,
        initial_guess,
        full_output=True
    )
    
    if ier == 1:
        # Check if solution is actually a fixed point
        residual = beta_func(solution)
        max_residual = np.max(np.abs(residual))
        
        if max_residual < 1e-6:
            print(f"  ✓ Fixed point found!")
            print(f"  Solution: {solution}")
            print(f"  Max residual: {max_residual:.2e}")
            
            return solution, {
                'success': True,
                'residual': residual,
                'max_residual': max_residual,
                'message': msg
            }
        else:
            print(f"  ⚠ Solution found but residual too large: {max_residual:.2e}")
            return None, {
                'success': False,
                'residual': residual,
                'max_residual': max_residual,
                'message': 'Residual too large'
            }
    else:
        print(f"  ✗ Search failed: {msg}")
        return None, {
            'success': False,
            'message': msg
        }

# ============================================================================
# STABILITY ANALYSIS
# ============================================================================

def compute_stability_matrix(fixed_point, method='LPA', delta=1e-5):
    """
    Compute stability matrix (Jacobian of beta functions at fixed point).
    
    Eigenvalues determine stability:
    - Negative eigenvalue → relevant (UV attractive)
    - Positive eigenvalue → irrelevant (UV repulsive)
    
    Args:
        fixed_point: Fixed point location
        method: 'LPA' or 'LPA_prime'
        delta: Step size for numerical derivative
        
    Returns:
        jacobian: Stability matrix
        eigenvalues: Critical exponents
        eigenvectors: Directions
    """
    if method == 'LPA':
        beta_func = lambda y: beta_functions_LPA(0, y)
        n = 5
    elif method == 'LPA_prime':
        beta_func = lambda y: beta_functions_LPA_prime(0, y)
        n = 7
    else:
        raise ValueError(f"Unknown method: {method}")
    
    # Compute Jacobian numerically
    jacobian = np.zeros((n, n))
    
    for i in range(n):
        # Perturb i-th component
        y_plus = fixed_point.copy()
        y_minus = fixed_point.copy()
        y_plus[i] += delta
        y_minus[i] -= delta
        
        # Central difference
        jacobian[:, i] = (beta_func(y_plus) - beta_func(y_minus)) / (2*delta)
    
    # Compute eigenvalues
    eigenvalues, eigenvectors = np.linalg.eig(jacobian)
    
    # Sort by real part (most negative first = most relevant)
    idx = np.argsort(-eigenvalues.real)
    eigenvalues = eigenvalues[idx]
    eigenvectors = eigenvectors[:, idx]
    
    return jacobian, eigenvalues, eigenvectors

# ============================================================================
# VISUALIZATION
# ============================================================================

def plot_RG_flow(sol, method='LPA', save_path=None):
    """
    Plot RG flow trajectories.
    
    Args:
        sol: Solution from solve_ivp
        method: 'LPA' or 'LPA_prime'
        save_path: If provided, save figure
    """
    t = sol.t
    k_array = k_UV * np.exp(-t)
    
    if method == 'LPA':
        m2_2, m2_3, lam2, lam3, lam23 = sol.y
        labels = ['m²₂', 'm²₃', 'λ₂', 'λ₃', 'λ₂₃']
        n_plots = 5
    else:  # LPA_prime
        m2_2, m2_3, lam2, lam3, lam23, Z2, Z3 = sol.y
        labels = ['m²₂', 'm²₃', 'λ₂', 'λ₃', 'λ₂₃', 'Z₂', 'Z₃']
        n_plots = 7
    
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()
    
    for i in range(n_plots):
        ax = axes[i]
        ax.plot(k_array, sol.y[i], 'b-', linewidth=2)
        ax.set_xlabel('RG scale k [eV]', fontsize=12)
        ax.set_ylabel(labels[i], fontsize=12)
        ax.set_xscale('log')
        ax.grid(True, alpha=0.3)
        ax.set_title(f'Running of {labels[i]}', fontsize=14)
    
    # Remove empty subplots
    if n_plots < len(axes):
        for i in range(n_plots, len(axes)):
            fig.delaxes(axes[i])
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
    
    plt.show()

def plot_phase_diagram(results, method='LPA', save_path=None):
    """
    Plot phase diagram showing basins of attraction to fixed points.
    
    Args:
        results: List of (initial_conditions, final_values) tuples
        method: 'LPA' or 'LPA_prime'
        save_path: If provided, save figure
    """
    # Extract data
    m2_init = []
    lam_init = []
    m2_final = []
    lam_final = []
    
    for ic, fv in results:
        m2_init.append(ic[0])
        lam_init.append(ic[2])
        m2_final.append(fv[0])
        lam_final.append(fv[2])
    
    m2_init = np.array(m2_init)
    lam_init = np.array(lam_init)
    m2_final = np.array(m2_final)
    lam_final = np.array(lam_final)
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
    
    # Initial conditions
    scatter1 = ax1.scatter(m2_init, lam_init, c=m2_final, cmap='viridis', s=50)
    ax1.set_xlabel('Initial m²/k² [dimensionless]', fontsize=12)
    ax1.set_ylabel('Initial λ [dimensionless]', fontsize=12)
    ax1.set_title('Initial Conditions Colored by Final m²', fontsize=14)
    ax1.set_xscale('log')
    ax1.set_yscale('log')
    ax1.grid(True, alpha=0.3)
    plt.colorbar(scatter1, ax=ax1, label='Final m²/k²')
    
    # Final values
    scatter2 = ax2.scatter(m2_final, lam_final, c=np.log10(lam_init), cmap='plasma', s=50)
    ax2.set_xlabel('Final m²/k² [dimensionless]', fontsize=12)
    ax2.set_ylabel('Final λ [dimensionless]', fontsize=12)
    ax2.set_title('Final Values Colored by log₁₀(Initial λ)', fontsize=14)
    ax2.set_xscale('log')
    ax2.set_yscale('log')
    ax2.grid(True, alpha=0.3)
    plt.colorbar(scatter2, ax=ax2, label='log₁₀(Initial λ)')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Phase diagram saved to {save_path}")
    
    plt.show()

# ============================================================================
# MAIN EXECUTION
# ============================================================================

def main():
    """
    Main execution: Run systematic search for fixed points.
    """
    print("=" * 70)
    print("ASYMPTOTIC SAFETY EXPLORER - LPA/LPA' Implementation")
    print("=" * 70)
    print()
    
    # Test 1: Single trajectory (LPA)
    print("TEST 1: Single RG flow (LPA)")
    print("-" * 70)
    
    # Initial conditions [m²₂/k², m²₃/k², λ₂, λ₃, λ₂₃]
    ic_test = np.array([0.1, 0.1, 0.5, 0.5, 0.3])
    
    sol_test = run_RG_flow(ic_test, t_max=15, method='LPA')
    
    if sol_test.success:
        plot_RG_flow(sol_test, method='LPA', 
                    save_path='/mnt/user-data/outputs/RG_flow_test_LPA.png')
    
    print()
    
    # Test 2: Fixed point search (LPA)
    print("TEST 2: Fixed point search (LPA)")
    print("-" * 70)
    
    # Try different initial guesses
    guesses = [
        np.array([0.0, 0.0, 0.0, 0.0, 0.0]),  # Trivial (Gaussian)
        np.array([0.1, 0.1, 0.5, 0.5, 0.2]),  # Non-trivial candidate
        np.array([1.0, 1.0, 1.0, 1.0, 0.5]),  # Strong coupling
    ]
    
    fixed_points = []
    for i, guess in enumerate(guesses):
        print(f"\nAttempt {i+1}:")
        fp, info = find_fixed_point(guess, method='LPA')
        if fp is not None:
            fixed_points.append(fp)
    
    print()
    print(f"Found {len(fixed_points)} fixed point(s)")
    
    # Stability analysis for found fixed points
    if fixed_points:
        print("\nSTABILITY ANALYSIS:")
        print("-" * 70)
        
        for i, fp in enumerate(fixed_points):
            print(f"\nFixed Point {i+1}: {fp}")
            
            jac, eigs, vecs = compute_stability_matrix(fp, method='LPA')
            
            print(f"  Eigenvalues (critical exponents):")
            for j, eig in enumerate(eigs):
                if np.abs(eig.imag) < 1e-10:
                    print(f"    θ_{j} = {eig.real:.4f}")
                else:
                    print(f"    θ_{j} = {eig.real:.4f} + {eig.imag:.4f}i")
            
            # Classify operators
            n_relevant = np.sum(eigs.real < 0)
            n_irrelevant = np.sum(eigs.real > 0)
            n_marginal = np.sum(np.abs(eigs.real) < 1e-6)
            
            print(f"\n  Classification:")
            print(f"    Relevant operators: {n_relevant}")
            print(f"    Irrelevant operators: {n_irrelevant}")
            print(f"    Marginal operators: {n_marginal}")
    
    print()
    print("=" * 70)
    print("EXPLORATION COMPLETE!")
    print("=" * 70)
    
    # Save results
    results_dict = {
        'timestamp': datetime.now().isoformat(),
        'method': 'LPA',
        'k_UV_eV': float(k_UV),
        'fixed_points': [fp.tolist() for fp in fixed_points],
        'test_trajectory': {
            'initial': ic_test.tolist(),
            'final': sol_test.y[:, -1].tolist() if sol_test.success else None
        }
    }
    
    with open('/mnt/user-data/outputs/asymptotic_safety_results.json', 'w') as f:
        json.dump(results_dict, f, indent=2)
    
    print("\nResults saved to: /mnt/user-data/outputs/asymptotic_safety_results.json")

if __name__ == '__main__':
    main()
