# -*- coding: utf-8 -*-
import numpy as np
import torch
from collections import OrderedDict

# import scipy.io as io
import hdf5storage

"""
# --------------------------------------------
# Convert matconvnet SimpleNN model into pytorch model
# --------------------------------------------
# Kai Zhang (cskaizhang@gmail.com)
# https://github.com/cszn
# 28/Nov/2019
# --------------------------------------------
"""


def weights2tensor(x, squeeze=False, in_features=None, out_features=None):
    """Modified version of https://github.com/albanie/pytorch-mcn
    Adjust memory layout and load weights as torch tensor
    Args:
        x (ndaray): a numpy array, corresponding to a set of network weights
           stored in column major order
        squeeze (bool) [False]: whether to squeeze the tensor (i.e. remove
           singletons from the trailing dimensions. So after converting to
           pytorch layout (C_out, C_in, H, W), if the shape is (A, B, 1, 1)
           it will be reshaped to a matrix with shape (A,B).
        in_features (int :: None): used to reshape weights for a linear block.
        out_features (int :: None): used to reshape weights for a linear block.
    Returns:
        torch.tensor: a permuted sets of weights, matching the pytorch layout
        convention
    """
    if x.ndim == 4:
        x = x.transpose((3, 2, 0, 1))
# for FFDNet, pixel-shuffle layer
#        if x.shape[1]==13:
#            x=x[:,[0,2,1,3,  4,6,5,7, 8,10,9,11, 12],:,:]
#        if x.shape[0]==12:
#            x=x[[0,2,1,3,  4,6,5,7, 8,10,9,11],:,:,:]
#        if x.shape[1]==5:
#            x=x[:,[0,2,1,3,  4],:,:]
#        if x.shape[0]==4:
#            x=x[[0,2,1,3],:,:,:]
# for SRMD, pixel-shuffle layer
#        if x.shape[0]==12:
#            x=x[[0,2,1,3,  4,6,5,7, 8,10,9,11],:,:,:]
#        if x.shape[0]==27:
#            x=x[[0,3,6,1,4,7,2,5,8, 0+9,3+9,6+9,1+9,4+9,7+9,2+9,5+9,8+9, 0+18,3+18,6+18,1+18,4+18,7+18,2+18,5+18,8+18],:,:,:]
#        if x.shape[0]==48:
#            x=x[[0,4,8,12,1,5,9,13,2,6,10,14,3,7,11,15,  0+16,4+16,8+16,12+16,1+16,5+16,9+16,13+16,2+16,6+16,10+16,14+16,3+16,7+16,11+16,15+16,  0+32,4+32,8+32,12+32,1+32,5+32,9+32,13+32,2+32,6+32,10+32,14+32,3+32,7+32,11+32,15+32],:,:,:]

    elif x.ndim == 3:  # add by Kai
        x = x[:, :, :, None]
        x = x.transpose((3, 2, 0, 1))
    elif x.ndim == 2:
        if x.shape[1] == 1:
            x = x.flatten()
    if squeeze:
        if in_features and out_features:
            x = x.reshape((out_features, in_features))
        x = np.squeeze(x)
    return torch.from_numpy(np.ascontiguousarray(x))


def save_model(network, save_path):
    state_dict = network.state_dict()
    for key, param in state_dict.items():
        state_dict[key] = param.cpu()
    torch.save(state_dict, save_path)


if __name__ == '__main__':

    #    from utils import utils_logger
    #    import logging
    #    utils_logger.logger_info('a', 'a.log')
    #    logger = logging.getLogger('a')
    #
    # mcn = hdf5storage.loadmat('/model_zoo/matfile/FFDNet_Clip_gray.mat')
    mcn = hdf5storage.loadmat('models/modelcolor.mat')

    # logger.info(mcn['CNNdenoiser'][0][0][0][1][0][0][0][0])

    mat_net = OrderedDict()
    for idx in range(25):
        mat_net[str(idx)] = OrderedDict()
        count = -1

        print(idx)
        for i in range(13):

            if mcn['CNNdenoiser'][0][idx][0][i][0][0][0][0] == 'conv':

                count += 1
                w = mcn['CNNdenoiser'][0][idx][0][i][0][1][0][0]
                # print(w.shape)
                w = weights2tensor(w)
                # print(w.shape)

                b = mcn['CNNdenoiser'][0][idx][0][i][0][1][0][1]
                b = weights2tensor(b)
                print(b.shape)

                mat_net[str(idx)]['model.{:d}.weight'.format(count * 2)] = w
                mat_net[str(idx)]['model.{:d}.bias'.format(count * 2)] = b

    torch.save(mat_net, 'model_zoo/modelcolor.pth')


#    from models.network_dncnn import IRCNN as net
#    network = net(in_nc=3, out_nc=3, nc=64)
#    state_dict = network.state_dict()
#
#    #show_kv(state_dict)
#
#    for i in range(len(mcn['net'][0][0][0])):
#        print(mcn['net'][0][0][0][i][0][0][0][0])
#
#    count = -1
#    mat_net = OrderedDict()
#    for i in range(len(mcn['net'][0][0][0])):
#        if mcn['net'][0][0][0][i][0][0][0][0] == 'conv':
#
#            count += 1
#            w = mcn['net'][0][0][0][i][0][1][0][0]
#            print(w.shape)
#            w = weights2tensor(w)
#            print(w.shape)
#
#            b = mcn['net'][0][0][0][i][0][1][0][1]
#            b = weights2tensor(b)
#            print(b.shape)
#
#            mat_net['model.{:d}.weight'.format(count*2)] = w
#            mat_net['model.{:d}.bias'.format(count*2)] = b
#
#    torch.save(mat_net, 'E:/pytorch/KAIR_ongoing/model_zoo/ffdnet_gray_clip.pth')
#
#
#
#    crt_net = torch.load('E:/pytorch/KAIR_ongoing/model_zoo/imdn_x4.pth')
#    def show_kv(net):
#        for k, v in net.items():
#            print(k)
#
#    show_kv(crt_net)


#    from models.network_dncnn import DnCNN as net
#    network = net(in_nc=2, out_nc=1, nc=64, nb=20, act_mode='R')

#    from models.network_srmd import SRMD as net
#    #network = net(in_nc=1, out_nc=1, nc=64, nb=15, act_mode='R')
#    network = net(in_nc=19, out_nc=3, nc=128, nb=12, upscale=4, act_mode='R', upsample_mode='pixelshuffle')
#
#    from models.network_rrdb import RRDB as net
#    network = net(in_nc=3, out_nc=3, nc=64, nb=23, gc=32, upscale=4, act_mode='L', upsample_mode='upconv')
#
#    state_dict = network.state_dict()
#    for key, param in state_dict.items():
#        print(key)
#    from models.network_imdn import IMDN as net
#    network = net(in_nc=3, out_nc=3, nc=64, nb=8, upscale=4, act_mode='L', upsample_mode='pixelshuffle')
#    state_dict = network.state_dict()
#    mat_net = OrderedDict()
#    for ((key, param),(key2, param2)) in zip(state_dict.items(), crt_net.items()):
#        mat_net[key] = param2
#    torch.save(mat_net, 'model_zoo/imdn_x4_1.pth')
#

#    net_old = torch.load('net_old.pth')
#    def show_kv(net):
#        for k, v in net.items():
#            print(k)
#
#    show_kv(net_old)
#    from models.network_dpsr import MSRResNet_prior as net
#    model = net(in_nc=4, out_nc=3, nc=96, nb=16, upscale=4, act_mode='R', upsample_mode='pixelshuffle')
#    state_dict = network.state_dict()
#    net_new = OrderedDict()
#    for ((key, param),(key_old, param_old)) in zip(state_dict.items(), net_old.items()):
#        net_new[key] = param_old
#    torch.save(net_new, 'net_new.pth')

    # print(key)
    #  print(param.size())

    # run utils/utils_matconvnet.py
