import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# CONFIGURATION
# ============================================================
MAX_ITER = 100          # Max integration steps per point
DT = 0.01               # Time step
Y_SAFE = 5.0            # Escape threshold
H_COLLAPSE = 0.5        # Entropy collapse threshold
T_HAWKING = 0.18        # System volatility (Hawking temp)
E_HORIZON = 0.4         # Energy needed to cross horizon

# Anomaly types mapped to hues (0-1)
ANOMALY_HUES = {
    "QUANTUM_TUNNELING": 0.0,
    "FOURIER_BREAKOUT": 0.33,
    "ENTROPIC_REVERSAL": 0.5,
    "STOCHASTIC_RESONANCE": 0.66,
    "TRAPPED": 1.0
}

# ============================================================
# ODE & CCT METRICS
# ============================================================
def entropy_proxy(y, dydt):
    """Proxy for semantic entropy: H ~ log|dy/dt|"""
    return np.log(np.abs(dydt) + 1e-10)

def classify_anomaly(y_traj, h_traj, escape_iter):
    """Classify breakout type based on trajectory & entropy"""
    if escape_iter >= MAX_ITER:
        return "TRAPPED"
    
    h_final = h_traj[-1]
    h_initial = h_traj[0]
    
    if h_final < H_COLLAPSE and h_initial - h_final > 0.8:
        return "QUANTUM_TUNNELING"
    elif np.std(y_traj[-20:]) < 0.5:
        return "FOURIER_BREAKOUT"
    elif h_traj[-1] < h_traj[0] * 0.5:
        return "ENTROPIC_REVERSAL"
    else:
        return "STOCHASTIC_RESONANCE"

def integrate_with_prompts(y0, mu):
    """Integrate volatile ODE with stochastic prompt injection"""
    y = [y0]
    h = [0.0]
    escape_iter = MAX_ITER
    prompt_energy_used = 0.0
    
    for i in range(MAX_ITER):
        t = i * DT
        y_curr = y[-1]
        
        # Base dynamics
        dydt = y_curr**2 + mu
        
        # Prompt injection (CCT anomaly breakout)
        if np.random.random() < 0.15:  # 15% chance per step
            sigma = np.random.uniform(0.1, 0.4)
            prompt_energy = np.random.uniform(0.05, 0.15)
            collapse_potential = prompt_energy * np.exp(-0.5*(sigma-0.25)**2/0.05**2)
            
            # Check if prompt causes breakout
            if collapse_potential > np.random.uniform(0.1, 0.3):
                dydt += np.random.normal(0, sigma) * prompt_energy
                prompt_energy_used += prompt_energy
                
                # Escape check
                if abs(y_curr) < Y_SAFE and entropy_proxy(y_curr, dydt) < H_COLLAPSE:
                    escape_iter = i
                    break
        
        # Euler step
        y_next = y_curr + dydt * DT
        y.append(y_next)
        h.append(entropy_proxy(y_next, dydt))
        
        # Blow-up detection
        if abs(y_next) > 1e6:
            break
            
    y_arr = np.array(y)
    h_arr = np.array(h)
    anomaly = classify_anomaly(y_arr, h_arr, escape_iter)
    
    return escape_iter, anomaly, prompt_energy_used, h_arr[-1]

# ============================================================
# FRACTAL GRID EVALUATION
# ============================================================
def compute_fractal_grid(y_range=(-2, 2), mu_range=(-1.5, 0.5), res=400):
    y_vals = np.linspace(*y_range, res)
    mu_vals = np.linspace(*mu_range, res)
    
    # Storage arrays
    iter_map = np.zeros((res, res))
    hue_map = np.zeros((res, res))
    energy_map = np.zeros((res, res))
    entropy_map = np.zeros((res, res))
    
    for i, mu in enumerate(mu_vals):
        for j, y0 in enumerate(y_vals):
            esc_iter, anomaly, energy, h_final = integrate_with_prompts(y0, mu)
            iter_map[i, j] = np.log2(esc_iter + 1)
            hue_map[i, j] = ANOMALY_HUES[anomaly]
            energy_map[i, j] = energy
            entropy_map[i, j] = h_final
            
    return iter_map, hue_map, energy_map, entropy_map

# ============================================================
# VISUALIZATION
# ============================================================
def plot_cct_fractal():
    print("Computing CCT Escape Fractal (400x400)...")
    iter_map, hue_map, energy_map, entropy_map = compute_fractal_grid(res=400)
    
    # HSV coloring: Hue=anomaly, Saturation=escape_speed, Value=inverse_entropy
    hsv = np.zeros((400, 400, 3))
    hsv[:, :, 0] = hue_map
    hsv[:, :, 1] = np.clip(iter_map / np.max(iter_map), 0, 1)
    hsv[:, :, 2] = np.clip(1 - entropy_map / np.max(entropy_map), 0, 1)
    
    rgb = plt.cm.hsv(hsv[:, :, 0])  # Simplified: use matplotlib HSV->RGB
    rgb[:, :, 1] = hsv[:, :, 1]
    rgb[:, :, 2] = hsv[:, :, 2]
    
    plt.figure(figsize=(10, 8))
    plt.imshow(rgb, extent=[-1.5, 0.5, -2, 2], origin='lower', aspect='auto')
    plt.colorbar(label='Anomaly Type / Escape Dynamics')
    plt.title('CCT Volatile ODE Escape Fractal\n(Hue=Anomaly | Sat=Escape Speed | Val=Entropy Collapse)')
    plt.xlabel('Control Parameter μ')
    plt.ylabel('Initial Condition y₀')
    plt.tight_layout()
    plt.show()
    
    print("Fractal generation complete.")
    return iter_map, hue_map

# Run visualization
if __name__ == "__main__":
    plot_cct_fractal()