"""Minimal inference for FairFace-MonkeyCake-3Task-224."""
import torch
import timm
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from torchvision.transforms import InterpolationMode

ARCH = "eva02_large_patch14_224.mim_in22k"
FEATURE_DIM = 1024
SIZE = 224
MEAN = (0.48145466, 0.4578275, 0.40821073)
STD = (0.26862954, 0.26130258, 0.27577711)
RACE = ('White', 'Black', 'Latino_Hispanic', 'East Asian', 'Southeast Asian', 'Indian', 'Middle Eastern')
GENDER = ("Male", "Female")
AGE = ("0-2", "3-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "more than 70")


class Student(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = timm.create_model(ARCH, pretrained=False, num_classes=0)
        self.race_head = nn.Linear(FEATURE_DIM, 7)
        self.gender_head = nn.Linear(FEATURE_DIM, 2)
        self.age_head = nn.Linear(FEATURE_DIM, 8)

    def forward(self, x):
        f = self.backbone(x)
        return {"race": self.race_head(f), "gender": self.gender_head(f), "age": self.age_head(f)}


def load(path="model.pt"):
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    model = Student()
    model.load_state_dict(ckpt["ema"] if ckpt.get("selected_ema") else ckpt["model"])
    return model.eval()


def predict(model, image_path):
    transform = transforms.Compose([
        transforms.Resize((SIZE, SIZE), interpolation=InterpolationMode.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize(mean=MEAN, std=STD),
    ])
    with torch.no_grad():
        tensor = transform(Image.open(image_path).convert("RGB")).unsqueeze(0)
        out = model(tensor)
    age_index = int((out["age"][0] > 0).sum())
    return {"race": RACE[int(out["race"][0].argmax())],
            "race_confidence": float(F.softmax(out["race"][0], dim=0).max()),
            "gender": GENDER[int(out["gender"][0].argmax())],
            "age": AGE[age_index]}