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 Circuit:
    """Represents a circuit (combination of multiple gates)"""
    gates: List[Tuple[str, float]]  # List of (gate_id, weight) tuples
    
    def evaluate(self, X: np.ndarray, gate_registry: Dict[str, FunctionGate]) -> np.ndarray:
        """Evaluate the circuit by combining gate outputs weighted by their weights"""
        result = np.zeros(len(X))
        for gate_id, weight in self.gates:
            if gate_id in gate_registry:
                result += weight * gate_registry[gate_id].function(X)
        return result
    
    def __str__(self):
        parts = [f"{gate_id}({w:.4f})" for gate_id, w in self.gates]
        return f"Circuit[{len(self.gates)}]: {' + '.join(parts)}"

@dataclass
class FittingResult:
    """Stores results from semantic function regression"""
    status: str
    best_gate: str  # Kept for backward compatibility
    best_circuit: Circuit  # New: best circuit combination
    confidence: float
    final_entropy: float
    r2_score: float
    r2_circuit: float  # R² score for the circuit
    mse: float
    mse_circuit: float  # MSE for the circuit
    iterations: int
    gates_evaluated: int
    gates_active: int
    circuit_size: int  # Number of gates in the circuit
    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,
                 circuit_size: int = 5,  # New: number of gates in circuit
                 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
            circuit_size: Number of gates to include in best circuit
            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.circuit_size = circuit_size  # New parameter
        self.detect_periodicity = 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 get_best_circuit(self, top_k: int = None) -> Circuit:
        """
        Get the best circuit (combination of top K gates by weight)
        
        Args:
            top_k: Number of gates to include (defaults to self.circuit_size)
        
        Returns:
            Circuit object with top K gates
        """
        if top_k is None:
            top_k = self.circuit_size
            
        if not self.active_indices:
            return Circuit(gates=[])
        
        # Sort active gates by weight (descending)
        sorted_gates = sorted(
            self.active_indices,
            key=lambda x: self.gates[x].weight,
            reverse=True
        )
        
        # Take top K gates
        top_gates = sorted_gates[:top_k]
        
        # Normalize weights so they sum to 1
        total_weight = sum(self.gates[idx].weight for idx in top_gates)
        if total_weight > 0:
            circuit_gates = [
                (self.gates[idx].gate_id, self.gates[idx].weight / total_weight)
                for idx in top_gates
            ]
        else:
            circuit_gates = [
                (self.gates[idx].gate_id, 1.0 / len(top_gates))
                for idx in top_gates
            ]
        
        return Circuit(gates=circuit_gates)
    
    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_cycle(self, iteration: int) -> bool:
        """
        Detect limit cycles in weight trajectory (ODE-CCT Periodicity)
        Returns True if periodicity detected
        """
        if not self.detect_periodicity:
            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_cycle(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_single = best_gate.function(X)
            mse = np.mean((y_pred_single - y) ** 2)
            ss_res = np.sum((y - y_pred_single) ** 2)
            ss_tot = np.sum((y - np.mean(y)) ** 2)
            r2 = 1 - (ss_res / (ss_tot + 1e-10))
            
            # Calculate circuit metrics
            best_circuit = self.get_best_circuit()
            if len(best_circuit.gates) > 0:
                # Create registry for circuit evaluation
                registry = {self.gates[i].gate_id: self.gates[i] for i in self.active_indices}
                y_pred_circuit = best_circuit.evaluate(X, registry)
                mse_circuit = np.mean((y_pred_circuit - y) ** 2)
                ss_res_circuit = np.sum((y - y_pred_circuit) ** 2)
                r2_circuit = 1 - (ss_res_circuit / (ss_tot + 1e-10))
                circuit_size = len(best_circuit.gates)
            else:
                mse_circuit = float('inf')
                r2_circuit = 0.0
                circuit_size = 0
        else:
            mse = float('inf')
            r2 = 0.0
            mse_circuit = float('inf')
            r2_circuit = 0.0
            best_gate = None
            best_circuit = Circuit(gates=[])
            circuit_size = 0

        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",
            best_circuit=best_circuit,
            confidence=best_gate.weight if best_gate else 0.0,
            final_entropy=self.entropy_history[-1],
            r2_score=r2,
            r2_circuit=r2_circuit,
            mse=mse,
            mse_circuit=mse_circuit,
            iterations=len(self.entropy_history) - 1,
            gates_evaluated=len(self.gates),
            gates_active=len(self.active_indices),
            circuit_size=circuit_size,
            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 Single Gate: {result.best_gate}")
            print(f"  Best Circuit: {result.best_circuit}")
            print(f"  Confidence (Single): {result.confidence:.2%}")
            print(f"  R² Score (Single Gate): {result.r2_score:.4f}")
            print(f"  R² Score (Circuit): {result.r2_circuit:.4f}")
            print(f"  MSE (Single Gate): {result.mse:.6f}")
            print(f"  MSE (Circuit): {result.mse_circuit:.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}")
            print(f"  Circuit Size: {result.circuit_size} gates")
            if self.period_detected:
                print(f"  Periodicity: Detected (period={self.period_detected})")
            print("=" * 70)

        return result
    
    # ============================================================
    # PREDICTION
    # ============================================================
    
    def predict(self,
                X: np.ndarray,
                mode: str = "circuit",
                return_uncertainty: bool = False) -> np.ndarray:
        """
        Make predictions using fitted function gates

        Args:
            X: Input features
            mode: Prediction mode:
                  - "circuit": Use best circuit (weighted combination, default)
                  - "ensemble": Use all active gates (full superposition)
                  - "single": Use best single gate only (legacy)
            return_uncertainty: Also return prediction uncertainty

        Returns:
            Predictions (and optionally uncertainty)
        """
        if not self.fitted:
            raise ValueError("Model must be fitted before prediction")

        if mode == "single":
            # Use best gate only (legacy behavior)
            best_idx = max(self.active_indices,
                          key=lambda x: self.gates[x].weight)
            y_pred = self.gates[best_idx].function(X)
        elif mode == "circuit":
            # Use best circuit (top K gates)
            best_circuit = self.get_best_circuit()
            registry = {self.gates[i].gate_id: self.gates[i] for i in self.active_indices}
            y_pred = best_circuit.evaluate(X, registry)
        else:  # mode == "ensemble"
            # Weighted ensemble (all active gates in 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 circuit
        X_smooth = np.linspace(self.X_train.min(), self.X_train.max(), 200)
        y_pred_circuit = self.predict(X_smooth, mode="circuit")
        best_circuit = self.get_best_circuit()
        plt.plot(X_smooth, y_pred_circuit, 'r-', linewidth=2,
                label=f'Circuit ({len(best_circuit.gates)} gates)')
        
        # Plot single gate for comparison
        y_pred_single = self.predict(X_smooth, mode="single")
        plt.plot(X_smooth, y_pred_single, 'g--', linewidth=1.5, alpha=0.7,
                label=f'Best Single: {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 (Circuit vs Single Gate)')
        plt.legend()
        plt.grid(True, alpha=0.3)

        # Plot residuals
        plt.subplot(1, 2, 2)
        y_train_pred_circuit = self.predict(self.X_train, mode="circuit")
        y_train_pred_single = self.predict(self.X_train, mode="single")
        residuals_circuit = self.y_train - y_train_pred_circuit
        residuals_single = self.y_train - y_train_pred_single
        
        plt.scatter(y_train_pred_circuit, residuals_circuit, alpha=0.5, 
                   color='red', label='Circuit residuals')
        plt.scatter(y_train_pred_single, residuals_single, alpha=0.5,
                   color='green', label='Single gate residuals')
        plt.axhline(y=0, color='r', linestyle='--')
        plt.xlabel('Predicted')
        plt.ylabel('Residual')
        plt.title('Residual Plot (Circuit vs Single)')
        plt.legend()
        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()
    
    def plot_circuit_composition(self):
        """Plot the best circuit showing gate contributions"""
        if not self.fitted:
            return
        
        best_circuit = self.get_best_circuit()
        
        if len(best_circuit.gates) == 0:
            print("No circuit found")
            return
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        # Plot 1: Circuit gate weights
        gate_ids = [g[0] for g in best_circuit.gates]
        gate_weights = [g[1] for g in best_circuit.gates]
        
        colors = plt.cm.RdYlGn(np.linspace(0.2, 0.8, len(gate_ids)))
        bars = axes[0].bar(range(len(gate_ids)), gate_weights, color=colors)
        
        # Add value labels on bars
        for bar, weight in zip(bars, gate_weights):
            axes[0].text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,
                        f'{weight:.3f}', ha='center', va='bottom', fontsize=9)
        
        axes[0].set_xticks(range(len(gate_ids)))
        axes[0].set_xticklabels(gate_ids, rotation=45, ha='right')
        axes[0].set_xlabel('Gate')
        axes[0].set_ylabel('Normalized Weight')
        axes[0].set_title(f'Circuit Composition ({len(gate_ids)} gates)')
        axes[0].grid(True, alpha=0.3, axis='y')
        
        # Plot 2: Individual gate contributions to the fit
        X_smooth = np.linspace(self.X_train.min(), self.X_train.max(), 200)
        
        axes[1].scatter(self.X_train, self.y_train, alpha=0.3, label='Data', color='black', s=10)
        
        for gate_id, weight in best_circuit.gates:
            # Find the gate in registry
            for idx in self.active_indices:
                if self.gates[idx].gate_id == gate_id:
                    y_gate = self.gates[idx].function(X_smooth) * weight
                    axes[1].plot(X_smooth, y_gate, '--', alpha=0.6, 
                               linewidth=1.5, label=f'{gate_id} (×{weight:.3f})')
                    break
        
        # Plot total circuit
        y_circuit = self.predict(X_smooth, mode="circuit")
        axes[1].plot(X_smooth, y_circuit, 'r-', linewidth=2.5, label='Circuit Total', alpha=0.9)
        
        axes[1].set_xlabel('X')
        axes[1].set_ylabel('y')
        axes[1].set_title('Circuit Gate Contributions')
        axes[1].legend(fontsize=7, loc='best')
        axes[1].grid(True, alpha=0.3)
        
        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)
        best_circuit = self.get_best_circuit()

        return {
            "status": "FITTED",
            "best_gate": self.gates[best_idx].gate_id,
            "best_circuit": str(best_circuit),
            "circuit_size": len(best_circuit.gates),
            "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,
        }