import argparse
import logging
import multiprocessing as mp
import logging
import os
from detectron2.evaluation import inference_context
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from detectron2.utils.collect_env import collect_env_info
from detectron2.utils.logger import setup_logger
from fvcore.common.file_io import PathManager
from pathlib import Path
from pytorch3d.io import save_obj

from shapenet.config.config import get_shapenet_cfg
from shapenet.data.utils import imagenet_preprocess
from shapenet.modeling.heads import voxel_head
from shapenet.modeling.mesh_arch import build_model
from shapenet.utils.checkpoint import clean_state_dict

import torchvision.transforms as T

import glob
from PIL import Image

import trimesh
import pyvista as pv
import pyacvd
import numpy as np

logger = logging.getLogger('demo')

def setup_cfgs(args):
    cfg = get_shapenet_cfg()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    return cfg

def get_parser():
    parser = argparse.ArgumentParser(description="MeshRCNN Demo")
    parser.add_argument(
        "--config-file",
        default="configs/shapenet/voxmesh_R50.yaml",
        metavar="FILE",
        help="path to config file",
    )
    parser.add_argument("--input", help="A path to an input main folder")
    # parser.add_argument("--output", help="A directory to save output visualizations")
    parser.add_argument(
        "--focal-length", type=float, default=20.0, help="Focal length for the image"
    )
    parser.add_argument(
        "--onlyhighest", action="store_true", help="will return only the highest scoring detection"
    )

    parser.add_argument(
        "opts",
        help="Modify model config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    return parser

def resample_mesh(mesh, count=2466):
    pv_mesh = pv.wrap(mesh)
    # logger.info('Original mesh:')
    # print(pv_mesh)
    
    clus = pyacvd.Clustering(pv_mesh)
    clus.subdivide(3)
    clus.cluster(count)

    # remesh
    remesh = clus.create_mesh()

    # verts = remesh.points
    # faces = remesh.faces.reshape((-1, 4))[:, 1:]
    
    return remesh

if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    args = get_parser().parse_args()

    device = torch.device("cuda:%d" % 0)

    logger = setup_logger(name="demo shapenet")
    logger.info("Arguments: " + str(args))

    cfg = setup_cfgs(args)

    # load checkpoing and build model
    if cfg.MODEL.CHECKPOINT == "":
        raise ValueError("Invalid checkpoing provided")
    logger.info("Loading model from checkpoint: %s" % (cfg.MODEL.CHECKPOINT))
    cp = torch.load(PathManager.get_local_path(cfg.MODEL.CHECKPOINT))
    state_dict = clean_state_dict(cp["best_states"]["model"])
    model = build_model(cfg)
    model.load_state_dict(state_dict)
    logger.info("Model loaded")
    model.to(device)

    sub_dir = sorted(os.listdir(args.input))

    for sd in sub_dir:
        curr_path = os.path.join(args.input, sd)
        images = glob.glob(curr_path + "/*.png")
        
        for img_dir in images:
            # load image
            transform = [T.ToTensor()]
            transform.append(imagenet_preprocess())
            transform = T.Compose(transform)
            
            im_name = img_dir.split("/")[-1].split(".")[0]

            with PathManager.open(img_dir, "rb") as f:
                img = Image.open(f).convert("RGB")

            img = transform(img)
            img = img[None, :, :, :]
            img = img.to(device)

            with inference_context(model):
                img_feats, voxel_scores, meshes_pred, P, cubified_meshes = model(img)

            # Save voxel_score
            voxel_odir = os.path.join(curr_path, "voxel_score")
            if not Path(voxel_odir).is_dir():
                os.mkdir(voxel_odir)

            voxel_file = os.path.join(voxel_odir, "%s.pt" % (im_name))
            torch.save(voxel_scores, voxel_file)

            # Save image features
            imgfeat_odir = os.path.join(curr_path, "img_feat")
            if not Path(imgfeat_odir).is_dir():
                os.mkdir(imgfeat_odir)

            img_feat_file = os.path.join(imgfeat_odir, "%s.pt" % (im_name))
            torch.save(img_feats, img_feat_file)

            # Save P
            p_odir = os.path.join(curr_path, "P")
            if not Path(p_odir).is_dir():
                os.mkdir(p_odir)

            p_file = os.path.join(p_odir, "%s.pt" % (im_name))
            torch.save(P, p_file)

            # Save cubified mesh
            cmesh_odir = os.path.join(curr_path, "cube_mesh")
            if not Path(cmesh_odir).is_dir():
                os.mkdir(cmesh_odir)

            cube_mesh_file = os.path.join(cmesh_odir, "%s_cube.obj" % (im_name))
            c_verts, c_faces = cubified_meshes[-1].get_mesh_verts_faces(0)
            save_obj(cube_mesh_file, c_verts, c_faces)

            # Save predicted mesh
            mesh_odir = os.path.join(curr_path, "final_mesh")
            if not Path(mesh_odir).is_dir():
                os.mkdir(mesh_odir)

            save_file = os.path.join(mesh_odir, "%s.obj" % (im_name))
            verts, faces = meshes_pred[-1].get_mesh_verts_faces(0)
            save_obj(save_file, verts, faces)
            logger.info("Predictions saved for %s/%s" % (curr_path.split('/')[-1], im_name))

[(32, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (35, 'shapenet.config.config.get_shapenet_cfg', 'get_shapenet_cfg', False, 'from shapenet.config.config import get_shapenet_cfg\n'), (42, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (67, 'pyvista.wrap', 'pv.wrap', True, 'import pyvista as pv\n'), (71, 'pyacvd.Clustering', 'pyacvd.Clustering', False, 'import pyacvd\n'), (84, 'torch.multiprocessing.set_start_method', 'mp.set_start_method', True, 'import torch.multiprocessing as mp\n'), (87, 'torch.device', 'torch.device', False, 'import torch\n'), (89, 'detectron2.utils.logger.setup_logger', 'setup_logger', False, 'from detectron2.utils.logger import setup_logger\n'), (99, 'shapenet.utils.checkpoint.clean_state_dict', 'clean_state_dict', False, 'from shapenet.utils.checkpoint import clean_state_dict\n'), (100, 'shapenet.modeling.mesh_arch.build_model', 'build_model', False, 'from shapenet.modeling.mesh_arch import build_model\n'), (98, 'fvcore.common.file_io.PathManager.get_local_path', 'PathManager.get_local_path', False, 'from fvcore.common.file_io import PathManager\n'), (105, 'os.listdir', 'os.listdir', False, 'import os\n'), (108, 'os.path.join', 'os.path.join', False, 'import os\n'), (109, 'glob.glob', 'glob.glob', False, 'import glob\n'), (115, 'torchvision.transforms.Compose', 'T.Compose', True, 'import torchvision.transforms as T\n'), (130, 'os.path.join', 'os.path.join', False, 'import os\n'), (134, 'os.path.join', 'os.path.join', False, 'import os\n'), (135, 'torch.save', 'torch.save', False, 'import torch\n'), (138, 'os.path.join', 'os.path.join', False, 'import os\n'), (142, 'os.path.join', 'os.path.join', False, 'import os\n'), (143, 'torch.save', 'torch.save', False, 'import torch\n'), (146, 'os.path.join', 'os.path.join', False, 'import os\n'), (150, 'os.path.join', 'os.path.join', False, 'import os\n'), (151, 'torch.save', 'torch.save', False, 'import torch\n'), (154, 'os.path.join', 'os.path.join', False, 'import os\n'), (158, 'os.path.join', 'os.path.join', False, 'import os\n'), (160, 'pytorch3d.io.save_obj', 'save_obj', False, 'from pytorch3d.io import save_obj\n'), (163, 'os.path.join', 'os.path.join', False, 'import os\n'), (167, 'os.path.join', 'os.path.join', False, 'import os\n'), (169, 'pytorch3d.io.save_obj', 'save_obj', False, 'from pytorch3d.io import save_obj\n'), (113, 'torchvision.transforms.ToTensor', 'T.ToTensor', True, 'import torchvision.transforms as T\n'), (114, 'shapenet.data.utils.imagenet_preprocess', 'imagenet_preprocess', False, 'from shapenet.data.utils import imagenet_preprocess\n'), (119, 'fvcore.common.file_io.PathManager.open', 'PathManager.open', False, 'from fvcore.common.file_io import PathManager\n'), (126, 'detectron2.evaluation.inference_context', 'inference_context', False, 'from detectron2.evaluation import inference_context\n'), (132, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (140, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (148, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (156, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (165, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (120, 'PIL.Image.open', 'Image.open', False, 'from PIL import Image\n'), (131, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (139, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (147, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (155, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (164, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n')]