import hashlib
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from collections import defaultdict
import numpy as np

@dataclass
class CascadeNode:
    prefix: List[int]
    prefix_hash: str
    level: int
    parent: Optional['CascadeNode'] = None
    children: List['CascadeNode'] = None
    cached_state: Any = None
    computation_cost: float = 0.0

    def __post_init__(self):
        if self.children is None:
            self.children = []

class CascadePrefixCache:
    def __init__(self, min_prefix_length: int = 2, max_cache_size: int = 10000):
        self.min_prefix_length = min_prefix_length
        self.max_cache_size = max_cache_size

        self.prefix_cache: Dict[str, Any] = {}
        self.computation_graph: Dict[str, CascadeNode] = {}

        self.cache_hits = 0
        self.cache_misses = 0

    def _compute_prefix_hash(self, prefix: List[int]) -> str:
        return hashlib.md5(str(prefix).encode()).hexdigest()

    def _build_cascade_hierarchy(self, queries: List[List[int]]) -> Dict[int, List[CascadeNode]]:
        prefix_map = {}
        level_map = defaultdict(list)
        all_prefixes = set()

        for query in queries:
            for i in range(self.min_prefix_length, len(query) + 1):
                prefix = tuple(query[:i])
                all_prefixes.add(prefix)

        for prefix_tuple in sorted(all_prefixes, key=len):
            prefix = list(prefix_tuple)
            prefix_hash = self._compute_prefix_hash(prefix)
            level = len(prefix)

            node = CascadeNode(
                prefix=prefix,
                prefix_hash=prefix_hash,
                level=level
            )

            if level > self.min_prefix_length:
                parent_prefix = prefix[:-1]
                parent_hash = self._compute_prefix_hash(parent_prefix)
                if parent_hash in prefix_map:
                    parent_node = prefix_map[parent_hash]
                    node.parent = parent_node
                    parent_node.children.append(node)

            prefix_map[prefix_hash] = node
            level_map[level].append(node)
            self.computation_graph[prefix_hash] = node

        return dict(level_map)

    def _compute_prefix_state(self, prefix: List[int], parent_state: Any = None) -> Any:
        if parent_state is not None:
            new_tokens = prefix[len(parent_state.get('processed_tokens', [])):]
            return {
                'processed_tokens': prefix,
                'embeddings': parent_state.get('embeddings', []) + [f"emb_{t}" for t in new_tokens],
                'hidden_state': f"hidden_state_{len(prefix)}",
                'attention_weights': np.random.rand(len(prefix), len(prefix)).tolist()
            }
        else:
            return {
                'processed_tokens': prefix,
                'embeddings': [f"emb_{t}" for t in prefix],
                'hidden_state': f"hidden_state_{len(prefix)}",
                'attention_weights': np.random.rand(len(prefix), len(prefix)).tolist()
            }

    def _cascade_compute(self, node: CascadeNode) -> Any:
        prefix_hash = node.prefix_hash

        if prefix_hash in self.prefix_cache:
            self.cache_hits += 1
            return self.prefix_cache[prefix_hash]

        self.cache_misses += 1

        parent_state = node.parent.cached_state if node.parent and node.parent.cached_state else None
        computed_state = self._compute_prefix_state(node.prefix, parent_state)

        node.cached_state = computed_state
        self.prefix_cache[prefix_hash] = computed_state

        if len(self.prefix_cache) > self.max_cache_size:
            self._evict_cache_entries()

        return computed_state

    def _evict_cache_entries(self):
        items_to_remove = len(self.prefix_cache) // 5
        keys_to_remove = list(self.prefix_cache.keys())[:items_to_remove]
        for key in keys_to_remove:
            del self.prefix_cache[key]

    def batch_process_cascade(self, queries: List[List[int]]) -> List[Any]:
        hierarchy = self._build_cascade_hierarchy(queries)
        cascade_results = {}

        for level in sorted(hierarchy.keys()):
            nodes = hierarchy[level]
            for node in nodes:
                result = self._cascade_compute(node)
                cascade_results[node.prefix_hash] = result

        final_results = []
        for query in queries:
            query_hash = self._compute_prefix_hash(query)
            if query_hash in cascade_results:
                final_results.append(cascade_results[query_hash])
            else:
                result = self._compute_prefix_state(query)
                final_results.append(result)

        return final_results

    def optimized_batch_process(self, ordered_queries: List[List[int]]) -> List[Any]:
        results = []
        previous_prefix = []

        for query in ordered_queries:
            common_len = 0
            min_len = min(len(previous_prefix), len(query))
            for j in range(min_len):
                if previous_prefix[j] == query[j]:
                    common_len += 1
                else:
                    break

            base_state = None
            if common_len >= self.min_prefix_length:
                common_prefix = query[:common_len]
                common_hash = self._compute_prefix_hash(common_prefix)
                if common_hash in self.prefix_cache:
                    base_state = self.prefix_cache[common_hash]
                    self.cache_hits += 1
                else:
                    self.cache_misses += 1
            else:
                self.cache_misses += 1

            result = self._compute_prefix_state(query, base_state)

            # Cache full query and intermediate prefixes
            for i in range(self.min_prefix_length, len(query) + 1):
                sub_prefix = query[:i]
                sub_hash = self._compute_prefix_hash(sub_prefix)
                if sub_hash not in self.prefix_cache:
                    sub_state = self._compute_prefix_state(sub_prefix)
                    self.prefix_cache[sub_hash] = sub_state

            results.append(result)
            previous_prefix = query

        return results

    def get_statistics(self) -> Dict[str, Any]:
        total_requests = self.cache_hits + self.cache_misses
        hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0
        return {
            'cache_hits': self.cache_hits,
            'cache_misses': self.cache_misses,
            'hit_rate': hit_rate,
            'cache_size': len(self.prefix_cache),
            'computation_graph_size': len(self.computation_graph)
        }

class MemoryEfficientCascadeCache(CascadePrefixCache):
    def __init__(self, min_prefix_length: int = 2, max_cache_size: int = 10000,
                 memory_threshold: float = 0.8):
        super().__init__(min_prefix_length, max_cache_size)
        self.memory_threshold = memory_threshold
        self.access_counts = defaultdict(int)
        self.access_times = {}
        self.current_time = 0

    def _smart_eviction(self):
        if len(self.prefix_cache) <= self.max_cache_size * self.memory_threshold:
            return

        scores = {}
        for key in self.prefix_cache:
            access_count = self.access_counts[key]
            last_access = self.access_times.get(key, 0)
            recency = self.current_time - last_access
            scores[key] = access_count / (1 + recency)

        items_to_remove = int(len(self.prefix_cache) * 0.2)
        keys_to_remove = sorted(scores, key=scores.get)[:items_to_remove]
        for key in keys_to_remove:
            self.prefix_cache.pop(key, None)
            self.access_counts.pop(key, None)
            self.access_times.pop(key, None)

    def _cascade_compute(self, node: CascadeNode) -> Any:
        self.current_time += 1
        prefix_hash = node.prefix_hash

        if prefix_hash in self.prefix_cache:
            self.cache_hits += 1
            self.access_counts[prefix_hash] += 1
            self.access_times[prefix_hash] = self.current_time
            return self.prefix_cache[prefix_hash]

        self.cache_misses += 1
        parent_state = node.parent.cached_state if node.parent and node.parent.cached_state else None
        computed_state = self._compute_prefix_state(node.prefix, parent_state)

        node.cached_state = computed_state
        self.prefix_cache[prefix_hash] = computed_state
        self.access_counts[prefix_hash] = 1
        self.access_times[prefix_hash] = self.current_time

        self._smart_eviction()
        return computed_state

    def optimized_batch_process(self, ordered_queries: List[List[int]]) -> List[Any]:
        results = []
        previous_prefix = []

        for query in ordered_queries:
            self.current_time += 1
            common_len = 0
            min_len = min(len(previous_prefix), len(query))
            for j in range(min_len):
                if previous_prefix[j] == query[j]:
                    common_len += 1
                else:
                    break

            base_state = None
            if common_len >= self.min_prefix_length:
                common_prefix = query[:common_len]
                common_hash = self._compute_prefix_hash(common_prefix)
                if common_hash in self.prefix_cache:
                    base_state = self.prefix_cache[common_hash]
                    self.cache_hits += 1
                    self.access_counts[common_hash] += 1
                    self.access_times[common_hash] = self.current_time
                else:
                    self.cache_misses += 1
            else:
                self.cache_misses += 1

            result = self._compute_prefix_state(query, base_state)

            for i in range(self.min_prefix_length, len(query) + 1):
                sub_prefix = query[:i]
                sub_hash = self._compute_prefix_hash(sub_prefix)
                if sub_hash not in self.prefix_cache:
                    sub_state = self._compute_prefix_state(sub_prefix)
                    self.prefix_cache[sub_hash] = sub_state
                    self.access_counts[sub_hash] = 1
                    self.access_times[sub_hash] = self.current_time

            query_hash = self._compute_prefix_hash(query)
            self.prefix_cache[query_hash] = result
            self.access_counts[query_hash] = 1
            self.access_times[query_hash] = self.current_time

            self._smart_eviction()
            results.append(result)
            previous_prefix = query

        return results

def demonstrate_cascade_cache():
    queries = [
        [1, 2, 3, 4, 5],
        [1, 2, 3, 6, 7],
        [1, 2, 8, 9, 0],
        [10, 11, 0, 9, 8],
        [10, 11, 0, 20, 21]
    ]

    print("=== Basic Cascade Cache ===")
    cache = CascadePrefixCache(min_prefix_length=2)
    results1 = cache.batch_process_cascade(queries)
    print(f"Cascade results: {len(results1)} queries processed")
    print(f"Statistics: {cache.get_statistics()}")

    print("\n=== Optimized Cascade Cache ===")
    cache2 = CascadePrefixCache(min_prefix_length=2)
    results2 = cache2.optimized_batch_process(queries)
    print(f"Optimized results: {len(results2)} queries processed")
    print(f"Statistics: {cache2.get_statistics()}")

    print("\n=== Memory Efficient Cascade Cache ===")
    cache3 = MemoryEfficientCascadeCache(min_prefix_length=2, max_cache_size=100)
    for batch_num in range(3):
        batch_queries = [q + [i + batch_num] for i, q in enumerate(queries)]
        results3 = cache3.optimized_batch_process(batch_queries)
        print(f"Batch {batch_num + 1} processed, Statistics: {cache3.get_statistics()}")

if __name__ == "__main__":
    demonstrate_cascade_cache()
