import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import entropy
from typing import List, Dict, Tuple, Optional, Callable
from dataclasses import dataclass, field
import hashlib

# ============================================================
# DATA STRUCTURES
# ============================================================

@dataclass
class FunctionGate:
    """Represents a single function gate in the superposition"""
    gate_id: str
    function: Callable[[np.ndarray], np.ndarray]
    weight: float = 0.0
    collapse_potential: float = 0.0
    age: int = 0
    active: bool = True
    metadata: Dict = field(default_factory=dict)
    
    def __str__(self):
        return f"Gate({self.gate_id}, w={self.weight:.4f})"

@dataclass
class FittingResult:
    """Stores results from semantic function regression"""
    status: str
    best_gate: str
    confidence: float
    final_entropy: float
    r2_score: float
    mse: float
    iterations: int
    gates_evaluated: int
    gates_active: int
    entropy_trajectory: List[float]
    weight_history: List[np.ndarray]

# ============================================================
# SEMANTIC FUNCTION REGRESSOR
# ============================================================

class SemanticFunctionRegressor:
    """
    🛸 CCT-ODE Semantic Function Regressor
    
    Fits time series data to a superposition of function gates
    using Conditional Collapse Theory dynamics.
    
    Key Features:
    - Maintains probability distribution over function space
    - Uses replicator ODE for weight evolution
    - Entropy-gated memory pruning
    - Symbolic function discovery
    - ODE-CCT periodicity detection
    """
    
    # ============================================================
    # CONFIGURATION
    # ============================================================
    
    def __init__(self,
                 alpha: float = 0.05,
                 entropy_threshold: float = 0.1,
                 prune_threshold: float = 0.01,
                 max_active_gates: int = 50,
                 decay_rate: float = 0.01,
                 gain_rate: float = 0.1,
                 mutation_rate: float = 0.05,
                 periodicity_detection: bool = True,
                 verbose: bool = True):
        """
        Args:
            alpha: Learning rate for replicator ODE
            entropy_threshold: Target entropy for collapse
            prune_threshold: Weight below which gates are pruned
            max_active_gates: Maximum gates to keep active
            decay_rate: Weight decay for unused gates
            gain_rate: Weight gain from collapse potential
            mutation_rate: Rate of symbolic mutation for new gates
            periodicity_detection: Enable ODE-CCT cycle detection
            verbose: Print progress during fitting
        """
        self.alpha = alpha
        self.epsilon = entropy_threshold
        self.prune_thresh = prune_threshold
        self.MAX_GATES = max_active_gates
        self.decay_rate = decay_rate
        self.gain_rate = gain_rate
        self.mutation_rate = mutation_rate
        self._detect_periodicity_enabled = periodicity_detection
        self.verbose = verbose
        
        # Function Gate Superposition
        self.gates: List[FunctionGate] = []
        self.active_indices: List[int] = []
        
        # CCT Metrics
        self.entropy_history: List[float] = []
        self.weight_history: List[np.ndarray] = []
        self.collapse_potentials: np.ndarray = np.array([])
        
        # Periodicity Detection (ODE-CCT)
        self.state_hashes: List[int] = []
        self.period_detected: Optional[int] = None
        self.period_start_iter: Optional[int] = None
        
        # Fitting State
        self.fitted = False
        self.X_train: Optional[np.ndarray] = None
        self.y_train: Optional[np.ndarray] = None
        
        # Initialize standard function gate library
        self._initialize_gate_library()
        
    # ============================================================
    # FUNCTION GATE LIBRARY (Float-Logic Primitives)
    # ============================================================
    
    def _initialize_gate_library(self):
        """Initialize standard function gates for time series"""
        self.gates = []
        
        # ===== Linear Gates =====
        self.gates.append(FunctionGate(
            gate_id="linear",
            function=lambda x: x,
            weight=1.0/20,
            metadata={"type": "linear", "complexity": 1}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="linear_offset",
            function=lambda x: x + 0.5,
            weight=1.0/20,
            metadata={"type": "linear", "complexity": 1}
        ))
        
        # ===== Polynomial Gates =====
        self.gates.append(FunctionGate(
            gate_id="quadratic",
            function=lambda x: x ** 2,
            weight=1.0/20,
            metadata={"type": "polynomial", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="cubic",
            function=lambda x: x ** 3,
            weight=1.0/20,
            metadata={"type": "polynomial", "complexity": 3}
        ))
        
        # ===== Periodic Gates =====
        self.gates.append(FunctionGate(
            gate_id="sine",
            function=lambda x: np.sin(2 * np.pi * x),
            weight=1.0/20,
            metadata={"type": "periodic", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="cosine",
            function=lambda x: np.cos(2 * np.pi * x),
            weight=1.0/20,
            metadata={"type": "periodic", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="sine_half",
            function=lambda x: np.sin(np.pi * x),
            weight=1.0/20,
            metadata={"type": "periodic", "complexity": 2}
        ))
        
        # ===== Exponential Gates =====
        self.gates.append(FunctionGate(
            gate_id="exp",
            function=lambda x: np.exp(x - 1),
            weight=1.0/20,
            metadata={"type": "exponential", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="exp_decay",
            function=lambda x: np.exp(-x),
            weight=1.0/20,
            metadata={"type": "exponential", "complexity": 2}
        ))
        
        # ===== Logarithmic Gates =====
        self.gates.append(FunctionGate(
            gate_id="log",
            function=lambda x: np.log(np.abs(x) + 0.1),
            weight=1.0/20,
            metadata={"type": "logarithmic", "complexity": 2}
        ))
        
        # ===== Composite Gates =====
        self.gates.append(FunctionGate(
            gate_id="sin_exp",
            function=lambda x: np.sin(x) * np.exp(-x / 5),
            weight=1.0/20,
            metadata={"type": "composite", "complexity": 3}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="poly_sin",
            function=lambda x: x * np.sin(2 * np.pi * x),
            weight=1.0/20,
            metadata={"type": "composite", "complexity": 3}
        ))
        
        # ===== Sigmoid Gates =====
        self.gates.append(FunctionGate(
            gate_id="sigmoid",
            function=lambda x: 1 / (1 + np.exp(-5 * (x - 0.5))),
            weight=1.0/20,
            metadata={"type": "sigmoid", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="tanh",
            function=lambda x: np.tanh(3 * (x - 0.5)),
            weight=1.0/20,
            metadata={"type": "sigmoid", "complexity": 2}
        ))
        
        # ===== Step Gates =====
        self.gates.append(FunctionGate(
            gate_id="step",
            function=lambda x: (x > 0.5).astype(float),
            weight=1.0/20,
            metadata={"type": "step", "complexity": 1}
        ))
        
        # ===== Constant Gates =====
        self.gates.append(FunctionGate(
            gate_id="constant",
            function=lambda x: np.ones_like(x) * 0.5,
            weight=1.0/20,
            metadata={"type": "constant", "complexity": 0}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="zero",
            function=lambda x: np.zeros_like(x),
            weight=1.0/20,
            metadata={"type": "constant", "complexity": 0}
        ))
        
        # ===== Advanced Gates =====
        self.gates.append(FunctionGate(
            gate_id="gaussian",
            function=lambda x: np.exp(-((x - 0.5) ** 2) / 0.1),
            weight=1.0/20,
            metadata={"type": "gaussian", "complexity": 2}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="abs",
            function=lambda x: np.abs(x - 0.5),
            weight=1.0/20,
            metadata={"type": "piecewise", "complexity": 1}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="sqrt",
            function=lambda x: np.sqrt(np.abs(x)),
            weight=1.0/20,
            metadata={"type": "radical", "complexity": 1}
        ))
        
        self.gates.append(FunctionGate(
            gate_id="inv",
            function=lambda x: 1 / (np.abs(x) + 0.1),
            weight=1.0/20,
            metadata={"type": "rational", "complexity": 2}
        ))
        
        # Initialize active indices
        self.active_indices = list(range(len(self.gates)))
        
    def add_custom_gate(self,
                        gate_id: str,
                        function: Callable[[np.ndarray], np.ndarray],
                        complexity: int = 2):
        """Add a custom function gate to the superposition"""
        gate = FunctionGate(
            gate_id=gate_id,
            function=function,
            weight=1.0 / (len(self.gates) + 1),
            metadata={"type": "custom", "complexity": complexity}
        )
        self.gates.append(gate)
        self.active_indices.append(len(self.gates) - 1)
        
    # ============================================================
    # CCT-ODE CORE DYNAMICS
    # ============================================================
    
    def calculate_collapse_potential(self,
                                     X: np.ndarray,
                                     y: np.ndarray,
                                     gate_idx: int) -> float:
        """
        Calculate collapse potential Δ_f for a function gate.
        Higher = better fit (negative MSE)
        """
        if gate_idx >= len(self.gates):
            return -1e10
            
        gate = self.gates[gate_idx]
        try:
            y_pred = gate.function(X)
            # Handle NaN/Inf
            y_pred = np.nan_to_num(y_pred, nan=0.0, posinf=1e10, neginf=-1e10)
            y_pred = np.clip(y_pred, -1e10, 1e10)
            
            # Negative MSE as collapse potential
            mse = np.mean((y_pred - y) ** 2)
            return -mse
        except Exception as e:
            return -1e10
    
    def replicator_update(self, X: np.ndarray, y: np.ndarray) -> None:
        """
        CCT Replicator ODE: dp/dt = α·p·(Δ - Δ̄)
        Updates gate weights based on collapse potential
        """
        if len(self.active_indices) == 0:
            return
            
        # Calculate collapse potentials for all active gates
        deltas = np.zeros(len(self.gates))
        for idx in self.active_indices:
            deltas[idx] = self.calculate_collapse_potential(X, y, idx)
        
        self.collapse_potentials = deltas
        
        # Get active weights and deltas
        active_weights = np.array([self.gates[i].weight for i in self.active_indices])
        active_deltas = deltas[self.active_indices]
        
        # Average collapse potential
        avg_delta = np.sum(active_weights * active_deltas)
        
        # Replicator ODE step
        dp = self.alpha * active_weights * (active_deltas - avg_delta)
        
        # Update weights
        for i, idx in enumerate(self.active_indices):
            self.gates[idx].weight += dp[i]
            self.gates[idx].collapse_potential = deltas[idx]
        
        # Renormalize active weights
        total_weight = sum(self.gates[i].weight for i in self.active_indices)
        if total_weight > 0:
            for idx in self.active_indices:
                self.gates[idx].weight /= total_weight
        
        # Ensure positivity and clip
        for idx in self.active_indices:
            self.gates[idx].weight = np.clip(self.gates[idx].weight, 1e-10, 1.0)
    
    def calculate_entropy(self) -> float:
        """Calculate Shannon entropy over function gate space"""
        if len(self.active_indices) == 0:
            return 0.0
            
        weights = np.array([self.gates[i].weight for i in self.active_indices])
        weights = np.clip(weights, 1e-10, 1.0)
        weights /= np.sum(weights)
        
        return entropy(weights, base=2)
    
    # ============================================================
    # MEMORY PRUNING (Entropy-Gated Forgetting)
    # ============================================================
    
    def prune_memory(self, iteration: int) -> int:
        """
        Prune low-potential gates to save compute
        Returns number of gates pruned
        """
        gates_to_prune = []
        
        for idx in self.active_indices:
            gate = self.gates[idx]
            
            # Update age
            gate.age += 1
            
            # Weight decay for unused gates
            if gate.collapse_potential < -1.0:  # Poor performer
                gate.weight *= (1 - self.decay_rate)
            
            # Check pruning threshold
            if gate.weight < self.prune_thresh:
                gates_to_prune.append(idx)
        
        # Execute pruning
        for idx in gates_to_prune:
            self.active_indices.remove(idx)
            self.gates[idx].active = False
        
        # Enforce max active gates
        if len(self.active_indices) > self.MAX_GATES:
            # Sort by weight and keep top K
            sorted_indices = sorted(
                self.active_indices,
                key=lambda x: self.gates[x].weight,
                reverse=True
            )
            to_archive = sorted_indices[self.MAX_GATES:]
            for idx in to_archive:
                self.active_indices.remove(idx)
                self.gates[idx].active = False
        
        return len(gates_to_prune)
    
    def mutate_and_discover(self, iteration: int, X: np.ndarray, y: np.ndarray):
        """
        Symbolic Discovery: Generate new gates via mutation
        Only when entropy is high (exploration needed)
        """
        current_entropy = self.calculate_entropy()
        
        # Only mutate if entropy is high and we have room
        if current_entropy < self.epsilon or len(self.active_indices) >= self.MAX_GATES:
            return
            
        if np.random.random() > self.mutation_rate:
            return
            
        # Mutate top performing gate
        if len(self.active_indices) > 0:
            best_idx = max(self.active_indices, 
                          key=lambda x: self.gates[x].weight)
            best_gate = self.gates[best_idx]
            
            # Create mutated version
            new_gate_id = f"{best_gate.gate_id}_mut_{iteration}"
            
            # Simple mutation: add small perturbation
            base_func = best_gate.function
            perturbation = np.random.uniform(-0.5, 0.5)
            
            def mutated_function(x, base=base_func, p=perturbation):
                return base(x) + p * np.sin(2 * np.pi * x * np.random.uniform(0.5, 2.0))
            
            new_gate = FunctionGate(
                gate_id=new_gate_id,
                function=mutated_function,
                weight=0.01,
                metadata={"type": "mutated", "parent": best_gate.gate_id}
            )
            
            self.gates.append(new_gate)
            self.active_indices.append(len(self.gates) - 1)
    
    # ============================================================
    # PERIODICITY DETECTION (ODE-CCT)
    # ============================================================
    
    def detect_periodicity(self, iteration: int) -> bool:
        """
        Detect limit cycles in weight trajectory (ODE-CCT Periodicity)
        Returns True if periodicity detected
        """
        if not self._detect_periodicity_enabled:
            return False
            
        # Hash current weight state
        weight_hash = hash(tuple(round(self.gates[i].weight, 4) 
                                 for i in self.active_indices[:10]))
        
        self.state_hashes.append(weight_hash)
        
        # Check for cycle (look back up to 50 iterations)
        if len(self.state_hashes) > 10:
            for k in range(5, min(50, len(self.state_hashes) - 1)):
                if self.state_hashes[-1] == self.state_hashes[-(k+1)]:
                    # Verify cycle continues
                    if len(self.state_hashes) > k + 2:
                        if self.state_hashes[-2] == self.state_hashes[-(k+2)]:
                            self.period_detected = k
                            self.period_start_iter = iteration - k
                            return True
        
        return False
    
    # ============================================================
    # FITTING INTERFACE
    # ============================================================
    
    def fit(self,
            X: np.ndarray,
            y: np.ndarray,
            max_iters: int = 100,
            early_stop: bool = True) -> FittingResult:
        """
        Fit time series data to function gate superposition
        
        Args:
            X: Input features (time or independent variable)
            y: Target values
            max_iters: Maximum fitting iterations
            early_stop: Stop when entropy collapses
            
        Returns:
            FittingResult with metrics and trajectories
        """
        self.X_train = X
        self.y_train = y
        self.fitted = False
        
        if self.verbose:
            print("=" * 70)
            print("🛸 SEMANTIC FUNCTION REGRESSOR: CCT-ODE FITTING")
            print("=" * 70)
            print(f"Samples: {len(X)} | Initial Gates: {len(self.active_indices)}")
            print(f"Entropy Threshold: {self.epsilon} | Max Iterations: {max_iters}")
            print("-" * 70)
        
        # Initial entropy
        H0 = self.calculate_entropy()
        self.entropy_history.append(H0)
        self.weight_history.append(np.array([self.gates[i].weight 
                                              for i in self.active_indices]))
        
        if self.verbose:
            print(f"Initial Entropy: {H0:.4f} bits")
        
        # Fitting loop
        for t in range(max_iters):
            # 1. Replicator ODE Update
            self.replicator_update(X, y)
            
            # 2. Calculate Entropy
            H = self.calculate_entropy()
            self.entropy_history.append(H)
            self.weight_history.append(np.array([self.gates[i].weight 
                                                  for i in self.active_indices]))
            
            # 3. Memory Pruning
            pruned = self.prune_memory(t)
            
            # 4. Symbolic Discovery
            self.mutate_and_discover(t, X, y)
            
            # 5. Periodicity Detection
            if self.detect_periodicity(t):
                if self.verbose:
                    print(f"[PERIOD] Cycle detected at iteration {t} (period={self.period_detected})")
                # Can early exit if periodicity stable
                if H < self.epsilon * 2:
                    break
            
            # 6. Verbose Logging
            if self.verbose and t % 10 == 0:
                best_gate = max(self.active_indices, 
                               key=lambda x: self.gates[x].weight) if self.active_indices else None
                best_weight = self.gates[best_gate].weight if best_gate else 0
                print(f"Iter {t:3d}: H={H:.4f} | Gates={len(self.active_indices)} | "
                      f"Best={self.gates[best_gate].gate_id if best_gate else 'N/A'} "
                      f"({best_weight:.2%}) | Pruned={pruned}")
            
            # 7. Early Stopping
            if early_stop and H < self.epsilon:
                if self.verbose:
                    print(f"\n[✓] ENTROPY COLLAPSE at iteration {t}")
                break
        
        # Calculate final metrics
        best_gate_idx = max(self.active_indices, 
                           key=lambda x: self.gates[x].weight) if self.active_indices else None
        
        if best_gate_idx is not None:
            best_gate = self.gates[best_gate_idx]
            y_pred = best_gate.function(X)
            mse = np.mean((y_pred - y) ** 2)
            ss_res = np.sum((y - y_pred) ** 2)
            ss_tot = np.sum((y - np.mean(y)) ** 2)
            r2 = 1 - (ss_res / (ss_tot + 1e-10))
        else:
            mse = float('inf')
            r2 = 0.0
            best_gate = None
        
        self.fitted = True
        
        result = FittingResult(
            status="COLLAPSED" if self.entropy_history[-1] < self.epsilon else "PARTIAL",
            best_gate=best_gate.gate_id if best_gate else "NONE",
            confidence=best_gate.weight if best_gate else 0.0,
            final_entropy=self.entropy_history[-1],
            r2_score=r2,
            mse=mse,
            iterations=len(self.entropy_history) - 1,
            gates_evaluated=len(self.gates),
            gates_active=len(self.active_indices),
            entropy_trajectory=self.entropy_history,
            weight_history=self.weight_history
        )
        
        if self.verbose:
            print("-" * 70)
            print("FITTING RESULTS:")
            print(f"  Status: {result.status}")
            print(f"  Best Gate: {result.best_gate}")
            print(f"  Confidence: {result.confidence:.2%}")
            print(f"  R² Score: {result.r2_score:.4f}")
            print(f"  MSE: {result.mse:.6f}")
            print(f"  Final Entropy: {result.final_entropy:.4f} bits")
            print(f"  Iterations: {result.iterations}")
            print(f"  Gates Active: {result.gates_active} / {result.gates_evaluated}")
            if self.period_detected:
                print(f"  Periodicity: Detected (period={self.period_detected})")
            print("=" * 70)
        
        return result
    
    # ============================================================
    # PREDICTION
    # ============================================================
    
    def predict(self,
                X: np.ndarray,
                collapsed: bool = True,
                return_uncertainty: bool = False) -> np.ndarray:
        """
        Make predictions using fitted function gates
        
        Args:
            X: Input features
            collapsed: Use best gate only (True) or weighted ensemble (False)
            return_uncertainty: Also return prediction uncertainty
            
        Returns:
            Predictions (and optionally uncertainty)
        """
        if not self.fitted:
            raise ValueError("Model must be fitted before prediction")
        
        if collapsed:
            # Use best gate only
            best_idx = max(self.active_indices, 
                          key=lambda x: self.gates[x].weight)
            y_pred = self.gates[best_idx].function(X)
        else:
            # Weighted ensemble (superposition)
            y_pred = np.zeros(len(X))
            for idx in self.active_indices:
                y_pred += self.gates[idx].weight * self.gates[idx].function(X)
        
        if return_uncertainty:
            # Uncertainty = entropy-weighted variance across gates
            if len(self.active_indices) > 1:
                predictions = np.array([self.gates[i].function(X) 
                                       for i in self.active_indices])
                uncertainty = np.std(predictions, axis=0) * self.calculate_entropy()
            else:
                uncertainty = np.zeros(len(X))
            return y_pred, uncertainty
        
        return y_pred
    
    # ============================================================
    # VISUALIZATION
    # ============================================================
    
    def plot_entropy_trajectory(self):
        """Plot entropy collapse over fitting iterations"""
        plt.figure(figsize=(10, 5))
        plt.plot(self.entropy_history, 'b-o', linewidth=2, markersize=6)
        plt.axhline(y=self.epsilon, color='r', linestyle='--', 
                   label=f'Collapse Threshold ({self.epsilon})')
        plt.xlabel('Iteration')
        plt.ylabel('Semantic Entropy H(Ψ) [bits]')
        plt.title('CCT-ODE Entropy Collapse Trajectory')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    def plot_weight_evolution(self, top_n: int = 10):
        """Plot evolution of top gate weights"""
        if len(self.weight_history) == 0:
            return
            
        plt.figure(figsize=(12, 6))
        
        # Get top N gates by final weight
        final_weights = [(i, self.gates[i].weight) for i in self.active_indices]
        top_indices = sorted(final_weights, key=lambda x: x[1], reverse=True)[:top_n]
        
        for idx, _ in top_indices:
            weights_over_time = [w[idx] if idx < len(w) else 0 
                                for w in self.weight_history]
            plt.plot(weights_over_time, linewidth=2, 
                    label=f"{self.gates[idx].gate_id}")
        
        plt.xlabel('Iteration')
        plt.ylabel('Gate Weight')
        plt.title('Function Gate Weight Evolution (Superposition Collapse)')
        plt.legend(loc='upper right', fontsize=8)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    def plot_fit(self, X_test: np.ndarray = None, y_test: np.ndarray = None):
        """Plot fitted function against data"""
        if not self.fitted:
            return
            
        plt.figure(figsize=(12, 5))
        
        # Plot training data
        plt.subplot(1, 2, 1)
        plt.scatter(self.X_train, self.y_train, alpha=0.5, 
                   label='Training Data', color='blue')
        
        # Plot fitted function
        X_smooth = np.linspace(self.X_train.min(), self.X_train.max(), 200)
        y_pred = self.predict(X_smooth, collapsed=True)
        plt.plot(X_smooth, y_pred, 'r-', linewidth=2, 
                label=f'Best Gate: {self.gates[max(self.active_indices, key=lambda x: self.gates[x].weight)].gate_id}')
        
        plt.xlabel('X')
        plt.ylabel('y')
        plt.title('Semantic Function Fit')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # Plot residuals
        plt.subplot(1, 2, 2)
        y_train_pred = self.predict(self.X_train, collapsed=True)
        residuals = self.y_train - y_train_pred
        plt.scatter(y_train_pred, residuals, alpha=0.5, color='green')
        plt.axhline(y=0, color='r', linestyle='--')
        plt.xlabel('Predicted')
        plt.ylabel('Residual')
        plt.title('Residual Plot')
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
    
    def plot_gate_distribution(self):
        """Plot final gate weight distribution"""
        if not self.fitted:
            return
            
        plt.figure(figsize=(14, 6))
        
        weights = [self.gates[i].weight for i in self.active_indices]
        gate_names = [self.gates[i].gate_id for i in self.active_indices]
        
        # Sort by weight
        sorted_idx = np.argsort(weights)[::-1]
        weights = np.array(weights)[sorted_idx]
        gate_names = np.array(gate_names)[sorted_idx]
        
        plt.bar(range(len(weights)), weights, color='steelblue')
        plt.xticks(range(len(weights)), gate_names, rotation=90, fontsize=8)
        plt.xlabel('Function Gate')
        plt.ylabel('Weight (Probability)')
        plt.title('Final Function Gate Distribution (Collapsed Superposition)')
        plt.grid(True, alpha=0.3, axis='y')
        plt.tight_layout()
        plt.show()
    
    # ============================================================
    # EXPORT / IMPORT
    # ============================================================
    
    def get_summary(self) -> Dict:
        """Get summary of fitted model"""
        if not self.fitted:
            return {"status": "NOT_FITTED"}
        
        best_idx = max(self.active_indices, 
                      key=lambda x: self.gates[x].weight)
        
        return {
            "status": "FITTED",
            "best_gate": self.gates[best_idx].gate_id,
            "confidence": self.gates[best_idx].weight,
            "entropy": self.entropy_history[-1],
            "active_gates": len(self.active_indices),
            "total_gates": len(self.gates),
            "periodicity_detected": self.period_detected,
            "r2_score": self.predict(self.X_train, collapsed=True),
        }


# Generate synthetic data
np.random.seed(42)
X = np.linspace(0, 1, 100)
y = np.random.randn(100)

# Initialize and fit regressor
regressor = SemanticFunctionRegressor(
    alpha=0.1,
    entropy_threshold=0.1,
    max_active_gates=30,
    verbose=True
)

result = regressor.fit(X, y, max_iters=50)

# Predict
y_pred = regressor.predict(X, collapsed=True)

# Visualize
regressor.plot_entropy_trajectory()
regressor.plot_fit()
regressor.plot_gate_distribution()
