# coding=utf-8
import torch
from models.codec_us import *
#from const import *

from p2voxel import *
import numpy

def getDH(x):
    torch.backends.cudnn.benchmark = True

    device = torch.device(f"cuda:0" if torch.cuda.is_available() else "cpu")


    model = convNN3d(1,18,8).to(device)

    ckpt_file = "./model_epoch3000.pth"

    checkpoint = torch.load(ckpt_file)

    state_dict = {}
    for k, v in checkpoint['model_state_dict'].items():
        if k in model.state_dict().keys():
            state_dict[k] = v
    model.load_state_dict(state_dict)

    # start_epoch = checkpoint['epoch']
    # loss = checkpoint['loss']
    # print(start_epoch, loss)


    #### 从p2voxel代码里生成voxel
    mtype = 387 - 1

    in_parameter_names, out_parameter_names = showPara(mtype)
    #print('in_parameters', in_parameter_names)
    #print('out_parameters', out_parameter_names)

    in_parameter_values = x[:7]

    out_parameter_values = x[7:]

    parameters = setPara(mtype, in_parameter_values, out_parameter_values)

    voxel = p2voxel(mtype, parameters, resolution=39)
    #drawVoxel_3(voxel)
    #print (voxel.shape)

    voxel = torch.from_numpy(voxel).unsqueeze(0).unsqueeze(0).to(device)

    ###

    ### 随机的输入
    # voxel = torch.rand((1,1,40,40,40)).to(device)

    ### 使用模型计算出DH流程

    U = model(voxel)

    batch_size = voxel.shape[0]
    ref18 = U.contiguous().view(batch_size, 18, -1)
    map0 = ref18[:, 0::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    map1 = ref18[:, 1::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    map2 = ref18[:, 2::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    map3 = ref18[:, 3::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    map4 = ref18[:, 4::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    map5 = ref18[:, 5::6].permute((0, 2, 1)).contiguous().view(batch_size, -1, 1)
    output_map = torch.cat([map0, map1, map2, map3, map4, map5], 2)
    output_U = output_map[:, edofMat, :]

    DHs = disp2DH(voxel, output_U, D00, intB2, h)

    #print(DHs)

    return DHs.cpu().detach().numpy().reshape(6,6)

def test():
    default = 0.1
    x=[default]*10
    DH=getDH(x)
    #print (DH)



#test()