#!/usr/bin/env python3
"""#9 Hyperbolic Poincaré Ball Embeddings for ABZ hierarchy.

Reference: Nickel & Kiela 2017 (arXiv:1705.08039);
Hyperbolic Image Embeddings (Khrulkov 2020 CVPR).

Poincaré ball doğal olarak hiyerarşik yapıları embed eder.
ABZ numaralandırması tree-like → hyperbolic space ideal.
"""
import torch
import torch.nn as nn

def exponential_map(v, c=1.0):
    """Exp map at origin: Euclidean → Poincaré ball."""
    v_norm = v.norm(dim=-1, keepdim=True).clamp(min=1e-8)
    return torch.tanh(v_norm * c**0.5) * v / (v_norm * c**0.5)

def poincare_distance(x, y, c=1.0):
    """Poincaré ball distance (batched)."""
    x = x.clamp(max=1-1e-5, min=-(1-1e-5))
    y = y.clamp(max=1-1e-5, min=-(1-1e-5))
    diff = x - y
    num = 2 * (diff * diff).sum(-1)
    denom = (1 - c * (x*x).sum(-1)) * (1 - c * (y*y).sum(-1)) + 1e-8
    return torch.acosh(torch.clamp(1 + num / denom, min=1+1e-7)) / c**0.5

class HyperbolicClassifier(nn.Module):
    """Maps features → Poincaré ball → distance to class prototypes."""
    
    def __init__(self, in_dim=1024, n_classes=198, hyperbolic_dim=128, c=1.0):
        super().__init__()
        self.project = nn.Linear(in_dim, hyperbolic_dim)
        # Class prototypes in Euclidean; mapped to hyperbolic during forward
        self.class_protos = nn.Parameter(torch.randn(n_classes, hyperbolic_dim) * 0.01)
        self.c = c
    
    def forward(self, features):
        euc = self.project(features)  # (B, H)
        hyp = exponential_map(euc, self.c)  # Poincaré ball
        # Prototypes also in hyperbolic
        protos_hyp = exponential_map(self.class_protos, self.c)
        # Distance → negative logits (smaller dist = higher logit)
        dists = poincare_distance(
            hyp.unsqueeze(1),  # (B, 1, H)
            protos_hyp.unsqueeze(0),  # (1, C, H)
            self.c
        )  # (B, C)
        return -dists  # negative distance = logits

if __name__ == '__main__':
    torch.manual_seed(0)
    hc = HyperbolicClassifier(in_dim=1024, n_classes=198, hyperbolic_dim=128)
    feats = torch.randn(8, 1024)
    logits = hc(feats)
    print(f"Hyperbolic logits: {logits.shape}, range: [{logits.min().item():.3f}, {logits.max().item():.3f}]")