import torch
import torch.nn.functional as F
from scipy import ndimage
import numpy as np

class VisualComplexityTokenizer:
    def __init__(self, target_tokens=96, patch_size=14):
        self.target_tokens = target_tokens
        self.patch_size = patch_size
        self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32)
        self.sobel_y = self.sobel_x.T
        
    def compute_edge_complexity(self, patch):
        """Compute Sobel edge magnitude"""
        gray = torch.mean(patch, dim=0)
        edge_x = F.conv2d(gray.unsqueeze(0).unsqueeze(0), 
                          self.sobel_x.unsqueeze(0).unsqueeze(0), padding=1)
        edge_y = F.conv2d(gray.unsqueeze(0).unsqueeze(0), 
                          self.sobel_y.unsqueeze(0).unsqueeze(0), padding=1)
        magnitude = torch.sqrt(edge_x**2 + edge_y**2)
        return magnitude.mean().item()
    
    def compute_frequency_complexity(self, patch):
        """Compute high-frequency energy using FFT"""
        gray = torch.mean(patch, dim=0).numpy()
        fft = np.fft.fft2(gray)
        fft_shift = np.fft.fftshift(fft)
        
        # High-frequency energy (outer regions)
        h, w = fft_shift.shape
        mask = np.ones_like(fft_shift)
        center = (h//2, w//2)
        radius = min(h, w) // 4
        y, x = np.ogrid[:h, :w]
        mask[(x - center[1])**2 + (y - center[0])**2 <= radius**2] = 0
        
        high_freq_energy = np.sum(np.abs(fft_shift * mask)**2)
        return high_freq_energy / (h * w)
    
    def compute_entropy(self, patch):
        """Compute Shannon entropy"""
        gray = torch.mean(patch, dim=0).numpy()
        hist, _ = np.histogram(gray, bins=256, range=(0, 1))
        hist = hist[hist > 0]  # Remove zero entries
        probs = hist / hist.sum()
        entropy = -np.sum(probs * np.log2(probs))
        return entropy
    
    def allocate_tokens(self, image, alpha=0.4, beta=0.3, gamma=0.3):
        """Main token allocation function"""
        B, C, H, W = image.shape
        
        # Extract patches
        patches = F.unfold(image, kernel_size=self.patch_size, 
                          stride=self.patch_size)
        patches = patches.reshape(B, C, self.patch_size, self.patch_size, -1)
        patches = patches.permute(0, 4, 1, 2, 3)  # [B, N_patches, C, H, W]
        
        complexities = []
        for i in range(patches.shape[1]):
            patch = patches[0, i]  # Process first batch item
            
            # Compute complexity metrics
            edge_c = self.compute_edge_complexity(patch)
            freq_c = self.compute_frequency_complexity(patch)
            entropy_c = self.compute_entropy(patch)
            
            # Weighted combination
            total_complexity = alpha * edge_c + beta * freq_c + gamma * entropy_c
            complexities.append(total_complexity)
        
        # Convert to probabilities
        complexities = torch.tensor(complexities)
        probs = F.softmax(complexities * 2.0, dim=0)  # Temperature scaling
        
        # Sample tokens based on complexity
        selected_indices = torch.multinomial(probs, 
                                            num_samples=self.target_tokens, 
                                            replacement=True)
        
        # Get selected patches
        selected_patches = patches[0, selected_indices]
        
        return selected_patches, selected_indices, complexities

# Usage example
tokenizer = VisualComplexityTokenizer(target_tokens=96)
image = torch.randn(1, 3, 224, 224)  # Batch of images

selected_tokens, indices, complexity_map = tokenizer.allocate_tokens(image)
print(f"Selected {len(indices)} tokens from {196} patches")
print(f"Complexity range: {complexity_map.min():.2f} - {complexity_map.max():.2f}")
print(f"Token reduction: {(1 - len(indices)/196)*100:.1f}%")