import torch
import torch.nn as nn
from torchvision import transforms

####################################################################################################################
# Expression Recognition Model
####################################################################################################################

# Definition of classes as dictionary
classes = {
    0: 'ANGER',
    1: 'DISGUST',
    2: 'FEAR',
    3: 'HAPPINESS',
    4: 'NEUTRAL',
    5: 'SADNESS',
    6: 'SURPRISE'
}


# ============================================================
# RESIDUAL BLOCK
# ============================================================

class ResidualBlock(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1):

        super().__init__()

        self.conv1 = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=False
        )

        self.bn1 = nn.BatchNorm2d(out_channels)

        self.relu = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(
            out_channels,
            out_channels,
            kernel_size=3,
            padding=1,
            bias=False
        )

        self.bn2 = nn.BatchNorm2d(out_channels)

        self.shortcut = nn.Sequential()

        if stride != 1 or in_channels != out_channels:

            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=stride,
                    bias=False
                ),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):

        identity = self.shortcut(x)

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity

        out = self.relu(out)

        return out


# ============================================================
# MODEL
# ============================================================

class facExpRec(nn.Module):

    def __init__(self, num_classes=7):

        super().__init__()

        self.features = nn.Sequential(

            nn.Conv2d(
                1,
                32,
                kernel_size=3,
                padding=1
            ),

            nn.BatchNorm2d(32),

            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            ResidualBlock(32, 64),

            nn.MaxPool2d(2),

            ResidualBlock(64, 128),

            nn.MaxPool2d(2),

            ResidualBlock(128, 256),

            nn.MaxPool2d(2),

            ResidualBlock(256, 512),

            nn.AdaptiveAvgPool2d((1, 1))
        )

        self.classifier = nn.Sequential(

            nn.Flatten(),

            nn.Linear(512, 256),

            nn.ReLU(inplace=True),

            nn.Dropout(0.4),

            nn.Linear(256, 128),

            nn.ReLU(inplace=True),

            nn.Dropout(0.3),

            nn.Linear(128, num_classes)
        )

    def forward(self, x):

        x = self.features(x)

        x = self.classifier(x)

        return x


# ============================================================
# HELPER FUNCTIONS
# ============================================================

def rgb2gray(image):
    return image.convert('L')


# ============================================================
# TRANSFORMS
# ============================================================

trnscm = transforms.Compose([

    rgb2gray,

    transforms.Resize((72, 72)),

    transforms.ToTensor(),

    transforms.Normalize(
        mean=[0.5],
        std=[0.5]
    )
])