import torch.nn as nn
import numpy as np
import itertools
import os
import sys

import torch.optim as optim
from torch.utils.data import DataLoader
from torch.autograd import Variable
from tqdm import tqdm
import time
from torch.utils.tensorboard import SummaryWriter
from pytorch3d.structures import Meshes
from pytorch3d import loss

from network_motion import *
from dataio_motion import *
from utils import *


lr = 1e-4
n_worker = 4
bs = 8
n_epoch = 400
base_err = 1000

w_smooth = 150
w_reg = 20
w_h = 0.5
width = 128
height = 128
depth = 64
sa_idx = [12, 17, 22, 27, 32, 37, 42, 47, 52]
temper = 3



model_save_path = './models/model_motion'
if not os.path.exists(model_save_path):
    os.makedirs(model_save_path)

# pytorch only saves the last model
Motion_save_path = os.path.join(model_save_path, 'motionEst.pth')
Motion_LA_save_path = os.path.join(model_save_path, 'multiview.pth')

flow_criterion = nn.MSELoss()
MotionNet = MotionMesh_25d().cuda()
MV_LA = Mesh_2d().cuda()

optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                              itertools.chain(MotionNet.parameters(), MV_LA.parameters())), lr=lr)
Tensor = torch.cuda.FloatTensor
TensorLong = torch.cuda.LongTensor

# visualisation
writer = SummaryWriter('./runs/model_motion')



def train(epoch):
    MotionNet.train()
    MV_LA.train()

    epoch_loss = []
    epoch_seg_loss = []
    epoch_smooth_loss = []
    epoch_reg_loss = []
    epoch_huber_loss = []

    Myo_dice_sa = []
    Myo_dice_2ch = []
    Myo_dice_4ch = []


    for batch_idx, batch in tqdm(enumerate(training_data_loader, 1),
                                 total=len(training_data_loader)):

        img_sa_t, img_sa_ed, img_2ch_t, img_2ch_ed, img_4ch_t, img_4ch_ed, \
        contour_sa, contour_2ch, contour_4ch, \
        vertex_ed, faces, affine_inv, affine, origin = batch


        x_sa_t = Variable(img_sa_t.type(Tensor))
        x_sa_ed = Variable(img_sa_ed.type(Tensor))
        x_2ch_t = Variable(img_2ch_t.type(Tensor))
        x_2ch_ed = Variable(img_2ch_ed.type(Tensor))
        x_4ch_t = Variable(img_4ch_t.type(Tensor))
        x_4ch_ed = Variable(img_4ch_ed.type(Tensor))

        x_sa_t_5D = Variable(img_sa_t.unsqueeze(1).type(Tensor))
        x_sa_ed_5D = Variable(img_sa_ed.unsqueeze(1).type(Tensor))


        con_sa = Variable(contour_sa.type(TensorLong))  # [bs, slices, H, W]
        con_2ch = Variable(contour_2ch.type(TensorLong))  # [bs, H, W]
        con_4ch = Variable(contour_4ch.type(TensorLong))  # [bs, H, W]

        aff_sa_inv = Variable(affine_inv[:, 0,:,:].type(Tensor))
        aff_sa = Variable(affine[:, 0,:,:].type(Tensor))
        aff_2ch_inv = Variable(affine_inv[:, 1,:,:].type(Tensor))
        aff_4ch_inv = Variable(affine_inv[:, 2,:,:].type(Tensor))

        origin_sa = Variable(origin[:, 0:1, :].type(Tensor))
        origin_2ch = Variable(origin[:, 1:2, :].type(Tensor))
        origin_4ch = Variable(origin[:, 2:3, :].type(Tensor))

        vertex_0 = Variable(vertex_ed.permute(0,2,1).type(Tensor)) # [bs, 3, number of vertices]
        faces_0 = Variable(faces.type(Tensor)) # [bs, number of faces, 3]


        optimizer.zero_grad()

        net_la = MV_LA(x_2ch_t, x_2ch_ed, x_4ch_t, x_4ch_ed)
        net_sa = MotionNet(x_sa_t, x_sa_ed, net_la['conv2_2ch'], net_la['conv2s_2ch'], net_la['conv2_4ch'], net_la['conv2s_4ch'])

        # ---------------sample from 3D motion fields
        #translate coordinate
        v_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], vertex_0) + aff_sa_inv[:, :3, 3:4]
        v_ed = v_ed_o.permute(0, 2, 1) - origin_sa  # [bs, number of vertices,3]

        # normalize translated coordinate (image space) to [-1,1]
        v_ed_x = (v_ed[:, :, 0:1] - (width / 2)) / (width / 2)
        v_ed_y = (v_ed[:, :, 1:2] - (height / 2)) / (height / 2)
        v_ed_z = (v_ed[:, :, 2:3] - (depth / 2)) / (depth / 2)
        v_ed_norm = torch.cat((v_ed_x, v_ed_y, v_ed_z), 2)
        v_ed_norm_expand = v_ed_norm.unsqueeze(1).unsqueeze(1)  # [bs, 1, 1,number of vertices,3]

        # sample from 3D motion field
        pxx = F.grid_sample(net_sa['out'][:, 0:1], v_ed_norm_expand, align_corners=True).transpose(4, 3)
        pyy = F.grid_sample(net_sa['out'][:, 1:2], v_ed_norm_expand, align_corners=True).transpose(4, 3)
        pzz = F.grid_sample(net_sa['out'][:, 2:3], v_ed_norm_expand, align_corners=True).transpose(4, 3)
        delta_p = torch.cat((pxx, pyy, pzz), 4)
        # updata coor (image space)
        # print (v_ed.shape, delta_p.shape)
        v_t_norm_expand = v_ed_norm_expand + delta_p  # [bs, 1, 1,number of vertices,3]
        # t frame
        v_t_norm = v_t_norm_expand.squeeze(1).squeeze(1)
        v_t_x = v_t_norm[:, :, 0:1] * (width / 2) + (width / 2)
        v_t_y = v_t_norm[:, :, 1:2] * (height / 2) + (height / 2)
        v_t_z = v_t_norm[:, :, 2:3] * (depth / 2) + (depth / 2)
        v_t_crop = torch.cat((v_t_x, v_t_y, v_t_z), 2)
        # translate back to mesh space
        v_t = v_t_crop + origin_sa  # [bs, number of vertices,3]
        pred_vertex_t = torch.matmul(aff_sa[:, :3, :3], v_t.permute(0,2,1)) + aff_sa[:, :3, 3:4] # [bs, 3, number of vertices]
        # print (pred_vertex_t.shape)


        pred_sa_ed = transform(x_sa_t_5D, net_sa['out'], mode='bilinear')

        # -------------- differentialable slicer

        # coordinate transformation np.dot(aff_sa_SR_inv[:3,:3], points_ED.T) + aff_sa_SR_inv[:3,3:4]
        v_sa_hat_t_o = torch.matmul(aff_sa_inv[:, :3, :3], pred_vertex_t) + aff_sa_inv[:, :3, 3:4]
        v_sa_hat_t = v_sa_hat_t_o.permute(0, 2, 1) - origin_sa
        # print (v_sa_hat_t.shape)
        v_2ch_hat_t_o = torch.matmul(aff_2ch_inv[:, :3, :3], pred_vertex_t) + aff_2ch_inv[:, :3, 3:4]
        v_2ch_hat_t = v_2ch_hat_t_o.permute(0, 2, 1) - origin_2ch
        v_4ch_hat_t_o = torch.matmul(aff_4ch_inv[:, :3, :3], pred_vertex_t) + aff_4ch_inv[:, :3, 3:4]
        v_4ch_hat_t = v_4ch_hat_t_o.permute(0,2, 1) - origin_4ch

        # project vertices satisfying threshood
        # project to SAX slices, project all vertices to a target plane,
        # vertices selection is moved to loss computation function
        v_sa_hat_t_x = torch.clamp(v_sa_hat_t[:, :, 0:1], min=0, max=height - 1)
        v_sa_hat_t_y = torch.clamp(v_sa_hat_t[:, :, 1:2], min=0, max=width - 1)
        v_sa_hat_t_cp = torch.cat((v_sa_hat_t_x, v_sa_hat_t_y, v_sa_hat_t[:, :, 2:3]), 2)

        v_sa_idx_t_0, w_sa_t_0 = projection(v_sa_hat_t_cp, 12, temper)
        # print (v_sa_idx_ed_0.shape, w_sa_ed_0.shape)
        v_sa_idx_t_1, w_sa_t_1 = projection(v_sa_hat_t_cp, 17, temper)
        v_sa_idx_t_2, w_sa_t_2 = projection(v_sa_hat_t_cp, 22, temper)
        v_sa_idx_t_3, w_sa_t_3 = projection(v_sa_hat_t_cp, 27, temper)
        v_sa_idx_t_4, w_sa_t_4 = projection(v_sa_hat_t_cp, 32, temper)
        v_sa_idx_t_5, w_sa_t_5 = projection(v_sa_hat_t_cp, 37, temper)
        v_sa_idx_t_6, w_sa_t_6 = projection(v_sa_hat_t_cp, 42, temper)
        v_sa_idx_t_7, w_sa_t_7 = projection(v_sa_hat_t_cp, 47, temper)
        v_sa_idx_t_8, w_sa_t_8 = projection(v_sa_hat_t_cp, 52, temper)

        # project to LAX 2CH view
        v_2ch_hat_t_x = torch.clamp(v_2ch_hat_t[:, :, 0:1], min=0, max=height - 1)
        v_2ch_hat_t_y = torch.clamp(v_2ch_hat_t[:, :, 1:2], min=0, max=width - 1)
        v_2ch_hat_t_cp = torch.cat((v_2ch_hat_t_x, v_2ch_hat_t_y, v_2ch_hat_t[:, :, 2:3]), 2)

        v_2ch_idx_t, w_2ch_t = projection(v_2ch_hat_t_cp, 0, temper)


        # project to LAX 4CH view
        v_4ch_hat_t_x = torch.clamp(v_4ch_hat_t[:, :, 0:1], min=0, max=height - 1)
        v_4ch_hat_t_y = torch.clamp(v_4ch_hat_t[:, :, 1:2], min=0, max=width - 1)
        v_4ch_hat_t_cp = torch.cat((v_4ch_hat_t_x, v_4ch_hat_t_y, v_4ch_hat_t[:, :, 2:3]), 2)

        v_4ch_idx_t, w_4ch_t = projection(v_4ch_hat_t_cp, 0, temper)



        # --------------------- Segmentation loss------------------
        loss_seg_sa_t_0 = weightedHausdorff_batch(v_sa_idx_t_0, w_sa_t_0, con_sa[:, 0, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_1 = weightedHausdorff_batch(v_sa_idx_t_1, w_sa_t_1, con_sa[:, 1, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_2 = weightedHausdorff_batch(v_sa_idx_t_2, w_sa_t_2, con_sa[:, 2, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_3 = weightedHausdorff_batch(v_sa_idx_t_3, w_sa_t_3, con_sa[:, 3, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_4 = weightedHausdorff_batch(v_sa_idx_t_4, w_sa_t_4, con_sa[:, 4, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_5 = weightedHausdorff_batch(v_sa_idx_t_5, w_sa_t_5, con_sa[:, 5, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_6 = weightedHausdorff_batch(v_sa_idx_t_6, w_sa_t_6, con_sa[:, 6, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_7 = weightedHausdorff_batch(v_sa_idx_t_7, w_sa_t_7, con_sa[:, 7, :, :], height, width, temper,
                                                  'train')
        loss_seg_sa_t_8 = weightedHausdorff_batch(v_sa_idx_t_8, w_sa_t_8, con_sa[:, 8, :, :], height, width, temper,
                                                  'train')
        loss_seg_2ch_t = weightedHausdorff_batch(v_2ch_idx_t, w_2ch_t, con_2ch, height, width, temper, 'train')
        loss_seg_4ch_t = weightedHausdorff_batch(v_4ch_idx_t, w_4ch_t, con_4ch, height, width, temper, 'train')

        loss_seg = (loss_seg_sa_t_0 + loss_seg_sa_t_1 + loss_seg_sa_t_2 + loss_seg_sa_t_3 +
                      loss_seg_sa_t_4 + loss_seg_sa_t_5 + loss_seg_sa_t_6 + loss_seg_sa_t_7 + loss_seg_sa_t_8) / 9.0 + \
                     loss_seg_2ch_t + loss_seg_4ch_t




        #----------------smoothness loss------------
        # print (pred_vertex_t.permute(0,2,1).shape)
        trg_mesh = Meshes(verts=list(pred_vertex_t.permute(0, 2, 1)), faces=list(faces_0))
        loss_smooth = loss.mesh_laplacian_smoothing(trg_mesh, method='uniform')

        # ----------------regularization loss------------

        # define image registration as a regularization term
        loss_reg = flow_criterion(pred_sa_ed, x_sa_ed_5D)

        loss_huber = huber_loss_3d(net_sa['out'])


        loss_all = loss_seg + w_reg*loss_reg + w_smooth * loss_smooth + w_h * loss_huber

        loss_all.backward()
        optimizer.step()


        epoch_loss.append(loss_all.item())
        epoch_seg_loss.append(loss_seg.item())
        epoch_smooth_loss.append(loss_smooth.item())
        epoch_reg_loss.append(loss_reg.item())
        epoch_huber_loss.append(loss_huber.item())



        # tensorboard visulisation
        writer.add_scalar("Loss/train", loss_all, epoch * len(training_data_loader) + batch_idx)
        writer.add_scalar("Loss/train_seg", loss_seg, epoch * len(training_data_loader) + batch_idx)
        writer.add_scalar("Loss/train_reg", loss_reg, epoch * len(training_data_loader) + batch_idx)
        writer.add_scalar("Loss/train_smooth", loss_smooth, epoch * len(training_data_loader) + batch_idx)
        writer.add_scalar("Loss/train_huber", loss_huber, epoch * len(training_data_loader) + batch_idx)


        if batch_idx % 40 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss all: {:.6f}, '
                  'Seg Loss: {:.6f}, Reg Loss: {:.6f}, Smooth Loss: {:.6f}, Huber Loss: {:.6f}'.format(
                epoch, batch_idx * len(img_sa_t), len(training_data_loader.dataset),
                100. * batch_idx / len(training_data_loader), np.mean(epoch_loss),
                np.mean(epoch_seg_loss), np.mean(epoch_reg_loss), np.mean(epoch_smooth_loss), np.mean(epoch_huber_loss), np.mean(Myo_dice_sa), np.mean(Myo_dice_2ch), np.mean(Myo_dice_4ch)))

            # torch.save(model.state_dict(), model_save_path)
            # print("Checkpoint saved to {}".format(model_save_path))

def val(epoch):
    MotionNet.eval()
    MV_LA.eval()

    val_loss = []
    val_seg_loss = []
    val_smooth_loss = []
    val_reg_loss = []
    val_huber_loss = []

    global base_err
    for batch_idx, batch in tqdm(enumerate(val_data_loader, 1),
                                 total=len(val_data_loader)):

        img_sa_t, img_sa_ed, img_2ch_t, img_2ch_ed, img_4ch_t, img_4ch_ed, \
        contour_sa, contour_2ch, contour_4ch, \
        vertex_ed, faces, affine_inv, affine, origin = batch

        with torch.no_grad():

            x_sa_t = img_sa_t.type(Tensor)
            x_sa_ed = img_sa_ed.type(Tensor)
            x_2ch_t = img_2ch_t.type(Tensor)
            x_2ch_ed = img_2ch_ed.type(Tensor)
            x_4ch_t = img_4ch_t.type(Tensor)
            x_4ch_ed = img_4ch_ed.type(Tensor)

            x_sa_t_5D = img_sa_t.unsqueeze(1).type(Tensor)
            x_sa_ed_5D = img_sa_ed.unsqueeze(1).type(Tensor)


            con_sa = contour_sa.type(TensorLong)  # [bs, slices, H, W]
            con_2ch = contour_2ch.type(TensorLong)  # [bs, H, W]
            con_4ch = contour_4ch.type(TensorLong)  # [bs, H, W]

            aff_sa_inv = affine_inv[:, 0, :, :].type(Tensor)
            aff_sa = affine[:, 0, :, :].type(Tensor)
            aff_2ch_inv = affine_inv[:, 1, :, :].type(Tensor)
            aff_4ch_inv = affine_inv[:, 2, :, :].type(Tensor)

            origin_sa = origin[:, 0:1, :].type(Tensor)
            origin_2ch = origin[:, 1:2, :].type(Tensor)
            origin_4ch = origin[:, 2:3, :].type(Tensor)

            vertex_0 = vertex_ed.permute(0, 2, 1).type(Tensor)  # [bs, 3, number of vertices]
            faces_0 = faces.type(Tensor)  # [bs, number of faces, 3]


            net_la = MV_LA(x_2ch_t, x_2ch_ed, x_4ch_t, x_4ch_ed)
            net_sa = MotionNet(x_sa_t, x_sa_ed, net_la['conv2_2ch'], net_la['conv2s_2ch'], net_la['conv2_4ch'],
                               net_la['conv2s_4ch'])

            # ---------------sample from 3D motion fields
            # translate coordinate
            v_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], vertex_0) + aff_sa_inv[:, :3, 3:4]
            v_ed = v_ed_o.permute(0, 2, 1) - origin_sa  # [bs, number of vertices,3]
            v_ed_x = (v_ed[:, :, 0:1] - (width / 2)) / (width / 2)
            v_ed_y = (v_ed[:, :, 1:2] - (height / 2)) / (height / 2)
            v_ed_z = (v_ed[:, :, 2:3] - (depth / 2)) / (depth / 2)
            v_ed_norm = torch.cat((v_ed_x, v_ed_y, v_ed_z), 2)
            v_ed_norm_expand = v_ed_norm.unsqueeze(1).unsqueeze(1)  # [bs, 1, 1,number of vertices,3]

            # sample from 3D motion field
            pxx = F.grid_sample(net_sa['out'][:, 0:1], v_ed_norm_expand, align_corners=True).transpose(4, 3)
            pyy = F.grid_sample(net_sa['out'][:, 1:2], v_ed_norm_expand, align_corners=True).transpose(4, 3)
            pzz = F.grid_sample(net_sa['out'][:, 2:3], v_ed_norm_expand, align_corners=True).transpose(4, 3)
            delta_p = torch.cat((pxx, pyy, pzz), 4)
            # updata coor (image space)
            # print (v_ed.shape, delta_p.shape)
            v_t_norm_expand = v_ed_norm_expand + delta_p  # [bs, 1, 1,number of vertices,3]
            # t frame
            v_t_norm = v_t_norm_expand.squeeze(1).squeeze(1)
            v_t_x = v_t_norm[:, :, 0:1] * (width / 2) + (width / 2)
            v_t_y = v_t_norm[:, :, 1:2] * (height / 2) + (height / 2)
            v_t_z = v_t_norm[:, :, 2:3] * (depth / 2) + (depth / 2)
            v_t_crop = torch.cat((v_t_x, v_t_y, v_t_z), 2)
            # translate back to mesh space
            v_t = v_t_crop + origin_sa  # [bs, number of vertices,3]
            pred_vertex_t = torch.matmul(aff_sa[:, :3, :3], v_t.permute(0, 2, 1)) + aff_sa[:, :3,
                                                                                    3:4]  # [bs, 3, number of vertices]
            # print (pred_vertex_t.shape)


            pred_sa_ed = transform(x_sa_t_5D, net_sa['out'], mode='bilinear')

            # -------------- differentialable slicer

            # coordinate transformation np.dot(aff_sa_SR_inv[:3,:3], points_ED.T) + aff_sa_SR_inv[:3,3:4]
            v_sa_hat_t_o = torch.matmul(aff_sa_inv[:, :3, :3], pred_vertex_t) + aff_sa_inv[:, :3, 3:4]
            v_sa_hat_t = v_sa_hat_t_o.permute(0, 2, 1) - origin_sa
            # print (v_sa_hat_t.shape)
            v_2ch_hat_t_o = torch.matmul(aff_2ch_inv[:, :3, :3], pred_vertex_t) + aff_2ch_inv[:, :3, 3:4]
            v_2ch_hat_t = v_2ch_hat_t_o.permute(0, 2, 1) - origin_2ch
            v_4ch_hat_t_o = torch.matmul(aff_4ch_inv[:, :3, :3], pred_vertex_t) + aff_4ch_inv[:, :3, 3:4]
            v_4ch_hat_t = v_4ch_hat_t_o.permute(0, 2, 1) - origin_4ch

            # project vertices satisfying threshood
            # project to SAX slices, project all vertices to a target plane,
            # vertices selection is moved to loss computation function
            v_sa_hat_t_x = torch.clamp(v_sa_hat_t[:, :, 0:1], min=0, max=height - 1)
            v_sa_hat_t_y = torch.clamp(v_sa_hat_t[:, :, 1:2], min=0, max=width - 1)
            v_sa_hat_t_cp = torch.cat((v_sa_hat_t_x, v_sa_hat_t_y, v_sa_hat_t[:, :, 2:3]), 2)

            v_sa_idx_t_0, w_sa_t_0 = projection(v_sa_hat_t_cp, 12, temper)
            # print (v_sa_idx_ed_0.shape, w_sa_ed_0.shape)
            v_sa_idx_t_1, w_sa_t_1 = projection(v_sa_hat_t_cp, 17, temper)
            v_sa_idx_t_2, w_sa_t_2 = projection(v_sa_hat_t_cp, 22, temper)
            v_sa_idx_t_3, w_sa_t_3 = projection(v_sa_hat_t_cp, 27, temper)
            v_sa_idx_t_4, w_sa_t_4 = projection(v_sa_hat_t_cp, 32, temper)
            v_sa_idx_t_5, w_sa_t_5 = projection(v_sa_hat_t_cp, 37, temper)
            v_sa_idx_t_6, w_sa_t_6 = projection(v_sa_hat_t_cp, 42, temper)
            v_sa_idx_t_7, w_sa_t_7 = projection(v_sa_hat_t_cp, 47, temper)
            v_sa_idx_t_8, w_sa_t_8 = projection(v_sa_hat_t_cp, 52, temper)

            # project to LAX 2CH view
            v_2ch_hat_t_x = torch.clamp(v_2ch_hat_t[:, :, 0:1], min=0, max=height - 1)
            v_2ch_hat_t_y = torch.clamp(v_2ch_hat_t[:, :, 1:2], min=0, max=width - 1)
            v_2ch_hat_t_cp = torch.cat((v_2ch_hat_t_x, v_2ch_hat_t_y, v_2ch_hat_t[:, :, 2:3]), 2)

            v_2ch_idx_t, w_2ch_t = projection(v_2ch_hat_t_cp, 0, temper)

            # project to LAX 4CH view
            v_4ch_hat_t_x = torch.clamp(v_4ch_hat_t[:, :, 0:1], min=0, max=height - 1)
            v_4ch_hat_t_y = torch.clamp(v_4ch_hat_t[:, :, 1:2], min=0, max=width - 1)
            v_4ch_hat_t_cp = torch.cat((v_4ch_hat_t_x, v_4ch_hat_t_y, v_4ch_hat_t[:, :, 2:3]), 2)

            v_4ch_idx_t, w_4ch_t = projection(v_4ch_hat_t_cp, 0, temper)

            # --------------------- Segmentation loss------------------
            loss_seg_sa_t_0 = weightedHausdorff_batch(v_sa_idx_t_0, w_sa_t_0, con_sa[:, 0, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_1 = weightedHausdorff_batch(v_sa_idx_t_1, w_sa_t_1, con_sa[:, 1, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_2 = weightedHausdorff_batch(v_sa_idx_t_2, w_sa_t_2, con_sa[:, 2, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_3 = weightedHausdorff_batch(v_sa_idx_t_3, w_sa_t_3, con_sa[:, 3, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_4 = weightedHausdorff_batch(v_sa_idx_t_4, w_sa_t_4, con_sa[:, 4, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_5 = weightedHausdorff_batch(v_sa_idx_t_5, w_sa_t_5, con_sa[:, 5, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_6 = weightedHausdorff_batch(v_sa_idx_t_6, w_sa_t_6, con_sa[:, 6, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_7 = weightedHausdorff_batch(v_sa_idx_t_7, w_sa_t_7, con_sa[:, 7, :, :], height, width, temper,
                                                      'val')
            loss_seg_sa_t_8 = weightedHausdorff_batch(v_sa_idx_t_8, w_sa_t_8, con_sa[:, 8, :, :], height, width, temper,
                                                      'val')
            loss_seg_2ch_t = weightedHausdorff_batch(v_2ch_idx_t, w_2ch_t, con_2ch, height, width, temper, 'val')
            loss_seg_4ch_t = weightedHausdorff_batch(v_4ch_idx_t, w_4ch_t, con_4ch, height, width, temper, 'val')

            loss_seg = (loss_seg_sa_t_0 + loss_seg_sa_t_1 + loss_seg_sa_t_2 + loss_seg_sa_t_3 +
                        loss_seg_sa_t_4 + loss_seg_sa_t_5 + loss_seg_sa_t_6 + loss_seg_sa_t_7 + loss_seg_sa_t_8) / 9.0 + \
                       loss_seg_2ch_t + loss_seg_4ch_t

            # ----------------smoothness loss------------
            # print (pred_vertex_t.permute(0,2,1).shape)
            trg_mesh = Meshes(verts=list(pred_vertex_t.permute(0, 2, 1)), faces=list(faces_0))
            loss_smooth = loss.mesh_laplacian_smoothing(trg_mesh, method='uniform')

            # ----------------regularization loss------------

            loss_reg = flow_criterion(pred_sa_ed, x_sa_ed_5D)

            loss_huber = huber_loss_3d(net_sa['out'])

            loss_all = loss_seg + w_reg * loss_reg + w_smooth * loss_smooth + w_h * loss_huber


            val_loss.append(loss_all.item())
            val_seg_loss.append(loss_seg.item())
            val_smooth_loss.append(loss_smooth.item())
            val_reg_loss.append(loss_reg.item())
            val_huber_loss.append(loss_huber.item())

            if batch_idx == 1:
                # tensorboard visulisation
                writer.add_scalar("Loss/val", loss_all, epoch * len(training_data_loader) + batch_idx)
                writer.add_scalar("Loss/val_seg", loss_seg, epoch * len(training_data_loader) + batch_idx)
                writer.add_scalar("Loss/val_reg", loss_reg, epoch * len(training_data_loader) + batch_idx)
                writer.add_scalar("Loss/val_smooth", loss_smooth, epoch * len(training_data_loader) + batch_idx)
                writer.add_scalar("Loss/val_huber", loss_huber, epoch * len(training_data_loader) + batch_idx)


    if np.mean(val_loss) < base_err:
        torch.save(MotionNet.state_dict(), Motion_save_path)
        torch.save(MV_LA.state_dict(), Motion_LA_save_path)
        base_err = np.mean(val_loss)



data_path = '/train_data_path'
train_set = TrainDataset(data_path)
# loading the data
training_data_loader = DataLoader(dataset=train_set, num_workers=n_worker, batch_size=bs, shuffle=True)

val_data_path = '/val_data_pathl'
val_set = ValDataset(val_data_path)
val_data_loader = DataLoader(dataset=val_set, num_workers=n_worker, batch_size=bs, shuffle=False)


for epoch in range(0, n_epoch + 1):
    start = time.time()
    train(epoch)
    end = time.time()
    print("training took {:.8f}".format(end-start))

    print('Epoch {}'.format(epoch))
    start = time.time()
    val(epoch)
    end = time.time()
    print("validation took {:.8f}".format(end - start))