import torch
import torch.nn.functional as F

def info_nce_loss(features, temperature=0.07):
    """InfoNCE Loss 구현

    Args:
        features: (2N, dim) — 양성 쌍이 연속으로 배치
                  [z_1, z_1', z_2, z_2', ..., z_N, z_N']
        temperature: 온도 파라미터
    """
    batch_size = features.shape[0] // 2

    # 정규화
    features = F.normalize(features, dim=1)

    # 유사도 행렬 (2N x 2N)
    similarity_matrix = torch.matmul(features, features.T) / temperature

    # 양성 쌍 마스크: (i, i+N) 과 (i+N, i) 가 양성 쌍
    labels = torch.cat([
        torch.arange(batch_size, 2 * batch_size),
        torch.arange(batch_size)
    ]).to(features.device)

    # 자기 자신과의 유사도 제거 (대각선)
    mask = torch.eye(2 * batch_size, dtype=torch.bool).to(features.device)
    similarity_matrix.masked_fill_(mask, float('-inf'))

    # CrossEntropy로 InfoNCE 계산
    loss = F.cross_entropy(similarity_matrix, labels)
    return loss

import torch
import torch.nn as nn
import torchvision.transforms as T

class SimCLR(nn.Module):
    """SimCLR 프레임워크"""
    def __init__(self, base_encoder, projection_dim=128):
        super().__init__()
        self.encoder = base_encoder  # 예: ResNet-50 (fc 제거)
        encoder_dim = 2048  # ResNet-50 출력 차원

        # 프로젝션 헤드 (학습 시에만 사용)
        self.projector = nn.Sequential(
            nn.Linear(encoder_dim, encoder_dim),
            nn.ReLU(),
            nn.Linear(encoder_dim, projection_dim),
        )

    def forward(self, x):
        h = self.encoder(x)          # 표현 (다운스트림 사용)
        z = self.projector(h)         # 프로젝션 (학습 시에만 사용)
        return h, z


# SimCLR 데이터 증강 파이프라인
class SimCLRAugmentation:
    """하나의 이미지에서 두 개의 증강 뷰 생성"""
    def __init__(self, size=224):
        self.transform = T.Compose([
            T.RandomResizedCrop(size, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(),
            T.RandomApply([T.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.GaussianBlur(kernel_size=23, sigma=(0.1, 2.0)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]),
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)

import torch
import torch.nn.functional as F

def clip_loss(image_features, text_features, temperature=0.07):
    """CLIP 대조 손실 (간소화 버전)

    Args:
        image_features: (N, dim) 이미지 임베딩
        text_features: (N, dim) 텍스트 임베딩
        temperature: 학습 가능한 온도 파라미터
    """
    # 정규화
    image_features = F.normalize(image_features, dim=1)
    text_features = F.normalize(text_features, dim=1)

    # 유사도 행렬 (N x N)
    logits = torch.matmul(image_features, text_features.T) / temperature

    # 대각선이 양성 쌍 → labels = [0, 1, 2, ..., N-1]
    labels = torch.arange(logits.shape[0]).to(logits.device)

    # 이미지→텍스트 + 텍스트→이미지 양방향 손실
    loss_i2t = F.cross_entropy(logits, labels)
    loss_t2i = F.cross_entropy(logits.T, labels)

    return (loss_i2t + loss_t2i) / 2

# CLIP 제로샷 분류 (OpenAI clip 라이브러리)
# import clip
# model, preprocess = clip.load("ViT-B/32", device="cuda")
#
# image = preprocess(image).unsqueeze(0).to("cuda")
# text = clip.tokenize(["고양이", "강아지", "자동차"]).to("cuda")
#
# with torch.no_grad():
#     image_features = model.encode_image(image)
#     text_features = model.encode_text(text)
#     similarity = (image_features @ text_features.T).softmax(dim=-1)
#     print(f"예측: {similarity}")  # 가장 높은 유사도 클래스 선택