import numpy as np
import hashlib
import zlib
import bz2
import lzma
from pathlib import Path
from typing import List, Dict, Tuple
from dataclasses import dataclass, field

# ============================================================
# DATA STRUCTURES
# ============================================================
@dataclass
class CompressionAlgorithm:
    """Represents one of the 100+ potential algorithms"""
    name: str
    id: int
    category: str  # 'dictionary', 'statistical', 'transform', 'semantic'
    cost_weight: float  # Compute energy cost
    active: bool = True
    success_history: List[float] = field(default_factory=list)
    weight: float = 0.5  # SuperBoolean probability weight

@dataclass
class ChunkResult:
    """Result of compressing a single chunk"""
    chunk_id: int
    algo_id: int
    original_size: int
    compressed_size: int
    entropy_before: float
    entropy_after: float
    checksum_part: str

# ============================================================
# CCT COMPRESSION ENGINE
# ============================================================
class CCT_Compression_Engine:
    """
    Conditional Collapse Theory Compression Engine
    Uses SuperBoolean selection + ODE trajectory + Memory Pruning
    """
    def __init__(self,
                 num_algorithms: int = 100,
                 chunk_size: int = 4096,
                 prune_threshold: float = 0.1,
                 collapse_threshold: float = 0.01):
        
        self.CHUNK_SIZE = chunk_size
        self.PRUNE_THRESHOLD = prune_threshold
        self.COLLAPSE_THRESHOLD = collapse_threshold
        
        # Initialize 100 Algorithms (Simulated subset for demo)
        self.algorithms = self._initialize_algorithm_manifold(num_algorithms)
        self.active_algos = set(range(num_algorithms))
        
        # 16-Element State Vector
        self.state_vector = np.zeros(16)
        self.entropy_trajectory = []
        self.algorithm_sequence = []
        self.total_original_size = 0
        self.total_compressed_size = 0
        
    def _initialize_algorithm_manifold(self, n: int) -> List[CompressionAlgorithm]:
        """Create the SuperBoolean Manifold of Compressors"""
        algos = []
        categories = ['dictionary', 'statistical', 'transform', 'semantic', 'hybrid']
        for i in range(n):
            algos.append(CompressionAlgorithm(
                name=f"Algo_{i}_{categories[i % len(categories)]}",
                id=i,
                category=categories[i % len(categories)],
                cost_weight=np.random.uniform(0.1, 1.0)
            ))
        return algos

    def calculate_entropy(self, data: bytes) -> float:
        """Calculate Shannon Entropy of chunk (E01)"""
        if not data: return 0.0
        freq = np.bincount(np.frombuffer(data, dtype=np.uint8))
        p = freq / np.sum(freq)
        p = p[p > 0]
        return -np.sum(p * np.log2(p))

    def superboolean_select(self, chunk: bytes) -> int:
        """
        SuperBoolean Selection: Collapse algorithm superposition based on chunk content
        """
        # 1. Analyze Chunk Semantic Signature (Simulated)
        chunk_entropy = self.calculate_entropy(chunk)
        self.state_vector[0] = chunk_entropy / 8.0  # Normalize
        
        # 2. Update Algorithm Weights (ODE Dynamics)
        # Algorithms that match the entropy profile gain weight
        for idx in self.active_algos:
            algo = self.algorithms[idx]
            # Simulate compatibility score
            compatibility = np.random.uniform(0.1, 1.0) 
            if algo.category == 'statistical' and chunk_entropy < 4.0:
                compatibility *= 1.5 # Boost for low entropy
            
            # ODE Weight Update: dw = alpha * compatibility - decay
            algo.weight += 0.01 * compatibility - 0.001 * algo.weight
            algo.weight = np.clip(algo.weight, 0, 1)
        
        # 3. Collapse to Best Algorithm
        best_algo = max(self.active_algos, key=lambda i: self.algorithms[i].weight)
        return best_algo

    def compress_chunk(self, chunk: bytes, algo_id: int) -> Tuple[bytes, float]:
        """Execute the collapsed algorithm"""
        algo = self.algorithms[algo_id]
        
        # Simulate compression using standard libs mapped to algo categories
        if 'dictionary' in algo.name:
            compressed = zlib.compress(chunk)
        elif 'statistical' in algo.name:
            compressed = bz2.compress(chunk)
        elif 'transform' in algo.name:
            compressed = lzma.compress(chunk)
        else:
            compressed = chunk # No compression
        
        ratio = len(compressed) / len(chunk) if len(chunk) > 0 else 1.0
        return compressed, ratio

    def prune_algorithms(self):
        """Entropy-Gated Forgetting of Poor Algorithms (Memory Pruning)"""
        to_prune = []
        for idx in self.active_algos:
            algo = self.algorithms[idx]
            # Prune if weight is low (low collapse potential)
            if algo.weight < self.PRUNE_THRESHOLD:
                to_prune.append(idx)
        
        for idx in to_prune:
            self.active_algos.remove(idx)
            self.algorithms[idx].active = False
            
        if to_prune:
            print(f"  [PRUNE] Removed {len(to_prune)} low-performance algorithms")

    def run(self, file_path: str, verbose: bool = True) -> Dict:
        """
        Main CCT Compression Loop
        """
        if verbose:
            print("="*70)
            print("CCT COMPRESSION ENGINE: SEMANTIC ENTROPY COLLAPSE")
            print("="*70)
            
        data = Path(file_path).read_bytes()
        self.total_original_size = len(data)
        
        chunks = [data[i:i+self.CHUNK_SIZE] for i in range(0, len(data), self.CHUNK_SIZE)]
        compressed_data = b''
        algo_sequence = []
        
        initial_entropy = self.calculate_entropy(data)
        
        for t, chunk in enumerate(chunks):
            # 1. SuperBoolean Selection (Collapse)
            selected_algo = self.superboolean_select(chunk)
            
            # 2. Execute Compression (Work Investment)
            compressed_chunk, ratio = self.compress_chunk(chunk, selected_algo)
            
            # 3. Update Metrics
            self.algorithm_sequence.append(selected_algo)
            compressed_data += compressed_chunk
            self.total_compressed_size += len(compressed_chunk)
            
            # 4. Feedback Loop (Update Algorithm Weights)
            # Good ratio -> Increase weight
            self.algorithms[selected_algo].success_history.append(ratio)
            self.algorithms[selected_algo].weight += 0.05 * (1.0 - ratio)
            
            # 5. Track Entropy Trajectory
            current_entropy = self.calculate_entropy(compressed_chunk)
            self.entropy_trajectory.append(current_entropy)
            
            # 6. Memory Pruning (Every 10 chunks)
            if t % 10 == 0:
                self.prune_algorithms()
                
            if verbose and t % 50 == 0:
                print(f"  Chunk {t}: Algo={self.algorithms[selected_algo].name} | Ratio={ratio:.2f} | Active Algos={len(self.active_algos)}")
        
        # 7. Final Checksum (Proof Stability - E16)
        full_signature = compressed_data + bytes(self.algorithm_sequence)
        final_checksum = hashlib.sha256(full_signature).hexdigest()
        
        # 8. Calculate Global Metrics
        global_ratio = self.total_compressed_size / self.total_original_size
        entropy_reduction = (initial_entropy - np.mean(self.entropy_trajectory)) / initial_entropy
        
        result = {
            "status": "COLLAPSED",
            "original_size": self.total_original_size,
            "compressed_size": self.total_compressed_size,
            "compression_ratio": global_ratio,
            "entropy_reduction": entropy_reduction,
            "checksum": final_checksum,
            "algorithm_sequence_length": len(self.algorithm_sequence),
            "active_algorithms_remaining": len(self.active_algos),
            "entropy_trajectory": self.entropy_trajectory
        }
        
        if verbose:
            print("-"*70)
            print(f"FINAL COMPRESSION METRICS:")
            print(f"  Original Size: {self.total_original_size} bytes")
            print(f"  Compressed Size: {self.total_compressed_size} bytes")
            print(f"  Compression Ratio: {global_ratio:.4f}")
            print(f"  Entropy Reduction: {entropy_reduction*100:.1f}%")
            print(f"  Final Checksum: {final_checksum[:16]}...")
            print(f"  Algorithms Pruned: {100 - len(self.active_algos)}")
            print("="*70)
            
        return result