import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

class FractalEscapeTimeCCT:
    """
    Implements Fractal Escape Time Theory (FETT) within CCT Framework.
    Maps ML Seed + Fractal Iteration -> Escape Time -> Life Regions (Red Stars)
    """
    def __init__(self, resolution=500, max_iter=100, escape_radius=4.0):
        self.resolution = resolution
        self.max_iter = max_iter
        self.escape_radius = escape_radius
        self.escape_times = np.zeros((resolution, resolution))
        self.life_map = np.zeros((resolution, resolution))
        
        # CCT Thresholds for "Sustainable Life"
        self.min_escape = 10   # Too fast = Trivial
        self.max_escape = 80   # Too slow = Singularity (Black Hole)
        
    def generate_noise_blob(self):
        """
        Generates the 'Random Noise Blob' as the initial semantic manifold.
        Represents High Entropy Initial State.
        """
        re = np.linspace(-2.0, 1.0, self.resolution)
        im = np.linspace(-1.5, 1.5, self.resolution)
        C = np.empty((self.resolution, self.resolution), dtype=complex)
        C.real, C.imag = np.meshgrid(re, im)
        # Add noise to simulate semantic uncertainty
        noise = np.random.normal(0, 0.01, C.shape) + 1j * np.random.normal(0, 0.01, C.shape)
        return C + noise

    def ml_seed_function(self, C):
        """
        The 'ML Function' learns the seed manifold.
        Here we simulate a learned transformation theta based on the noise blob.
        In real AI, this would be weights from a neural network.
        """
        # Simulate learning: Adjust the fractal constant based on local density
        # This represents the 'Stationary' law learned from data
        theta = C * 0.8 + 0.2 * np.sin(C.real * 5) 
        return theta

    def compute_fractal_escape(self, C, theta):
        """
        The 'Fractal Function' computes time signal.
        Iterates ODE-like dynamics to find Escape Time (Compute Work).
        """
        Z = np.zeros_like(C)
        escape_times = np.zeros(C.shape)
        mask = np.ones(C.shape, dtype=bool)
        
        for t in range(1, self.max_iter + 1):
            # ODE-CCT Iteration: z_{t+1} = z_t^2 + theta
            # This is the 'Probability' component evolving over time
            Z[mask] = Z[mask]**2 + theta[mask]
            
            # Check Escape Condition (Entropy Collapse)
            escaped = np.abs(Z) > self.escape_radius
            newly_escaped = escaped & mask
            
            escape_times[newly_escaped] = t
            mask[newly_escaped] = False
            
            # Early exit if all escaped
            if not np.any(mask):
                break
                
        # Points that never escaped are Logic Singularities (Black Holes)
        escape_times[mask] = self.max_iter + 1
        return escape_times

    def identify_red_stars(self, escape_times):
        """
        Identifies 'Life Regions' where escape time is sustainable.
        Matches Red Star Theory: Escape Probability > Critical AND Entropy Collapses.
        """
        # Sustainable Life Band
        life_mask = (escape_times > self.min_escape) & (escape_times < self.max_escape)
        
        # Black Holes (Singularities)
        singularity_mask = (escape_times >= self.max_iter)
        
        # Trivial Regions (Too simple)
        trivial_mask = (escape_times <= self.min_escape)
        
        return life_mask, singularity_mask, trivial_mask

    def run_simulation(self):
        print("🛸 Initializing Fractal Escape Time Theory (FETT)...")
        
        # 1. Noise Blob (Initial State)
        C = self.generate_noise_blob()
        
        # 2. ML Seed (Stationary Law)
        theta = self.ml_seed_function(C)
        
        # 3. Fractal Time (Probability Evolution)
        print("Computing Escape Times (Compute Work)...")
        self.escape_times = self.compute_fractal_escape(C, theta)
        
        # 4. Red Star Identification (Life Regions)
        life, singularities, trivial = self.identify_red_stars(self.escape_times)
        self.life_map = life.astype(float)
        
        # 5. Visualization
        self.plot_results(C, life, singularities, trivial)
        
        # 6. Metrics
        total_points = self.resolution ** 2
        life_percent = np.sum(life) / total_points * 100
        singularity_percent = np.sum(singularities) / total_points * 100
        
        print(f"\n📊 FETT Universe Metrics:")
        print(f"  Total Manifold Points: {total_points}")
        print(f"  Red Star (Life) Regions: {life_percent:.2f}%")
        print(f"  Black Hole (Singularity) Regions: {singularity_percent:.2f}%")
        print(f"  Average Escape Time (Work): {np.mean(self.escape_times[life]):.2f} steps")
        
    def plot_results(self, C, life, singularities, trivial):
        fig, axs = plt.subplots(1, 3, figsize=(18, 5))
        
        # Plot 1: Escape Time Heatmap (Compute Work Landscape)
        im1 = axs[0].imshow(self.escape_times, extent=[C.real.min(), C.real.max(), C.imag.min(), C.imag.max()], 
                            origin='lower', cmap='magma', norm=LogNorm(vmin=1, vmax=self.max_iter))
        axs[0].set_title("Escape Time Landscape (Compute Work $W$)")
        axs[0].set_xlabel("Real Axis (Semantic Dimension 1)")
        axs[0].set_ylabel("Imaginary Axis (Semantic Dimension 2)")
        plt.colorbar(im1, ax=axs[0], label="Iterations (Work)")
        
        # Plot 2: Red Star Life Regions
        im2 = axs[1].imshow(self.life_map, extent=[C.real.min(), C.real.max(), C.imag.min(), C.imag.max()], 
                            origin='lower', cmap='Reds')
        axs[1].set_title("Red Star Regions (Sustainable Life)")
        axs[1].set_xlabel("Real Axis")
        axs[1].set_ylabel("Imaginary Axis")
        plt.colorbar(im2, ax=axs[1], label="Life Probability")
        
        # Plot 3: Singularity vs Life vs Trivial
        universe_map = np.zeros_like(self.life_map)
        universe_map[life] = 0.5   # Greenish (Life)
        universe_map[singularities] = 1.0 # Black (Hole)
        universe_map[trivial] = 0.2  # Blue (Trivial)
        
        im3 = axs[2].imshow(universe_map, extent=[C.real.min(), C.real.max(), C.imag.min(), C.imag.max()], 
                            origin='lower', cmap='RdYlGn')
        axs[2].set_title("Universe Classification (CCT)")
        axs[2].set_xlabel("Real Axis")
        axs[2].set_ylabel("Imaginary Axis")
        
        plt.tight_layout()
        plt.show()

# --- EXECUTION ---
if __name__ == "__main__":
    # Initialize FETT Engine
    fett = FractalEscapeTimeCCT(resolution=500, max_iter=100)
    # Run Simulation
    fett.run_simulation()