# import torch
# import torchvision

# from models import basicblock as B


def show_kv(net):
    for k, v in net.items():
        print(k)

# should run train debug mode first to get an initial model
# crt_net = torch.load('../../experiments/debug_SRResNet_bicx4_in3nf64nb16/models/8_G.pth')
#
# for k, v in crt_net.items():
#    print(k)
# for k, v in crt_net.items():
#    if k in pretrained_net:
#        crt_net[k] = pretrained_net[k]
#        print('replace ... ', k)

# x2 -> x4
# crt_net['model.5.weight'] = pretrained_net['model.2.weight']
# crt_net['model.5.bias'] = pretrained_net['model.2.bias']
# crt_net['model.8.weight'] = pretrained_net['model.5.weight']
# crt_net['model.8.bias'] = pretrained_net['model.5.bias']
# crt_net['model.10.weight'] = pretrained_net['model.7.weight']
# crt_net['model.10.bias'] = pretrained_net['model.7.bias']
# torch.save(crt_net, '../pretrained_tmp.pth')


# x2 -> x3
'''
in_filter = pretrained_net['model.2.weight'] # 256, 64, 3, 3
new_filter = torch.Tensor(576, 64, 3, 3)
new_filter[0:256, :, :, :] = in_filter
new_filter[256:512, :, :, :] = in_filter
new_filter[512:, :, :, :] = in_filter[0:576-512, :, :, :]
crt_net['model.2.weight'] = new_filter

in_bias = pretrained_net['model.2.bias']  # 256, 64, 3, 3
new_bias = torch.Tensor(576)
new_bias[0:256] = in_bias
new_bias[256:512] = in_bias
new_bias[512:] = in_bias[0:576 - 512]
crt_net['model.2.bias'] = new_bias

torch.save(crt_net, '../pretrained_tmp.pth')
'''

# x2 -> x8
'''
crt_net['model.5.weight'] = pretrained_net['model.2.weight']
crt_net['model.5.bias'] = pretrained_net['model.2.bias']
crt_net['model.8.weight'] = pretrained_net['model.2.weight']
crt_net['model.8.bias'] = pretrained_net['model.2.bias']
crt_net['model.11.weight'] = pretrained_net['model.5.weight']
crt_net['model.11.bias'] = pretrained_net['model.5.bias']
crt_net['model.13.weight'] = pretrained_net['model.7.weight']
crt_net['model.13.bias'] = pretrained_net['model.7.bias']
torch.save(crt_net, '../pretrained_tmp.pth')
'''

# x3/4/8 RGB -> Y


def rgb2gray_net(net, only_input=True):

    if only_input:
        in_filter = net['0.weight']
        in_new_filter = in_filter[:,
                                  0,
                                  :,
                                  :] * 0.2989 + in_filter[:,
                                                          1,
                                                          :,
                                                          :] * 0.587 + in_filter[:,
                                                                                 2,
                                                                                 :,
                                                                                 :] * 0.114
        in_new_filter.unsqueeze_(1)
        net['0.weight'] = in_new_filter

#    out_filter = pretrained_net['model.13.weight']
#    out_new_filter = out_filter[0, :, :, :] * 0.2989 + out_filter[1, :, :, :] * 0.587 + \
#        out_filter[2, :, :, :] * 0.114
#    out_new_filter.unsqueeze_(0)
#    crt_net['model.13.weight'] = out_new_filter
#    out_bias = pretrained_net['model.13.bias']
#    out_new_bias = out_bias[0] * 0.2989 + out_bias[1] * 0.587 + out_bias[2] * 0.114
#    out_new_bias = torch.Tensor(1).fill_(out_new_bias)
#    crt_net['model.13.bias'] = out_new_bias

#    torch.save(crt_net, '../pretrained_tmp.pth')

    return net


# if __name__ == '__main__':

#     net = torchvision.models.vgg19(pretrained=True)
#     for k, v in net.features.named_parameters():
#         if k == '0.weight':
#             in_new_filter = v[:, 0, :, :] * 0.2989 + \
#                 v[:, 1, :, :] * 0.587 + v[:, 2, :, :] * 0.114
#             in_new_filter.unsqueeze_(1)
#             v = in_new_filter
#             print(v.shape)
#             print(v[0, 0, 0, 0])
#         if k == '0.bias':
#             in_new_bias = v
#             print(v[0])

#     print(net.features[0])

#     net.features[0] = B.conv(1, 64, mode='C')

#     print(net.features[0])
#     net.features[0].weight.data = in_new_filter
#     net.features[0].bias.data = in_new_bias

#     for k, v in net.features.named_parameters():
#         if k == '0.weight':
#             print(v[0, 0, 0, 0])
#         if k == '0.bias':
#             print(v[0])

#     # transfer parameters of old model to new one
#     model_old = torch.load(model_path)
#     state_dict = model.state_dict()
#     for ((key, param), (key2, param2)) in zip(
#             model_old.items(), state_dict.items()):
#         state_dict[key2] = param
#         print([key, key2])
#         # print([param.size(), param2.size()])
#     torch.save(state_dict, 'model_new.pth')

#     # rgb2gray_net(net)