# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# MIXING AUGMENTATIONS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def alphamix_data(x, y, alpha_range=(0.3, 0.7), spatial_ratio=0.25):
    """
    Standard AlphaMix: Single spatially localized transparent overlay.
    """
    batch_size = x.size(0)
    index = torch.randperm(batch_size, device=x.device)
    
    y_a, y_b = y, y[index]
    
    # Sample alpha from Beta distribution
    alpha_min, alpha_max = alpha_range
    beta_sample = torch.distributions.Beta(2.0, 2.0).sample().item()
    alpha = alpha_min + (alpha_max - alpha_min) * beta_sample
    
    # Compute overlay region
    _, _, H, W = x.shape
    overlay_ratio = torch.sqrt(torch.tensor(spatial_ratio)).item()
    overlay_h = int(H * overlay_ratio)
    overlay_w = int(W * overlay_ratio)
    
    top = torch.randint(0, H - overlay_h + 1, (1,), device=x.device).item()
    left = torch.randint(0, W - overlay_w + 1, (1,), device=x.device).item()
    
    # Blend
    composited_x = x.clone()
    overlay_region = alpha * x[:, :, top:top+overlay_h, left:left+overlay_w]
    background_region = (1 - alpha) * x[index, :, top:top+overlay_h, left:left+overlay_w]
    composited_x[:, :, top:top+overlay_h, left:left+overlay_w] = overlay_region + background_region
    
    return composited_x, y_a, y_b, alpha


def alphamix_fractal(
    x: torch.Tensor,
    y: torch.Tensor,
    alpha_range=(0.3, 0.7),
    steps_range=(1, 3),
    triad_scales=(1/3, 1/9, 1/27),
    beta_shape=(2.0, 2.0),
    seed: int | None = None,
):
    """
    Fractal AlphaMix: Triadic multi-patch overlays aligned to Cantor geometry.
    Pure torch, GPU-compatible.
    """
    if seed is not None:
        torch.manual_seed(seed)
    
    B, C, H, W = x.shape
    device = x.device
    
    # Permutation for mixing
    idx = torch.randperm(B, device=device)
    y_a, y_b = y, y[idx]
    
    x_mix = x.clone()
    total_area = H * W
    
    # Beta distribution for transparency sampling
    k1, k2 = beta_shape
    beta_dist = torch.distributions.Beta(k1, k2)
    alpha_min, alpha_max = alpha_range
    
    # Storage for effective alpha calculation
    alpha_elems = []
    area_weights = []
    
    # Sample number of patches (same for all images in batch)
    steps = torch.randint(steps_range[0], steps_range[1] + 1, (1,), device=device).item()
    
    for _ in range(steps):
        # Choose triadic scale
        scale_idx = torch.randint(0, len(triad_scales), (1,), device=device).item()
        scale = triad_scales[scale_idx]
        
        # Compute patch dimensions (triadic area)
        patch_area = max(1, int(total_area * scale))
        side = int(torch.sqrt(torch.tensor(patch_area, dtype=torch.float32)).item())
        h = max(1, min(H, side))
        w = max(1, min(W, side))
        
        # Random position
        top = torch.randint(0, H - h + 1, (1,), device=device).item()
        left = torch.randint(0, W - w + 1, (1,), device=device).item()
        
        # Sample transparency from Beta distribution
        alpha_raw = beta_dist.sample().item()
        alpha = alpha_min + (alpha_max - alpha_min) * alpha_raw
        
        # Track for effective alpha
        alpha_elems.append(alpha)
        area_weights.append(h * w)
        
        # Blend patches
        fg = alpha * x[:, :, top:top + h, left:left + w]
        bg = (1 - alpha) * x[idx, :, top:top + h, left:left + w]
        x_mix[:, :, top:top + h, left:left + w] = fg + bg
    
    # Compute area-weighted effective alpha
    alpha_t = torch.tensor(alpha_elems, dtype=torch.float32, device=device)
    area_t = torch.tensor(area_weights, dtype=torch.float32, device=device)
    alpha_eff = (alpha_t * area_t).sum() / (area_t.sum() + 1e-12)
    alpha_eff = alpha_eff.item()
    
    return x_mix, y_a, y_b, alpha_eff


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# DEVIL'S STAIRCASE PE
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class DevilStaircasePE(nn.Module):
    """Devil's Staircase PE - VECTORIZED for GPU."""
    
    def __init__(self, levels=20, features_per_level=4, smooth_tau=0.25, base=3):
        super().__init__()
        self.levels = levels
        self.features_per_level = features_per_level
        self.tau = smooth_tau
        self.base = base
        
        self.alpha = nn.Parameter(torch.tensor(0.1))
        
        # Precompute level scales and powers
        self.register_buffer('k_range', torch.arange(1, levels + 1, dtype=torch.float32))
        self.register_buffer('cantor_powers', 0.5 ** self.k_range)
        
        self.base_features = 2
        if features_per_level > 2:
            self.feature_expansion = nn.Linear(self.base_features, features_per_level)
        else:
            self.feature_expansion = None
        
    def forward(self, positions, seq_len):
        B = positions.shape[0]
        device = positions.device
        
        x = positions.float() / max(1, (seq_len - 1))
        x = x.clamp(1e-6, 1.0 - 1e-6)  # [B]
        
        # VECTORIZED: Compute all levels at once
        scales = self.base ** self.k_range.to(device)  # [levels]
        y = (x.unsqueeze(1) * scales.unsqueeze(0)) % self.base  # [B, levels]
        
        # VECTORIZED: Triadic softmax for all levels
        centers = torch.tensor([0.5, 1.5, 2.5], device=device, dtype=x.dtype)
        d2 = (y.unsqueeze(-1) - centers) ** 2  # [B, levels, 3]
        logits = -d2 / (self.tau + 1e-8)
        p = F.softmax(logits, dim=-1)  # [B, levels, 3]
        
        # VECTORIZED: Cantor bits
        bit_k = p[..., 2] + self.alpha * p[..., 1]  # [B, levels]
        
        # VECTORIZED: Cantor sum (single matmul instead of loop)
        Cx = (bit_k * self.cantor_powers.to(device).unsqueeze(0)).sum(dim=1)  # [B]
        
        # VECTORIZED: Entropy and PDF
        ent = -(p * p.clamp_min(1e-8).log()).sum(dim=-1)  # [B, levels]
        pdf_proxy = 1.1 - ent / math.log(3.0)  # [B, levels]
        
        # Stack features
        base_feat = torch.stack([bit_k, pdf_proxy], dim=-1)  # [B, levels, 2]
        
        if self.feature_expansion is not None:
            # [B, levels, 2] -> [B, levels, features_per_level]
            pe_levels = self.feature_expansion(base_feat)
        else:
            pe_levels = base_feat
        
        return pe_levels, Cx


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# GEOMETRIC BASIN COMPATIBILITY
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class GeometricBasinCompatibility(nn.Module):
    """Compute geometric compatibility scores - 4-factor product."""
    
    def __init__(self, num_classes=100, pe_levels=20, features_per_level=4):
        super().__init__()
        
        self.num_classes = num_classes
        self.pe_levels = pe_levels
        self.features_per_level = features_per_level
        
        self.class_signatures = nn.Parameter(
            torch.randn(num_classes, pe_levels, features_per_level) * 0.1
        )
        
        self.cantor_prototypes = nn.Parameter(
            torch.linspace(0.0, 1.0, num_classes)
        )
        
        self.level_resonance = nn.Parameter(
            torch.ones(num_classes, pe_levels) / pe_levels
        )
    
    def forward(self, pe_levels, cantor_measures):
        B = pe_levels.shape[0]
        
        # 1. TRIADIC COMPATIBILITY
        pe_norm = F.normalize(pe_levels, p=2, dim=-1)
        sig_norm = F.normalize(self.class_signatures, p=2, dim=-1)
        
        similarities = torch.einsum('blf,clf->bcl', pe_norm, sig_norm)
        similarities = (similarities + 1) / 2
        
        resonance = F.softmax(self.level_resonance, dim=-1)
        triadic_compat = (similarities * resonance.unsqueeze(0)).sum(dim=-1)
        
        # 2. SELF-SIMILARITY - VECTORIZED
        level_k = pe_levels[:, :-1, :]   # [B, 19, features] - all levels except last
        level_k1 = pe_levels[:, 1:, :]    # [B, 19, features] - all levels except first

        # Compute all pairwise similarities at once
        sim = F.cosine_similarity(level_k, level_k1, dim=-1, eps=1e-8)  # [B, 19]
        sim = (sim + 1) / 2
        self_sim_pattern = sim  # No stack needed, already [B, levels-1]
        
        expected_patterns = torch.sigmoid(
            self.level_resonance[:, :-1] - self.level_resonance[:, 1:]
        )
        
        pattern_diff = torch.abs(
            self_sim_pattern.unsqueeze(1) - expected_patterns.unsqueeze(0)
        )
        self_sim_compat = 1 - pattern_diff.mean(dim=-1)
        self_sim_compat = torch.clamp(self_sim_compat, 0.0, 1.0)
        
        # 3. CANTOR COHERENCE
        distances = torch.abs(
            cantor_measures.unsqueeze(1) - self.cantor_prototypes.unsqueeze(0)
        )
        cantor_compat = torch.exp(-distances ** 2 / 0.1) + 1e-8
        
        # 4. HIERARCHICAL CHECK
        split_point = self.pe_levels // 2
        early_levels = pe_levels[:, :split_point, :].mean(dim=1)
        late_levels = pe_levels[:, split_point:, :].mean(dim=1)
        
        early_targets = self.class_signatures[:, :split_point, :].mean(dim=1)
        late_targets = self.class_signatures[:, split_point:, :].mean(dim=1)
        
        early_levels_norm = F.normalize(early_levels, p=2, dim=-1)
        late_levels_norm = F.normalize(late_levels, p=2, dim=-1)
        early_targets_norm = F.normalize(early_targets, p=2, dim=-1)
        late_targets_norm = F.normalize(late_targets, p=2, dim=-1)
        
        early_compat = torch.matmul(early_levels_norm, early_targets_norm.t())
        late_compat = torch.matmul(late_levels_norm, late_targets_norm.t())
        
        early_compat = (early_compat + 1) / 2
        late_compat = (late_compat + 1) / 2
        hier_compat = (early_compat + late_compat) / 2
        
        # 5. COMBINE (geometric mean)
        eps = 1e-6
        triadic_compat = torch.clamp(triadic_compat, eps, 1.0)
        self_sim_compat = torch.clamp(self_sim_compat, eps, 1.0)
        cantor_compat = torch.clamp(cantor_compat, eps, 1.0)
        hier_compat = torch.clamp(hier_compat, eps, 1.0)
        
        compatibility_scores = (
            triadic_compat * 
            self_sim_compat * 
            cantor_compat * 
            hier_compat
        ) ** 0.25
        
        return compatibility_scores


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# GEOMETRIC BASIN LOSS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class GeometricBasinLoss(nn.Module):
    """Loss supervising geometric basin stability field."""
    
    def __init__(self, temperature=0.1):
        super().__init__()
        self.temperature = temperature
        
    def forward(self, compatibility_scores, labels, mixed_labels=None, lam=None):
        batch_size = compatibility_scores.shape[0]
        
        if mixed_labels is not None and lam is not None:
            primary_compat = compatibility_scores[torch.arange(batch_size), labels]
            secondary_compat = compatibility_scores[torch.arange(batch_size), mixed_labels]
            
            primary_loss = F.mse_loss(primary_compat, torch.full_like(primary_compat, lam))
            secondary_loss = F.mse_loss(secondary_compat, torch.full_like(secondary_compat, 1 - lam))
            
            soft_targets = torch.zeros_like(compatibility_scores)
            soft_targets[torch.arange(batch_size), labels] = lam
            soft_targets[torch.arange(batch_size), mixed_labels] = 1 - lam
            
            compat_normalized = compatibility_scores / (compatibility_scores.sum(dim=1, keepdim=True) + 1e-8)
            kl_loss = F.kl_div(
                compat_normalized.log(),
                soft_targets,
                reduction='batchmean'
            )
            
            total_loss = primary_loss + secondary_loss + 0.1 * kl_loss
            
        else:
            correct_compat = compatibility_scores[torch.arange(batch_size), labels]
            correct_loss = -torch.log(correct_compat + 1e-8).mean()
            
            mask = torch.ones_like(compatibility_scores)
            mask[torch.arange(batch_size), labels] = 0
            
            incorrect_compat = compatibility_scores * mask
            incorrect_loss = torch.log(1 - incorrect_compat + 1e-8).mean()
            incorrect_loss = -incorrect_loss
            
            scaled_scores = compatibility_scores / self.temperature
            log_probs = F.log_softmax(scaled_scores, dim=1)
            contrastive_loss = F.nll_loss(log_probs, labels)
            
            total_loss = correct_loss + 0.5 * incorrect_loss + 0.5 * contrastive_loss
        
        return total_loss


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# GEOMETRIC BASIN CLASSIFIER
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class GeometricBasinClassifier(nn.Module):
    """Geometric basin classifier with ResNet18 backbone + Cantor PE."""
    
    def __init__(self, num_classes=100, pe_levels=20, pe_features_per_level=4, dropout=0.1, pretrained=False):
        super().__init__()
        
        self.num_classes = num_classes
        self.pe_levels = pe_levels
        self.pe_features_per_level = pe_features_per_level
        
        # ResNet18 backbone from torchvision
        from torchvision.models import resnet18, ResNet18_Weights
        if pretrained:
            resnet = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        else:
            resnet = resnet18(weights=None) # will be running both types of train labeled
        
        # Extract feature extractor (everything except fc layer)
        self.backbone = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4,
            resnet.avgpool
        )
        
        # ResNet18 outputs 512 features
        self.feature_dim = 512
        self.dropout = nn.Dropout(dropout)
        
        # Devil's Staircase PE
        self.pe = DevilStaircasePE(pe_levels, pe_features_per_level)
        
        # PE modulator (adjusted for ResNet18's 512 features)
        self.pe_modulator = nn.Sequential(
            nn.Linear(self.feature_dim, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, pe_levels * pe_features_per_level)
        )
        
        # Geometric basin
        self.basin = GeometricBasinCompatibility(
            num_classes, 
            pe_levels, 
            pe_features_per_level
        )
        
    def forward(self, x, return_details=False):
        batch_size = x.shape[0]
        
        # ResNet18 backbone
        cnn_features = self.backbone(x)
        cnn_features = torch.flatten(cnn_features, 1)
        cnn_features = self.dropout(cnn_features)
        
        # Generate PE
        positions = torch.arange(batch_size, device=x.device)
        pe_levels, cantor_measures = self.pe(positions, seq_len=batch_size)
        
        # Modulate PE with CNN features
        modulation = self.pe_modulator(cnn_features)
        modulation = modulation.view(batch_size, self.pe_levels, self.pe_features_per_level)
        pe_levels = pe_levels + 0.1 * modulation
        
        # Geometric basin compatibility
        compatibility_scores = self.basin(pe_levels, cantor_measures)
        
        if return_details:
            return {
                'compatibility_scores': compatibility_scores,
                'pe_levels': pe_levels,
                'cantor_measures': cantor_measures,
                'cnn_features': cnn_features
            }
        
        return compatibility_scores