
import numpy as np

class SemanticAnalyzer:
    """
    Analyzes the semantic behavior of symbolic trees.
    Computes semantic signatures, distances, and diversity scores.
    """
    def __init__(self, X_data, sample_size=256, random_state=None):
        self.rng = np.random.default_rng(random_state)
        self.sample_size = min(sample_size, len(X_data))
        
        # 1. Select evaluation set S (fixed subset)
        # Using simple random sampling for now. 
        # Ideally robust quasi-random (Latin Hypercube) if possible, but random is sufficient for now.
        indices = self.rng.choice(len(X_data), size=self.sample_size, replace=False)
        self.S = X_data[indices]
        
        # Parameters for definitions
        self.epsilon = 1e-12
        self.lambda_inv = 1.0 # Penalty for invalidity
        self.k_nn = 5 # Neighbors for diversity score

    def compute_signature(self, tree):
        """
        Computes the raw signature vector v_T (evaluation on S).
        Returns raw vector and invalid fraction.
        """
        try:
            # Evaluate tree on S
            v = tree.evaluate(self.S)
            
            # Check for invalid values (NaN, Inf)
            # Mask them?
            # We treat them as invalid entries.
            # Convert to float array ensuring compatibility
            v = np.asarray(v, dtype=float)
            
            # Check validity
            is_valid = np.isfinite(v)
            valid_fraction = np.mean(is_valid)
            
            # Fill invalid with 0.0 or nan for now? 
            # Normalization needs to handle them.
            # Let's replace NaNs with median of valid? Or 0?
            # Standard practice: if completely invalid, special handling.
            
            if valid_fraction < 1e-6:
                # Totally invalid
                return np.zeros_like(v), 1.0 # full invalidity
                
            return v, 1.0 - valid_fraction
            
        except Exception:
            # Crash during evaluation
            return np.zeros(self.sample_size), 1.0

    def normalize_signature(self, v):
        """
        Computes robust Z-score: (v - median) / (MAD + epsilon).
        Handles NaNs by ignoring them in stats, then filling.
        """
        # Mask invalids for stats
        # Assuming v might have NaNs/Infs if we passed them through?
        # Let's clean v first just in case.
        v_clean = np.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0)
        
        med = np.median(v_clean)
        abs_dev = np.abs(v_clean - med)
        mad = np.median(abs_dev)
        
        v_norm = (v_clean - med) / (mad + self.epsilon)
        return v_norm

    def compute_distance(self, v1, v2, p1_inv=0.0, p2_inv=0.0):
        """
        Robust L1 distance with invalidity penalty.
        d = mean(|v1 - v2|) + lambda * (p1 + p2)
        Note: The user formula says d <- d + lambda * p.
        If both are partially invalid, we penalize both?
        Let's take p = max(p1, p2) or sum?
        Usually distance is between logic. If one is invalid, distance is high.
        Let's perform L1 on the vectors (which are normalized).
        """
        # Normalize first? Or assume passed normalized?
        # Signatures passed here should be normalized?
        # User blueprint says: 1. Evaluate (v), 2. Normalize (v_tilde), 3. Distance.
        # So we assume v1, v2 are ALREADY normalized.
        
        # Robust L1
        diff = np.abs(v1 - v2)
        # We might want to trim outliers in the distance calculation itself? 
        # User said "Robust L1... behaves well with outliers". 
        # So mean(|diff|) is fine.
        d = np.mean(diff)
        
        # Penalty
        # If p > 0, add penalty.
        # Use average invalidity?
        p_total = (p1_inv + p2_inv) / 2.0 # or sum
        # User: d <- d + lambda * p.
        # Let's use sum of invalid fractions as a safe upper bound penalty.
        d += self.lambda_inv * (p1_inv + p2_inv)
        
        return d

    def compute_diversity_score(self, target_signature, population_signatures):
        """
        Computes mean distance to k nearest neighbors in population.
        High score = high novelty/diversity.
        
        target_signature: tuple (normalized_vec, invalid_fraction)
        population_signatures: list of tuples
        """
        v_t, p_t = target_signature
        
        distances = []
        for v_pop, p_pop in population_signatures:
            dist = self.compute_distance(v_t, v_pop, p_t, p_pop)
            distances.append(dist)
            
        # Sort
        distances.sort()
        
        # kNN mean
        # Skip self if present? (Distance 0). 
        # Usually we compare candidate vs Archive or Population (excluding self).
        # Assuming population_signatures might include self, the first distance is 0.
        # We want diversity relative to OTHERS.
        # If dist=0 exists, it's likely self or identical clone.
        # We should take k smallest distances.
        
        if not distances:
            return 0.0
            
        k = min(self.k_nn, len(distances))
        if k == 0: return 0.0
        
        # k nearest
        nearest = distances[:k]
        return np.mean(nearest)
