import torch
import torch.nn as nn
import torchvision.models as models
from PIL import Image
from torchvision import transforms

####################################################################################################################
# Define your model and transform and all necessary helper functions here                                          #
# They will be imported to the exp_recognition.py file                                                             #
####################################################################################################################

# Must match ImageFolder's alphabetical class_to_idx from the training notebook.
classes = {
    0: 'ANGER',
    1: 'DISGUST',
    2: 'FEAR',
    3: 'HAPPINESS',
    4: 'NEUTRAL',
    5: 'SADNESS',
    6: 'SURPRISE',
}

IMG_SIZE = 100
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


class facExpRec(nn.Module):
    """ResNet18 expression classifier trained in the expression notebook."""

    def __init__(self, num_classes=len(classes)):
        super().__init__()
        self.backbone = models.resnet18(weights=None)
        in_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(in_features, num_classes)

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


def ensure_rgb(image):
    if isinstance(image, Image.Image):
        return image.convert('RGB')
    return Image.fromarray(image).convert('RGB')


trnscm = transforms.Compose([
    transforms.Lambda(ensure_rgb),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])


def _extract_state_dict(checkpoint):
    if not isinstance(checkpoint, dict):
        return checkpoint

    for key in ('state_dict', 'model_state_dict', 'net_dict'):
        if key in checkpoint:
            return checkpoint[key]

    return checkpoint


def _normalize_state_dict_keys(state_dict):
    normalized = {}
    for key, value in state_dict.items():
        if key.startswith('module.'):
            key = key[len('module.'):]
        if key.startswith('model.'):
            key = key[len('model.'):]
        normalized[key] = value
    return normalized


def load_model(checkpoint_path, device, num_classes=len(classes)):
    model = facExpRec(num_classes=num_classes).to(device)
    checkpoint = torch.load(checkpoint_path, map_location=device)
    state_dict = _normalize_state_dict_keys(_extract_state_dict(checkpoint))

    if any(key.startswith('backbone.') for key in state_dict):
        model.load_state_dict(state_dict, strict=True)
    else:
        model.backbone.load_state_dict(state_dict, strict=True)

    model.eval()
    return model