# .objファイルの表示
import os
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from skimage.io import imread
import cv2
# Util function for loading meshes
from pytorch3d.io import load_obj,load_objs_as_meshes
# Data structures and functions for rendering
from pytorch3d.structures import Meshes, Textures
from pytorch3d.renderer import (
    look_at_view_transform,
    OpenGLPerspectiveCameras, 
    PointLights, 
    DirectionalLights, 
    Materials, 
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    TexturedPhongShader
)
# 3D transformations functions
from pytorch3d.transforms import Rotate, Translate

# 予めplot_image_grid.pyをダウンロードして同じディレクトリに置いています
#!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/master/docs/tutorials/utils/plot_image_grid.py
from plot_image_grid import image_grid

# add path for demo utils functions 
import sys
import os
sys.path.append(os.path.abspath(''))

# Setup
device = torch.device("cuda:0")
torch.cuda.set_device(device)

# Set paths
DATA_DIR = "./data"

# クラス定義
class Model(nn.Module):
    def __init__(self, device, filename):
        super().__init__()
        # バッチサイズの設定 - レンダリングする異なる視点の数
        self.batch_size = 20
        # メッシュの読込み(バッチサイズ分作成)
        self.meshes = self.loadMeshes(device, filename)
        self.device = device

        # Get a batch of viewing angles.
        elev = torch.linspace(0, 360, self.batch_size)
        azim = torch.linspace(0, 360, self.batch_size)
 
        # camerasのヘルパー関数は、入力とブロードキャストの混合型をサポートしています。
        # そのため、距離は1.7で固定して、角度だけ変更するようなことも可能です。
        R, T = look_at_view_transform(dist=2.0, elev=elev, azim=azim)
        cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)

        # ラスタライザの設定
        raster_settings = RasterizationSettings(
            image_size=512, 
            blur_radius=0.0, 
            faces_per_pixel=1, 
            bin_size=0
        )

        # ポイントライトの設置
        lights = PointLights(device=self.device, location=[[0.5, 2.5, -4.0]])

        # フォンのレンダラの作成。フォンシェーダはテクスチャのuv座標を補間します。
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(
                cameras=cameras, 
                raster_settings=raster_settings
            ),
            shader=TexturedPhongShader(
                device=self.device, 
                cameras=cameras,
                lights=lights
            )
        )
        self.renderer = renderer
        self.filename = filename
        self.images = None

    # forward
    def forward(self):
        # メッシュの描画
        self.images = self.renderer(meshes_world=self.meshes)
        return self.images
    
    # objファイルの読込み
    def loadMeshes(self, device, filename):
        verts, faces, aux = load_obj(obj_filename)
        vertNum = verts.shape[0]
        faceNum =  faces.verts_idx.shape[0]
        print("vertNum=", vertNum)
        print("faceNum=", faceNum)

        # get vertex_uvs
        if aux.verts_uvs is not None:
            verts_uvs = aux.verts_uvs[None,...]
        else:
            # 0～1の乱数を適当に設定
            verts_uvs = torch.rand(1, vertNum, 2)
        print("verts_uvs.shape=", verts_uvs.shape)

        # get face_uvs
        if len(faces.textures_idx) > 0:
            faces_uvs = faces.textures_idx[None,...]
        else:
            # 面のインデックスで代用
            faces_uvs = faces.verts_idx[None]
        print("faces_uvs.shape=", faces_uvs.shape)

        # get texture data
        if aux.texture_images is not None:
            tex_maps = aux.texture_images
            texture_image = list(tex_maps.values())[0]
            texture_image = texture_image[None, ...]
        else:
            # ダミーの白色画像を準備
            texture_image = imread('data/white.jpg')
            texture_image = torch.from_numpy(texture_image/255.0)[None]
            texture_image = texture_image.float()
        print("texture_image.shape=", texture_image.shape)

        tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image)

        # centering, scaling
        center = verts.mean(0)
        verts = verts - center
        scale = max(verts.abs().max(0)[0])
        verts = verts / scale
        
        meshes = Meshes(verts=[verts], faces=[faces.verts_idx], textures=tex).to(device)

        # バッチサイズ分のメッシュを作る必要があります。
        # `extend`関数を使うと、簡単に作成できます。テクスチャも拡張してくれます。
        meshes = meshes.extend(self.batch_size)

        return meshes
    
    # タイル画像として保存
    # https://note.nkmk.me/python-opencv-hconcat-vconcat-np-tile/
    def saveTileImage(self, filename):
        im = self.images.cpu().numpy()
        for i in range(20):
            im[i] = cv2.cvtColor(255*im[i], cv2.COLOR_BGRA2RGBA)
        im_tile = self.concat_tile([[im[0], im[1], im[2], im[3], im[4]],
                                    [im[5], im[6], im[7], im[8], im[9]],
                                    [im[10], im[11], im[12], im[13], im[14]],
                                    [im[15], im[16], im[17], im[18], im[19]]])
        cv2.imwrite(filename, im_tile)
    
    def concat_tile(self, im_list_2d):
        return cv2.vconcat([cv2.hconcat(im_list_h) for im_list_h in im_list_2d])
        
        

# ファイル名の指定
#obj_filename = os.path.join(DATA_DIR, "tda_append/tda_append.obj")
obj_filename = os.path.join(DATA_DIR, "final_model.obj")
#obj_filename = os.path.join(DATA_DIR, "cow_mesh/cow.obj")

# ファイル名を指定してレンダリング準備
model = Model(device=device, filename=obj_filename).to(device)

vertNum= 2562
faceNum= 5120
verts_uvs.shape= torch.Size([1, 2562, 2])
faces_uvs.shape= torch.Size([1, 5120, 3])
texture_image.shape= torch.Size([1, 100, 100, 3])

# レンダリングの実行
images = model()
# グリッド画像の表示
image_grid(images.cpu().numpy(), rows=4, cols=5, rgb=True)
# 画像として保存
model.saveTileImage('tile_image.jpg')