import torch
import torch.nn as nn

from src.logger import get_logger

logger = get_logger(__name__)


class ResidualBlock(nn.Module):
    """
    Basic ResNet18 Block

           ┌───────────────┐
           │   Shortcut    │
           └──────┬────────┘
                  │
    Conv3x3 -> BN -> ReLU
                  │
    Conv3x3 -> BN
                  │
             Add Shortcut
                  │
                ReLU
    """

    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,
            stride=1,
            padding=1,
            bias=False
        )

        self.bn2 = nn.BatchNorm2d(
            out_channels
        )

        self.shortcut = nn.Sequential()

        # Downsample shortcut
        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


class ResNet18(nn.Module):

    def __init__(
        self,
        num_classes=11
    ):
        super().__init__()

        logger.info(
            "Building ResNet18"
        )

        # Initial Layer

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

        self.bn1 = nn.BatchNorm2d(
            64
        )

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

        # Stage 1

        self.layer1 = nn.Sequential(
            ResidualBlock(
                64,
                64
            ),
            ResidualBlock(
                64,
                64
            )
        )

        # Stage 2

        self.layer2 = nn.Sequential(
            ResidualBlock(
                64,
                128,
                stride=2
            ),
            ResidualBlock(
                128,
                128
            )
        )

        # Stage 3

        self.layer3 = nn.Sequential(
            ResidualBlock(
                128,
                256,
                stride=2
            ),
            ResidualBlock(
                256,
                256
            )
        )

        # Stage 4

        self.layer4 = nn.Sequential(
            ResidualBlock(
                256,
                512,
                stride=2
            ),
            ResidualBlock(
                512,
                512
            )
        )

        self.global_pool = (
            nn.AdaptiveAvgPool2d(
                (1, 1)
            )
        )

        self.dropout = nn.Dropout(
            0.5
        )

        self.fc = nn.Linear(
            512,
            num_classes
        )

        logger.info(
            "ResNet18 created successfully"
        )

    def forward(
        self,
        x
    ):

        x = self.conv1(
            x
        )

        x = self.bn1(
            x
        )

        x = self.relu(
            x
        )

        x = self.layer1(
            x
        )

        x = self.layer2(
            x
        )

        x = self.layer3(
            x
        )

        x = self.layer4(
            x
        )

        x = self.global_pool(
            x
        )

        x = torch.flatten(
            x,
            1
        )

        x = self.dropout(
            x
        )

        x = self.fc(
            x
        )

        return x


def build_model():

    model = ResNet18(
        num_classes=11
    )

    return model


if __name__ == "__main__":

    model = build_model()

    print(model)

    sample = torch.randn(
        8,
        3,
        32,
        32
    )

    output = model(
        sample
    )

    print(
        "\nOutput Shape:",
        output.shape
    )