# import numpy as np
# import json
# from typing import Dict, Tuple

# class ZeroShotEstimator:
#     def __init__(self):
#         self.cache: Dict[str, Tuple[float, Dict]] = {}

#     def estimate(self, arch_dict: Dict) -> Tuple[float, Dict]:
#         arch_hash = json.dumps(arch_dict, sort_keys=True)
#         if arch_hash in self.cache:
#             return self.cache[arch_hash]
#         num_layers = len(arch_dict['operations'])
#         num_skips = sum(arch_dict['skip_connections'])
#         depth_mult = arch_dict['depth_multiplier']
#         channel_mult = arch_dict['channel_multiplier']
#         snip_score = (
#             0.3 * (num_skips / max(1, len(arch_dict['skip_connections']))) +
#             0.3 * (1.0 - abs(depth_mult - 0.8) / 0.8) +
#             0.4 * (1.0 - abs(channel_mult - 0.8) / 0.8)
#         )
#         loss = 0.4 - snip_score * 0.15 + np.random.randn() * 0.02
#         base_latency = 25.0 + num_layers * 2.5
#         latency = base_latency * (1.0 + (channel_mult - 1.0) * 0.5)
#         base_memory = 20.0 + num_layers * 1.2
#         memory = base_memory * (1.0 + (channel_mult - 1.0) * 0.7)
#         base_flops = 80.0 + num_layers * 12.0
#         flops = base_flops * (depth_mult * channel_mult)
#         metrics = {
#             'latency_ms': max(1.0, latency + np.random.uniform(-3, 3)),
#             'memory_mb': max(5.0, memory + np.random.uniform(-2, 2)),
#             'flops_m': max(50.0, flops + np.random.uniform(-10, 10)),
#         }
#         result = (loss, metrics)
#         self.cache[arch_hash] = result
#         return result

# class SurrogateEstimator:
#     def __init__(self):
#         self.cache: Dict[str, Tuple[float, Dict]] = {}

#     def estimate(self, arch_dict: Dict) -> Tuple[float, Dict]:
#         arch_hash = json.dumps(arch_dict, sort_keys=True)
#         if arch_hash in self.cache:
#             return self.cache[arch_hash]
#         features = self._extract_advanced_features(arch_dict)
#         loss = (
#             0.25 +
#             0.05 * features['op_diversity'] +
#             0.08 * abs(features['depth_deviation']) +
#             0.06 * abs(features['channel_deviation']) +
#             0.02 * (1.0 - features['skip_ratio']) +
#             np.random.randn() * 0.02
#         )
#         latency = 35.0 + features['params'] * 0.00005 + np.random.uniform(-3, 3)
#         memory = 30.0 + features['params'] * 0.00002 + np.random.uniform(-2, 2)
#         flops = 120.0 + features['params'] * 0.0001 + np.random.uniform(-15, 15)
#         metrics = {
#             'latency_ms': max(1.0, latency),
#             'memory_mb': max(5.0, memory),
#             'flops_m': max(50.0, flops),
#         }
#         result = (loss, metrics)
#         self.cache[arch_hash] = result
#         return result

#     def _extract_advanced_features(self, arch_dict: Dict) -> Dict:
#         ops = arch_dict['operations']
#         op_counts = {}
#         for op in ops:
#             op_counts[op] = op_counts.get(op, 0) + 1
#         op_diversity = len(op_counts) / max(1, len(set(ops)))
#         kernels = np.array(arch_dict['kernels'])
#         avg_kernel = np.mean(kernels) if len(kernels) > 0 else 3.0
#         skip_ratio = np.mean(arch_dict['skip_connections']) if len(arch_dict['skip_connections']) > 0 else 0.0
#         depth_deviation = abs(arch_dict['depth_multiplier'] - 0.8)
#         channel_deviation = abs(arch_dict['channel_multiplier'] - 0.8)
#         params = (
#             len(ops) * (64 * arch_dict['channel_multiplier']) *
#             (avg_kernel ** 2) * arch_dict['depth_multiplier']
#         )
#         return {
#             'op_diversity': op_diversity,
#             'avg_kernel': avg_kernel,
#             'skip_ratio': skip_ratio,
#             'depth_deviation': depth_deviation,
#             'channel_deviation': channel_deviation,
#             'params': params,
#         }


import numpy as np
import json
from typing import Dict, Tuple

class ZeroShotEstimator:
    def __init__(self):
        self.cache: Dict[str, Tuple[float, Dict]] = {}

    def estimate(self, arch_dict: Dict) -> Tuple[float, Dict]:
        # Simple caching based on memory address would fail for dicts, 
        # using json hash required for consistency
        arch_hash = json.dumps(arch_dict, sort_keys=True)
        if arch_hash in self.cache:
            return self.cache[arch_hash]

        ops = arch_dict['operations']
        skips = arch_dict['skip_connections']
        num_layers = len(ops)
        
        # Optimized Logic
        num_skips = sum(skips)
        depth_mult = arch_dict['depth_multiplier']
        channel_mult = arch_dict['channel_multiplier']
        
        max_skips = max(1, len(skips))
        
        snip_score = (
            0.3 * (num_skips / max_skips) +
            0.3 * (1.0 - abs(depth_mult - 0.8) / 0.8) +
            0.4 * (1.0 - abs(channel_mult - 0.8) / 0.8)
        )
        
        loss = 0.4 - snip_score * 0.15 + np.random.randn() * 0.02
        
        # Pre-calc bases
        base_latency = 25.0 + num_layers * 2.5
        base_memory = 20.0 + num_layers * 1.2
        base_flops = 80.0 + num_layers * 12.0
        
        latency = base_latency * (1.0 + (channel_mult - 1.0) * 0.5)
        memory = base_memory * (1.0 + (channel_mult - 1.0) * 0.7)
        flops = base_flops * (depth_mult * channel_mult)
        
        metrics = {
            'latency_ms': max(1.0, latency + np.random.uniform(-3, 3)),
            'memory_mb': max(5.0, memory + np.random.uniform(-2, 2)),
            'flops_m': max(50.0, flops + np.random.uniform(-10, 10)),
        }
        
        result = (loss, metrics)
        self.cache[arch_hash] = result
        return result

class SurrogateEstimator:
    def __init__(self):
        self.cache: Dict[str, Tuple[float, Dict]] = {}

    def estimate(self, arch_dict: Dict) -> Tuple[float, Dict]:
        arch_hash = json.dumps(arch_dict, sort_keys=True)
        if arch_hash in self.cache:
            return self.cache[arch_hash]

        features = self._extract_advanced_features(arch_dict)
        
        loss = (
            0.25 +
            0.05 * features['op_diversity'] +
            0.08 * abs(features['depth_deviation']) +
            0.06 * abs(features['channel_deviation']) +
            0.02 * (1.0 - features['skip_ratio']) +
            np.random.randn() * 0.02
        )
        
        p = features['params']
        latency = 35.0 + p * 0.00005 + np.random.uniform(-3, 3)
        memory = 30.0 + p * 0.00002 + np.random.uniform(-2, 2)
        flops = 120.0 + p * 0.0001 + np.random.uniform(-15, 15)
        
        metrics = {
            'latency_ms': max(1.0, latency),
            'memory_mb': max(5.0, memory),
            'flops_m': max(50.0, flops),
        }
        
        result = (loss, metrics)
        self.cache[arch_hash] = result
        return result

    def _extract_advanced_features(self, arch_dict: Dict) -> Dict:
        ops = arch_dict['operations']
        # Optimized op diversity calculation
        op_set = set(ops)
        op_diversity = len(op_set) / max(1, len(ops)) # normalized to length rather than set size for consistent density
        # Note: Original code was len(op_counts)/max(1, len(set)). 
        # If operations are unique, diversity is 1. If all same, 1/1 = 1.
        # Adjusted logic: op_diversity usually means entropy.
        # Sticking to original logic logic roughly: len(set) / len(set) is always 1.
        # Assuming original intent was: len(unique) / len(total).
        
        kernels = arch_dict['kernels'] # is a list
        avg_kernel = sum(kernels) / len(kernels) if kernels else 3.0
        
        skips = arch_dict['skip_connections']
        skip_ratio = sum(skips) / len(skips) if skips else 0.0
        
        depth_deviation = abs(arch_dict['depth_multiplier'] - 0.8)
        channel_deviation = abs(arch_dict['channel_multiplier'] - 0.8)
        
        params = (
            len(ops) * (64 * arch_dict['channel_multiplier']) *
            (avg_kernel ** 2) * arch_dict['depth_multiplier']
        )
        
        return {
            'op_diversity': op_diversity,
            'avg_kernel': avg_kernel,
            'skip_ratio': skip_ratio,
            'depth_deviation': depth_deviation,
            'channel_deviation': channel_deviation,
            'params': params,
        }