import torch
import timm
import torch.nn as nn
from torchvision import models

def get_style_vector(feat):
    # feat: (B, C)
    mu = feat.mean(dim=1)    # shape: [B]
    sigma = feat.std(dim=1)  # shape: [B]
    style_vec = torch.cat([mu.unsqueeze(1), sigma.unsqueeze(1)], dim=1) # shape: [B, 2]
    return style_vec

class ASAR(nn.Module):
    def __init__(self, feature_dim=2048):
        super(ASAR, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(feature_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, feature_dim)
        )

    def forward(self, x):
        mu = x.mean(dim=1, keepdim=True)   # [B, 1]
        sigma = x.std(dim=1, keepdim=True) # [B, 1]

        content = (x - mu) / (sigma + 1e-5)

        z = self.encoder(content)
        recon_content = self.decoder(z)
        recon_x = recon_content * sigma + mu

        return recon_x

class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        backbone = models.resnet50(pretrained=True)
        backbone.fc = nn.Identity()
        self.backbone = backbone
        self.output_dim = 2048

    def forward(self, x):
        return self.backbone(x)