# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import megengine as mge
import megengine.functional as F
import megengine.hub as hub
import megengine.module as M
import official.vision.classification.resnet.model as resnet

import numpy as np


class DeconvLayers(M.Module):
    def __init__(self, nf1, nf2s, kernels, num_layers, bias=True, norm=M.BatchNorm2d):
        super(DeconvLayers, self).__init__()
        _body = []
        for i in range(num_layers):
            kernel = kernels[i]
            padding = (
                kernel // 3
            )  # padding=0 when kernel=2 and padding=1 when kernel=4 or kernel=3
            _body += [
                M.ConvTranspose2d(nf1, nf2s[i], kernel, 2, padding, bias=bias),
                norm(nf2s[i]),
                M.ReLU(),
            ]
            nf1 = nf2s[i]
        self.body = M.Sequential(*_body)

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


class SimpleBaseline(M.Module):
    def __init__(self, backbone, cfg):
        super(SimpleBaseline, self).__init__()
        norm = M.BatchNorm2d
        self.backbone = getattr(resnet, backbone)(
            norm=norm, pretrained=cfg.backbone_pretrained
        )
        del self.backbone.fc

        self.cfg = cfg

        self.deconv_layers = DeconvLayers(
            cfg.initial_deconv_channels,
            cfg.deconv_channels,
            cfg.deconv_kernel_sizes,
            cfg.num_deconv_layers,
            cfg.deconv_with_bias,
            norm,
        )
        self.last_layer = M.Conv2d(cfg.deconv_channels[-1], cfg.keypoint_num, 3, 1, 1)

        self._initialize_weights()

        self.inputs = {
            "image": mge.tensor(dtype="float32"),
            "heatmap": mge.tensor(dtype="float32"),
            "heat_valid": mge.tensor(dtype="float32"),
        }

    def calc_loss(self):
        out = self.forward(self.inputs["image"])
        valid = self.inputs["heat_valid"][:, :, None, None]
        label = self.inputs["heatmap"][:, -1]
        loss = 
F.square_loss(out * valid, label * valid)