import torch
from torch.nn import Conv2d, Sequential, ModuleList, ReLU
from ..nn.mobilenet import MobileNetV1

from .fpn_ssd import FPNSSD
from .predictor import Predictor
from .config import mobilenetv1_ssd_config as config


def create_fpn_mobilenetv1_ssd(num_classes):
    base_net = MobileNetV1(1001).features  # disable dropout layer

    source_layer_indexes = [
        (69, Conv2d(in_channels=512, out_channels=256, kernel_size=1)),
        (len(base_net), Conv2d(in_channels=1024, out_channels=256, kernel_size=1)),
    ]
    extras = ModuleList(
        [
            Sequential(
                Conv2d(in_channels=1024, out_channels=256, kernel_size=1),
                ReLU(),
                Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1),
                ReLU(),
            ),
            Sequential(
                Conv2d(in_channels=256, out_channels=128, kernel_size=1),
                ReLU(),
                Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
                ReLU(),
            ),
            Sequential(
                Conv2d(in_channels=256, out_channels=128, kernel_size=1),
                ReLU(),
                Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
                ReLU(),
            ),
            Sequential(
                Conv2d(in_channels=256, out_channels=128, kernel_size=1),
                ReLU(),
                Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
                ReLU(),
            ),
        ]
    )

    regression_headers = ModuleList(
        [
            Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
            Conv2d(
                in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1
            ),  #  TODO: change to kernel_size=1, padding=0?
        ]
    )

    classification_headers = ModuleList(
        [
            Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
            Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
            Conv2d(
                in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1
            ),  # TODO: change to kernel_size=1, padding=0?
        ]
    )

    return FPNSSD(num_classes, base_net, source_layer_indexes, extras, classification_headers, regression_headers)


def create_fpn_mobilenetv1_ssd_predictor(
    net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device("cpu")
):
    predictor = Predictor(
        net,
        config.image_size,
        config.image_mean,
        config.priors,
        config.center_variance,
        config.size_variance,
        nms_method=nms_method,
        iou_threshold=config.iou_threshold,
        candidate_size=candidate_size,
        sigma=sigma,
        device=device,
    )
    return predictor