from torch.autograd import Function
import torch.nn as nn
from transformers import ViTModel


class GradientReversalFn(Function):
    @staticmethod
    def forward(ctx, x, lambda_):
        ctx.lambda_ = lambda_
        return x.clone()

    @staticmethod
    def backward(ctx, grads):
        return -ctx.lambda_ * grads, None


class GradientReversal(nn.Module):
    def __init__(self, lambda_=1.0):
        super().__init__()
        self.lambda_ = lambda_

    def forward(self, x):
        return GradientReversalFn.apply(x, self.lambda_)

    def set_lambda(self, v):
        self.lambda_ = v


class FERModel(nn.Module):
    def __init__(self, num_classes=7, num_domains=2, domain_lambda=0.1):
        super().__init__()

        self.backbone = ViTModel.from_pretrained(
            'trpakov/vit-face-expression',
            ignore_mismatched_sizes=True,
            add_pooling_layer=False
        )
        D = 768  # ViT-base feature dim

        self.emotion_head = nn.Sequential(
            nn.LayerNorm(D),
            nn.Linear(D, 512),
            nn.GELU(),
            nn.Dropout(0.4),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )

        self.grl = GradientReversal(lambda_=domain_lambda)
        self.domain_head = nn.Sequential(
            nn.LayerNorm(D),
            nn.Linear(D, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_domains)
        )

    def forward(self, x):
        out = self.backbone(pixel_values=x)
        features = out.last_hidden_state[:, 0, :]  # CLS token [B, 768]
        return self.emotion_head(features), self.domain_head(self.grl(features))


# Config constants
BACKBONE = 'trpakov/vit-face-expression'
IMG_SIZE = 224
FEATURE_DIM = 768
NUM_CLASSES = 7
NUM_DOMAINS = 2

VIT_MEAN = [0.5, 0.5, 0.5]
VIT_STD = [0.5, 0.5, 0.5]

ID_TO_EMOTION = {
    0: 'angry',
    1: 'disgust',
    2: 'fear',
    3: 'happy',
    4: 'neutral',
    5: 'sad',
    6: 'surprise'
}

EMOTION_TO_ID = {v: k for k, v in ID_TO_EMOTION.items()}