# 0. モジュールのインポート

# 変形のサンプル
# .objの読込み方
# Meshesデータ構造の使い方
# PyTorch3dの4つの異なるmesh loss functions
# 最適化ループの設定方法
# chamfer_distanceだけではスムーズでない形状になってしまうので、shape regularizersを加える
import os
import torch
from pytorch3d.io import load_obj, save_obj
from pytorch3d.structures import Meshes
from pytorch3d.utils import ico_sphere
from pytorch3d.ops import sample_points_from_meshes
from pytorch3d.loss import (
    chamfer_distance,          # differentiably sampling points
    mesh_edge_loss,            # minimizes the length of the edges in the predicted mesh
    mesh_laplacian_smoothing,  # enforces consistency across the normals of neighboring faces
    mesh_normal_consistency,   # laplacian regularizer
)

import numpy as np
import tqdm
%matplotlib notebook 
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['savefig.dpi'] = 80
mpl.rcParams['figure.dpi'] = 80

# Set the device
device = torch.device("cuda:0")

# 1. objファイルの読込みとMeshesオブジェクトの作成
# !wget https://dl.fbaipublicfiles.com/pytorch3d/data/dolphin/dolphin.obj

# load the dolphin mesh.
trg_obj = os.path.join('data/dolphin.obj')

# pytorch3dの関数でobjファイルを読み込む
verts, faces, aux = load_obj(trg_obj)

# vertsは頂点座標。FloatTensor (V, 3)
# facesは面の情報。LongTensors: verts_idx, normals_idx, textures_idx
# このチュートリアルでは、法線とテクスチャは無視します
faces_idx = faces.verts_idx.to(device)
verts = verts.to(device)
print("verts.shape=", verts.shape)

# 中心が(0, 0, 0)で半径が1の球にフィットさせるために、メッシュサイズの正規化とセンタリングします
# (scale, center)は生成されたメッシュをオリジナルの中心と大きさにするために使います
# 注意:対象メッシュを正規化することは、最適化の高速化が可能ですが、必須というわけではありません
center = verts.mean(0)  # 頂点座標の平均値＝中心座標
print("center.shape={0}, center={1}".format(center.shape, center))
verts = verts - center  # 中心を(0, 0, 0)に移動
print("verts.abs().max(0)=", verts.abs().max(0)) # 値と、最大値を持つ頂点番号が取得できる
print("verts.abs().max(0)[0]=", verts.abs().max(0)[0]) # 今回は値が必要
scale = max(verts.abs().max(0)[0]) # Tensorの中で最大値をscaleに設定
print("scale=", scale)
verts = verts / scale  # 最大の絶対値で割ることで、最大値が1になる

# 正規化した頂点座標でメッシュを作成
trg_mesh = Meshes(verts=[verts], faces=[faces_idx])

# pytorch3d.utils.ico_sphere関数を使って、半径1の牛を生成
# ico_sphere(level, device): levelはメッシュの分割レベル
src_mesh = ico_sphere(4, device)
print("src_mesh.verts.shape=", src_mesh.verts_list()[0].shape)

verts.shape= torch.Size([2562, 3])
center.shape=torch.Size([3]), center=tensor([2.3951e-04, 3.4056e-01, 8.8234e-02], device='cuda:0')
verts.abs().max(0)= torch.return_types.max(
values=tensor([0.1417, 0.3086, 0.3716], device='cuda:0'),
indices=tensor([2023,  496, 1695], device='cuda:0'))
verts.abs().max(0)[0]= tensor([0.1417, 0.3086, 0.3716], device='cuda:0')
scale= tensor(0.3716, device='cuda:0')
src_mesh.verts.shape= torch.Size([2562, 3])

# 可視化
def plot_pointcloud(mesh, title=""):
    # Sample points uniformly from the surface of the mesh.
    points = sample_points_from_meshes(mesh, 5000) # メッシュの表面から5000点サンプリング
    x, y, z = points.clone().detach().cpu().squeeze().unbind(1)    # cpuに座標を転送？
    print("x.shape=", x.shape)
    fig = plt.figure(figsize=(5, 5))
    ax = Axes3D(fig)
    ax.scatter3D(x, z, -y)
    ax.set_xlabel('x')
    ax.set_ylabel('z')
    ax.set_zlabel('y')
    ax.set_title(title)
    ax.view_init(190, 30)
    plt.show()
    

# %matplotlib notebook
plot_pointcloud(trg_mesh, "Target mesh")
plot_pointcloud(src_mesh, "Source mesh")

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

# 3. 最適化ループ

# 頂点のオフセッティング（相殺？）による元メッシュの変形の仕方を学びます
# 変形のためのパラメータは、src_meshの全頂点数と等しくなります
# torch.full: 指定したサイズ行列を、指定した値で初期化する関数。
#             requires_grad: If autograd should record operations on the returned tensor. 
deform_verts = torch.full(src_mesh.verts_packed().shape, 0.0, device=device, requires_grad=True)
print("verts.shape=", src_mesh.verts_packed().shape)
print("deform_verts=", deform_verts.shape)

# optimizerはSGDです
optimizer = torch.optim.SGD([deform_verts], lr=1.0, momentum=0.9)

verts.shape= torch.Size([2562, 3])
deform_verts= torch.Size([2562, 3])

# 学習回数
# Number of optimization steps
Niter = 2000

# 各loss関数の重み
# Weight for the chamfer loss
w_chamfer = 1.0 
# Weight for mesh edge loss
w_edge = 1.0 
# Weight for mesh normal consistency
w_normal = 0.01 
# Weight for mesh laplacian smoothing
w_laplacian = 0.1 

# Plot period for the losses
plot_period = 250
loop = tqdm.notebook.tqdm(range(Niter))

chamfer_losses = []
laplacian_losses = []
edge_losses = []
normal_losses = []

%matplotlib inline

for i in loop:
    # Initialize optimizer
    optimizer.zero_grad()
    
    # Deform the mesh
    # offset_verts: 頂点座標にオフセットとして引数の値を加算する
    # deform_vertsは、生成時に0で初期化しているので、最初は初期形状と同じで、そこからdeform_vertsを最適化していく
    new_src_mesh = src_mesh.offset_verts(deform_verts)
    
    # We sample 5k points from the surface of each mesh 
    sample_trg = sample_points_from_meshes(trg_mesh, 5000)
    sample_src = sample_points_from_meshes(new_src_mesh, 5000)
    
    # We compare the two sets of pointclouds by computing (a) the chamfer loss
    # (a) ポイントクラウドのchamfer lossを計算
    loss_chamfer, _ = chamfer_distance(sample_trg, sample_src)
    
    # and (b) the edge length of the predicted mesh
    # (b) エッジの長さのlossを計算
    loss_edge = mesh_edge_loss(new_src_mesh)
    
    # mesh normal consistency
    # (c) 近接する法線ベクトルの一貫性のlossを計算
    loss_normal = mesh_normal_consistency(new_src_mesh)
    
    # mesh laplacian smoothing
    # (d) ラプラシアンスムージング
    loss_laplacian = mesh_laplacian_smoothing(new_src_mesh, method="uniform")
    
    # Weighted sum of the losses
    # 4つのlossの重み付き合計を計算
    loss = loss_chamfer * w_chamfer + loss_edge * w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian
    
    # Print the losses
    loop.set_description('total_loss = %.6f' % loss)
    
    # Save the losses for plotting
    chamfer_losses.append(loss_chamfer)
    edge_losses.append(loss_edge)
    normal_losses.append(loss_normal)
    laplacian_losses.append(loss_laplacian)
    
    # Plot mesh
    if i % plot_period == 0:
        plot_pointcloud(new_src_mesh, title="iter: %d" % i)
        
    # Optimization step
    # lossから勾配を計算
    loss.backward()
    # 最適化の実行
    optimizer.step()

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

x.shape= torch.Size([5000])

# 4. lossの可視化

fig = plt.figure(figsize=(13, 5))
ax = fig.gca() # get current axesの略？
ax.plot(chamfer_losses, label="chamfer loss")
ax.plot(edge_losses, label="edge loss")
ax.plot(normal_losses, label="normal loss")
ax.plot(laplacian_losses, label="laplacian loss")
ax.legend(fontsize="16")
ax.set_xlabel("Iteration", fontsize="16")
ax.set_ylabel("Loss", fontsize="16")
ax.set_title("Loss vs iterations", fontsize="16")

Text(0.5, 1.0, 'Loss vs iterations')

# 5. save the predicted mesh

# Fetch the verts and faces of the final predicted mesh
final_verts, final_faces = new_src_mesh.get_mesh_verts_faces(0)

# Scale normalize back to the original target size
final_verts = final_verts * scale + center

# Store the predicted mesh using save_obj
final_obj = os.path.join('./data/', 'final_model.obj')
save_obj(final_obj, final_verts, final_faces)