import hashlib
import time
from tgi.tgi import TGICognitiveKernel

class TopologicalCompressionEngine:
    """
    TGI Advanced Compression v1.0.
    Transforms data into a (Start Coordinate + Spike Sequence).
    Data is logically sharded along the manifold's Hamiltonian path.
    """
    def __init__(self, m=256, k=3):
        self.m = m
        self.k = k
        self.kernel = TGICognitiveKernel(m, k)
        self.storage_manifold = {} # {coord: shard}

    def _generate_stateless_gates(self):
        # Local logic gates for O(1) routing
        level = [[(0,1,2) for _ in range(self.m)] for _ in range(self.m)]
        identity, swap12, swap01 = (0, 1, 2), (0, 2, 1), (1, 0, 2)
        swap02 = lambda p: (p[2], p[1], p[0])
        P = [identity] * self.m
        if self.m >= 2: P[self.m-2], P[self.m-1] = swap12, swap01
        for s in range(self.m):
            for j in range(self.m):
                level[s][j] = swap02(P[s]) if j == 0 and s != (self.m - 2) else P[s]
        return level

    def compress_data(self, label, data_string, shard_size=32):
        """
        GEOMETRIC COMPRESSION:
        Saves only the start coordinate. The 'Spike' (gates) reconstructs the path.
        """
        # 1. Start Coordinate derived from label
        h = hashlib.sha256(label.encode()).digest()
        start_coord = tuple(h[i] % self.m for i in range(self.k))

        shards = [data_string[i:i+shard_size] for i in range(0, len(data_string), shard_size)]
        gates = self._generate_stateless_gates()

        curr = start_coord
        print(f"[*] Compressing '{label}' -> Starting at {start_coord} across {len(shards)} shards.")

        for i, shard in enumerate(shards):
            self.storage_manifold[curr] = shard

            # Step via O(1) Gates
            s = sum(curr) % self.m
            direction = gates[s][curr[0]][0]
            next_coord = list(curr)
            next_coord[direction] = (next_coord[direction] + 1) % self.m
            curr = tuple(next_coord)

        return start_coord, len(shards)

    def reconstruct_data(self, start_coord, num_shards):
        """
        O(1) Reconstruction Logic: Follows the breadcrumbs.
        """
        start_time = time.perf_counter()
        gates = self._generate_stateless_gates()

        reconstructed = ""
        curr = start_coord

        for _ in range(num_shards):
            shard = self.storage_manifold.get(curr)
            if not shard: break
            reconstructed += shard

            s = sum(curr) % self.m
            direction = gates[s][curr[0]][0]
            next_coord = list(curr)
            next_coord[direction] = (next_coord[direction] + 1) % self.m
            curr = tuple(next_coord)

        latency = (time.perf_counter() - start_time) * 1000
        return reconstructed, latency

if __name__ == "__main__":
    tce = TopologicalCompressionEngine(m=256, k=3)
    data = "ALGERIA_SOVEREIGN_INFRASTRUCTURE_v2.0_CODE_BLOCK_AGI"
    coord, count = tce.compress_data("agi_core", data)
    out, lat = tce.reconstruct_data(coord, count)
    print(f"[*] Reconstructed in {lat:.4f} ms: {out}")