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_reconstruction import *

from dataio_reconstruction import *

from utils import *

lr = 1e-4

n_worker = 4

bs = 5

n_epoch = 400

base_err = 10000

w_smooth = 20

w_surface = 0.5

w_h = 0.5

width = 128

height = 128

depth = 64

temper = 3

model_save_path = './models/model_reconstruction'

if not os.path.exists(model_save_path):

    os.makedirs(model_save_path)

# pytorch only saves the last model

Deform_save_path = os.path.join(model_save_path, 'deform.pth')

Motion_LA_save_path = os.path.join(model_save_path, 'multiview.pth')

DeformNet = deformnet().cuda()

MV_LA = Mesh_2d().cuda()

optimizer = optim.Adam(filter(lambda p: p.requires_grad,

                              itertools.chain(DeformNet.parameters(), MV_LA.parameters())), lr=lr)

Tensor = torch.cuda.FloatTensor

TensorLong = torch.cuda.LongTensor

# visualisation

writer = SummaryWriter('./runs/model_reconstruction')

def train(epoch):

    DeformNet.train()

    MV_LA.train()

    epoch_loss = []

    epoch_seg_loss = []

    epoch_smooth_loss = []

    epoch_surface_loss = []

    epoch_huber_loss = []

    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_ed, contour_2ch_ed, contour_4ch_ed, \

        vertex_tpl_ed, faces_tpl, affine_inv, affine, origin, vertex_ed, mesh2seg_sa, mesh2seg_2ch, mesh2seg_4ch = batch

        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))

        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_tpl_0 = Variable(vertex_tpl_ed.permute(0,2,1).type(Tensor)) # [bs, 3, number of vertices]

        faces_tpl_0 = Variable(faces_tpl.type(Tensor)) # [bs, number of faces, 3]

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

        mesh2seg_sa_gt = Variable(mesh2seg_sa.type(Tensor))

        mesh2seg_2ch_gt = Variable(mesh2seg_2ch.type(Tensor))

        mesh2seg_4ch_gt = Variable(mesh2seg_4ch.type(Tensor))

        optimizer.zero_grad()

        net_la = MV_LA(x_2ch_t, x_2ch_ed, x_4ch_t, x_4ch_ed)

        net_df = DeformNet(x_sa_ed, net_la['conv2s_2ch'], net_la['conv2s_4ch'])

        # ---------------sample from 3D motion fields

        # translate coordinate

        v_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], vertex_tpl_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_df['out_def_ed'][:, 0:1], v_ed_norm_expand, align_corners=True).transpose(4, 3)

        pyy = F.grid_sample(net_df['out_def_ed'][:, 1:2], v_ed_norm_expand, align_corners=True).transpose(4, 3)

        pzz = F.grid_sample(net_df['out_def_ed'][:, 2:3], v_ed_norm_expand, align_corners=True).transpose(4, 3)

        # print (pxx.shape, pyy.shape, pzz.shape)

        delta_p = torch.cat((pxx, pyy, pzz), 4)

        # updata coor (image space)

        # print (v_ed.shape, delta_p.shape)

        v_0_norm_expand = v_ed_norm_expand + delta_p  # [bs, 1, 1,number of vertices,3]

        # t frame

        v_0_norm = v_0_norm_expand.squeeze(1).squeeze(1)

        v_0_x = v_0_norm[:, :, 0:1] * (width / 2) + (width / 2)

        v_0_y = v_0_norm[:, :, 1:2] * (height / 2) + (height / 2)

        v_0_z = v_0_norm[:, :, 2:3] * (depth / 2) + (depth / 2)

        v_0_crop = torch.cat((v_0_x, v_0_y, v_0_z), 2)

        # translate back to mesh space

        v_0 = v_0_crop + origin_sa  # [bs, number of vertices,3]

        pred_v_0 = torch.matmul(aff_sa[:, :3, :3], v_0.permute(0, 2, 1)) + aff_sa[:, :3,3:4]  # [bs, 3, number of vertices]

        # print (pred_vertex_t.shape)

        # -------------- 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_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], pred_v_0) + aff_sa_inv[:, :3, 3:4]

        v_sa_hat_ed = v_sa_hat_ed_o.permute(0, 2, 1) - origin_sa

        # print (v_sa_hat_t.shape)

        v_2ch_hat_ed_o = torch.matmul(aff_2ch_inv[:, :3, :3], pred_v_0) + aff_2ch_inv[:, :3, 3:4]

        v_2ch_hat_ed = v_2ch_hat_ed_o.permute(0, 2, 1) - origin_2ch

        v_4ch_hat_ed_o = torch.matmul(aff_4ch_inv[:, :3, :3], pred_v_0) + aff_4ch_inv[:, :3, 3:4]

        v_4ch_hat_ed = v_4ch_hat_ed_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_ed_x = torch.clamp(v_sa_hat_ed[:, :, 0:1], min=0, max=height - 1)

        v_sa_hat_ed_y = torch.clamp(v_sa_hat_ed[:, :, 1:2], min=0, max=width - 1)

        v_sa_hat_ed_cp = torch.cat((v_sa_hat_ed_x, v_sa_hat_ed_y, v_sa_hat_ed[:, :, 2:3]), 2)

        # project to LAX 2CH view

        v_2ch_hat_ed_x = torch.clamp(v_2ch_hat_ed[:, :, 0:1], min=0, max=height - 1)

        v_2ch_hat_ed_y = torch.clamp(v_2ch_hat_ed[:, :, 1:2], min=0, max=width - 1)

        v_2ch_hat_ed_cp = torch.cat((v_2ch_hat_ed_x, v_2ch_hat_ed_y, v_2ch_hat_ed[:, :, 2:3]), 2)

        v_2ch_idx_ed, w_2ch_ed = projection(v_2ch_hat_ed_cp, 0, temper)

        # project to LAX 4CH view

        v_4ch_hat_ed_x = torch.clamp(v_4ch_hat_ed[:, :, 0:1], min=0, max=height - 1)

        v_4ch_hat_ed_y = torch.clamp(v_4ch_hat_ed[:, :, 1:2], min=0, max=width - 1)

        v_4ch_hat_ed_cp = torch.cat((v_4ch_hat_ed_x, v_4ch_hat_ed_y, v_4ch_hat_ed[:, :, 2:3]), 2)

        v_4ch_idx_ed, w_4ch_ed = projection(v_4ch_hat_ed_cp, 0, temper)

        # --------------------- Segmentation loss------------------

        loss_seg_sa_ed = projection_weightHD_loss_SA(v_sa_hat_ed_cp, temper, height, width, depth, mesh2seg_sa_gt, 'train')

        loss_seg_2ch_ed = weightedHausdorff_batch(v_2ch_idx_ed, w_2ch_ed, mesh2seg_2ch_gt, height, width, temper, 'train')

        loss_seg_4ch_ed = weightedHausdorff_batch(v_4ch_idx_ed, w_4ch_ed, mesh2seg_4ch_gt, height, width, temper, 'train')

        loss_seg = loss_seg_sa_ed + loss_seg_2ch_ed + loss_seg_4ch_ed

        #----------------smoothness loss------------

        trg_mesh_ed = Meshes(verts=list(pred_v_0.permute(0, 2, 1)), faces=list(faces_tpl_0))

        loss_laplacian_smooth = loss.mesh_laplacian_smoothing(trg_mesh_ed, method='uniform')

        loss_smooth = loss_laplacian_smooth

        # ------------------J loss---------------------

        loss_huber = huber_loss_3d(net_df['out_def_ed'])

        # ------------------Surface chamfer loss---------------------

        loss_surface, _ = loss.chamfer_distance(pred_v_0.permute(0, 2, 1), vertex_0.permute(0, 2, 1))

        loss_all = loss_seg + w_surface * loss_surface + 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_surface_loss.append(loss_surface.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_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)

        writer.add_scalar("Loss/train_surface", loss_surface, epoch * len(training_data_loader) + batch_idx)

        if batch_idx % 40 == 0:

            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss all: {:.6f}, '

                  'Seg Loss: {:.6f}, Smooth Loss: {:.6f}, Surface Loss: {:.6f}, Huger 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_smooth_loss), np.mean(epoch_surface_loss), np.mean(epoch_huber_loss)))

def val(epoch):

    DeformNet.eval()

    MV_LA.eval()

    val_loss = []

    val_seg_loss = []

    val_smooth_loss = []

    val_surface_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_ed, contour_2ch_ed, contour_4ch_ed, \

        vertex_tpl_ed, faces_tpl, affine_inv, affine, origin, vertex_ed, mesh2seg_sa, mesh2seg_2ch, mesh2seg_4ch = batch

        with torch.no_grad():

            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)

            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_tpl_0 = vertex_tpl_ed.permute(0, 2, 1).type(Tensor)  # [bs, 3, number of vertices]

            faces_tpl_0 = faces_tpl.type(Tensor)  # [bs, number of faces, 3]

            vertex_0 = vertex_ed.permute(0, 2, 1).cuda()  # [bs, 3, number of vertices]

            mesh2seg_sa_gt = Variable(mesh2seg_sa.type(Tensor))

            mesh2seg_2ch_gt = Variable(mesh2seg_2ch.type(Tensor))

            mesh2seg_4ch_gt = Variable(mesh2seg_4ch.type(Tensor))

            net_la = MV_LA(x_2ch_t, x_2ch_ed, x_4ch_t, x_4ch_ed)

            net_df = DeformNet(x_sa_ed, net_la['conv2s_2ch'], net_la['conv2s_4ch'])

            # ---------------sample from 3D motion fields

            # translate coordinate

            v_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], vertex_tpl_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_df['out_def_ed'][:, 0:1], v_ed_norm_expand, align_corners=True).transpose(4, 3)

            pyy = F.grid_sample(net_df['out_def_ed'][:, 1:2], v_ed_norm_expand, align_corners=True).transpose(4, 3)

            pzz = F.grid_sample(net_df['out_def_ed'][:, 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_0_norm_expand = v_ed_norm_expand + delta_p  # [bs, 1, 1,number of vertices,3]

            # t frame

            v_0_norm = v_0_norm_expand.squeeze(1).squeeze(1)

            v_0_x = v_0_norm[:, :, 0:1] * (width / 2) + (width / 2)

            v_0_y = v_0_norm[:, :, 1:2] * (height / 2) + (height / 2)

            v_0_z = v_0_norm[:, :, 2:3] * (depth / 2) + (depth / 2)

            v_0_crop = torch.cat((v_0_x, v_0_y, v_0_z), 2)

            # translate back to mesh space

            v_0 = v_0_crop + origin_sa  # [bs, number of vertices,3]

            pred_v_0 = torch.matmul(aff_sa[:, :3, :3], v_0.permute(0, 2, 1)) + aff_sa[:, :3,

                                                                               3:4]  # [bs, 3, number of vertices]

            # -------------- 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_ed_o = torch.matmul(aff_sa_inv[:, :3, :3], pred_v_0) + aff_sa_inv[:, :3, 3:4]

            v_sa_hat_ed = v_sa_hat_ed_o.permute(0, 2, 1) - origin_sa

            # print (v_sa_hat_t.shape)

            v_2ch_hat_ed_o = torch.matmul(aff_2ch_inv[:, :3, :3], pred_v_0) + aff_2ch_inv[:, :3, 3:4]

            v_2ch_hat_ed = v_2ch_hat_ed_o.permute(0, 2, 1) - origin_2ch

            v_4ch_hat_ed_o = torch.matmul(aff_4ch_inv[:, :3, :3], pred_v_0) + aff_4ch_inv[:, :3, 3:4]

            v_4ch_hat_ed = v_4ch_hat_ed_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_ed_x = torch.clamp(v_sa_hat_ed[:, :, 0:1], min=0, max=height - 1)

            v_sa_hat_ed_y = torch.clamp(v_sa_hat_ed[:, :, 1:2], min=0, max=width - 1)

            v_sa_hat_ed_cp = torch.cat((v_sa_hat_ed_x, v_sa_hat_ed_y, v_sa_hat_ed[:, :, 2:3]), 2)

            # project to LAX 2CH view

            v_2ch_hat_ed_x = torch.clamp(v_2ch_hat_ed[:, :, 0:1], min=0, max=height - 1)

            v_2ch_hat_ed_y = torch.clamp(v_2ch_hat_ed[:, :, 1:2], min=0, max=width - 1)

            v_2ch_hat_ed_cp = torch.cat((v_2ch_hat_ed_x, v_2ch_hat_ed_y, v_2ch_hat_ed[:, :, 2:3]), 2)

            v_2ch_idx_ed, w_2ch_ed = projection(v_2ch_hat_ed_cp, 0, temper)

            # project to LAX 4CH view

            v_4ch_hat_ed_x = torch.clamp(v_4ch_hat_ed[:, :, 0:1], min=0, max=height - 1)

            v_4ch_hat_ed_y = torch.clamp(v_4ch_hat_ed[:, :, 1:2], min=0, max=width - 1)

            v_4ch_hat_ed_cp = torch.cat((v_4ch_hat_ed_x, v_4ch_hat_ed_y, v_4ch_hat_ed[:, :, 2:3]), 2)

            v_4ch_idx_ed, w_4ch_ed = projection(v_4ch_hat_ed_cp, 0, temper)

            # --------------------- Segmentation loss------------------

            loss_seg_sa_ed = projection_weightHD_loss_SA(v_sa_hat_ed_cp, temper, height, width, depth, mesh2seg_sa_gt,

                                                         'val')

            loss_seg_2ch_ed = weightedHausdorff_batch(v_2ch_idx_ed, w_2ch_ed, mesh2seg_2ch_gt, height, width, temper,

                                                      'val')

            loss_seg_4ch_ed = weightedHausdorff_batch(v_4ch_idx_ed, w_4ch_ed, mesh2seg_4ch_gt, height, width, temper,

                                                      'val')

            loss_seg = loss_seg_sa_ed + loss_seg_2ch_ed + loss_seg_4ch_ed

            # ----------------smoothness loss------------

            # print (pred_vertex_t.permute(0,2,1).shape)

            trg_mesh_ed = Meshes(verts=list(pred_v_0.permute(0, 2, 1)), faces=list(faces_tpl_0))

            loss_laplacian_smooth = loss.mesh_laplacian_smoothing(trg_mesh_ed, method='uniform')

            loss_smooth = loss_laplacian_smooth

            # ------------------J loss---------------------

            loss_huber = huber_loss_3d(net_df['out_def_ed'])

            # ------------------Surface chamfer loss---------------------

            loss_surface, _ = loss.chamfer_distance(pred_v_0.permute(0, 2, 1), vertex_0.permute(0, 2, 1))

            loss_all = loss_seg + w_surface * loss_surface + 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_surface_loss.append(loss_surface.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_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)

                writer.add_scalar("Loss/val_surface", loss_surface, epoch * len(training_data_loader) + batch_idx)

    if np.mean(val_loss) < base_err:

        torch.save(DeformNet.state_dict(), Deform_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_path'

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))