import logging

import torch

from .networks.dlav0 import get_pose_net as get_dlav0
from .networks.large_hourglass import get_large_hourglass_net
from .networks.mobilenet import get_mobilenet_v3_small
from .networks.msra_resnet import get_pose_net
from .networks.pose_dla_dcn import get_pose_net as get_dla_dcn
from .networks.resnet_dcn import get_pose_net as get_pose_net_dcn

logger = logging.getLogger()


_model_factory = {
    "res": get_pose_net,  # default Resnet with deconv
    "dlav0": get_dlav0,  # default DLAup
    "dla": get_dla_dcn,
    "resdcn": get_pose_net_dcn,
    "hourglass": get_large_hourglass_net,
    "mobilenetv3small": get_mobilenet_v3_small,
}


def create_model(arch, heads, head_conv):
    num_layers = int(arch[arch.find("_") + 1 :]) if "_" in arch else 0
    arch = arch[: arch.find("_")] if "_" in arch else arch
    get_model = _model_factory[arch]
    model = get_model(num_layers=num_layers, heads=heads, head_conv=head_conv)
    return model


def load_model(model, model_path, opt, optimizer=None):
    start_epoch = 0
    checkpoint = torch.load(
        model_path, map_location=lambda storage, loc: storage
    )
    print("loaded {}, epoch {}".format(model_path, checkpoint["epoch"]))
    state_dict_ = checkpoint["state_dict"]
    state_dict = {}

    # convert data_parallal to model
    for k in state_dict_:
        if k.startswith("module") and not k.startswith("module_list"):
            state_dict[k[7:]] = state_dict_[k]
        else:
            state_dict[k] = state_dict_[k]
    model_state_dict = model.state_dict()

    # check loaded parameters and created model parameters
    for k in state_dict:
        if k in model_state_dict:
            if (state_dict[k].shape != model_state_dict[k].shape) or (
                opt.reset_hm
                and k.startswith("hm")
                and (state_dict[k].shape[0] in [80, 1])
            ):
                if opt.reuse_hm:
                    print(
                        "Reusing parameter {}, required shape{}, "
                        "loaded shape{}.".format(
                            k, model_state_dict[k].shape, state_dict[k].shape
                        )
                    )
                    if state_dict[k].shape[0] < state_dict[k].shape[0]:
                        model_state_dict[k][
                            : state_dict[k].shape[0]
                        ] = state_dict[k]
                    else:
                        model_state_dict[k] = state_dict[k][
                            : model_state_dict[k].shape[0]
                        ]
                    state_dict[k] = model_state_dict[k]
                else:
                    print(
                        "Skip loading parameter {}, required shape{}, "
                        "loaded shape{}.".format(
                            k, model_state_dict[k].shape, state_dict[k].shape
                        )
                    )
                    state_dict[k] = model_state_dict[k]
        else:
            print(f"Drop parameter {k}.")
    for k in model_state_dict:
        if not (k in state_dict):
            print(f"No param {k}.")
            state_dict[k] = model_state_dict[k]
    model.load_state_dict(state_dict, strict=False)

    # resume optimizer parameters
    if optimizer is not None and opt.resume:
        if "optimizer" in checkpoint:
            # optimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint["epoch"]
            start_lr = opt.lr
            for step in opt.lr_step:
                if start_epoch >= step:
                    start_lr *= 0.1
            for param_group in optimizer.param_groups:
                param_group["lr"] = start_lr
            print("Resumed optimizer with start lr", start_lr)
        else:
            print("No optimizer parameters in checkpoint.")
    if optimizer is not None:
        return model, optimizer, start_epoch
    else:
        return model


def save_model(path, epoch, model, optimizer=None):
    if isinstance(model, torch.nn.DataParallel):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()
    data = {"epoch": epoch, "state_dict": state_dict}
    if not (optimizer is None):
        data["optimizer"] = optimizer.state_dict()
    torch.save(data, path)
