import argparse
import glob
import os
import shutil
import time

import numpy as np
import open3d as o3d
import torch
import trimesh
from plyfile import PlyData
from pytorch3d.io import load_objs_as_meshes
from pytorch3d.ops import sample_points_from_meshes
from pytorch3d.structures import Meshes
from skimage import measure
from torch.utils.tensorboard import SummaryWriter

from src.optimization import Trainer
from src.utils import (
    AverageMeter,
    adjust_learning_rate,
    export_pointcloud,
    get_learning_rate_schedules,
    initialize_logger,
    load_config,
    update_config,
    update_optimizer,
)

np.set_printoptions(precision=4)


def main():
    parser = argparse.ArgumentParser(description="MNIST toy experiment")
    parser.add_argument("config", type=str, help="Path to config file.")
    parser.add_argument("--no_cuda", action="store_true", default=False, help="disables CUDA training")
    parser.add_argument("--seed", type=int, default=1457, metavar="S", help="random seed")

    args, unknown = parser.parse_known_args()
    cfg = load_config(args.config, "configs/default.yaml")
    cfg = update_config(cfg, unknown)
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    data_type = cfg["data"]["data_type"]
    cfg["data"]["class"]

    print(cfg["train"]["out_dir"])

    # PYTORCH VERSION > 1.0.0
    assert float(torch.__version__.split(".")[-3]) > 0

    # boiler-plate
    if cfg["train"]["timestamp"]:
        cfg["train"]["out_dir"] += "_" + time.strftime("%Y_%m_%d_%H_%M_%S")
    logger = initialize_logger(cfg)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    shutil.copyfile(args.config, os.path.join(cfg["train"]["out_dir"], "config.yaml"))

    # tensorboardX writer
    tblogdir = os.path.join(cfg["train"]["out_dir"], "tensorboard_log")
    if not os.path.exists(tblogdir):
        os.makedirs(tblogdir)
    SummaryWriter(log_dir=tblogdir)

    # initialize o3d visualizer
    vis = None
    if cfg["train"]["o3d_show"]:
        vis = o3d.visualization.Visualizer()
        vis.create_window(width=cfg["train"]["o3d_window_size"], height=cfg["train"]["o3d_window_size"])

    # initialize dataset
    if data_type == "point":
        if cfg["data"]["object_id"] != -1:
            data_paths = sorted(glob.glob(cfg["data"]["data_path"]))
            data_path = data_paths[cfg["data"]["object_id"]]
            print("Loaded %d/%d object" % (cfg["data"]["object_id"] + 1, len(data_paths)))
        else:
            data_path = cfg["data"]["data_path"]
            print("Data loaded")
        ext = data_path.split(".")[-1]
        if ext == "obj":  # have GT mesh
            mesh = load_objs_as_meshes([data_path], device=device)
            # scale the mesh into unit cube
            verts = mesh.verts_packed()
            N = verts.shape[0]
            center = verts.mean(0)
            mesh.offset_verts_(-center.expand(N, 3))
            scale = max((verts - center).abs().max(0)[0])
            mesh.scale_verts_(1.0 / float(scale))
            # important for our DPSR to have the range in [0, 1), not reaching 1
            mesh.scale_verts_(0.9)

            target_pts, target_normals = sample_points_from_meshes(mesh, num_samples=200000, return_normals=True)
        elif ext == "ply":  # only have the point cloud
            plydata = PlyData.read(data_path)
            vertices = np.stack([plydata["vertex"]["x"], plydata["vertex"]["y"], plydata["vertex"]["z"]], axis=1)
            normals = np.stack([plydata["vertex"]["nx"], plydata["vertex"]["ny"], plydata["vertex"]["nz"]], axis=1)
            N = vertices.shape[0]
            center = vertices.mean(0)
            scale = np.max(np.max(np.abs(vertices - center), axis=0))
            vertices -= center
            vertices /= scale
            vertices *= 0.9

            target_pts = torch.tensor(vertices, device=device)[None].float()
            target_normals = torch.tensor(normals, device=device)[None].float()
            mesh = None  # no GT mesh

        if not torch.is_tensor(center):
            center = torch.from_numpy(center)
        if not torch.is_tensor(scale):
            scale = torch.from_numpy(np.array([scale]))

        data = {
            "target_points": target_pts,
            "target_normals": target_normals,  # normals are never used
            "gt_mesh": mesh,
        }

    else:
        raise NotImplementedError

    # save the input point cloud
    if "target_points" in data.keys():
        outdir_pcl = os.path.join(cfg["train"]["out_dir"], "target_pcl.ply")
        if "target_normals" in data.keys():
            export_pointcloud(outdir_pcl, data["target_points"], data["target_normals"])
        else:
            export_pointcloud(outdir_pcl, data["target_points"])

    # save oracle PSR mesh (mesh from our PSR using GT point+normals)
    if data.get("gt_mesh") is not None:
        gt_verts, gt_faces = data["gt_mesh"].get_mesh_verts_faces(0)
        pts_gt, norms_gt = sample_points_from_meshes(data["gt_mesh"], num_samples=500000, return_normals=True)
        pts_gt = (pts_gt + 1) / 2
        from src.dpsr import DPSR

        dpsr_tmp = DPSR(
            res=(cfg["model"]["grid_res"], cfg["model"]["grid_res"], cfg["model"]["grid_res"]),
            sig=cfg["model"]["psr_sigma"],
        ).to(device)
        target = dpsr_tmp(pts_gt, norms_gt).unsqueeze(1).to(device)
        target = torch.tanh(target)
        s = target.shape[-1]  # size of psr_grid
        psr_grid_numpy = target.squeeze().detach().cpu().numpy()
        verts, faces, _, _ = measure.marching_cubes(psr_grid_numpy)
        verts = verts / s * 2.0 - 1  # [-1, 1]
        mesh = o3d.geometry.TriangleMesh()
        mesh.vertices = o3d.utility.Vector3dVector(verts)
        mesh.triangles = o3d.utility.Vector3iVector(faces)
        outdir_mesh = os.path.join(cfg["train"]["out_dir"], "oracle_mesh.ply")
        o3d.io.write_triangle_mesh(outdir_mesh, mesh)

    # initialize the source point cloud given an input mesh
    if "input_mesh" in cfg["train"].keys() and os.path.isfile(cfg["train"]["input_mesh"]):
        if cfg["train"]["input_mesh"].split("/")[-2] == "mesh":
            mesh_tmp = trimesh.load_mesh(cfg["train"]["input_mesh"])
            verts = torch.from_numpy(mesh_tmp.vertices[None]).float().to(device)
            faces = torch.from_numpy(mesh_tmp.faces[None]).to(device)
            mesh = Meshes(verts=verts, faces=faces)
            points, normals = sample_points_from_meshes(
                mesh, num_samples=cfg["data"]["num_points"], return_normals=True,
            )
            # mesh is saved in the original scale of the gt
            points -= center.float().to(device)
            points /= scale.float().to(device)
            points *= 0.9
            # make sure the points are within the range of [0, 1)
            points = points / 2.0 + 0.5
        else:
            # directly initialize from a point cloud
            pcd = o3d.io.read_point_cloud(cfg["train"]["input_mesh"])
            points = torch.from_numpy(np.array(pcd.points)[None]).float().to(device)
            normals = torch.from_numpy(np.array(pcd.normals)[None]).float().to(device)
            points -= center.float().to(device)
            points /= scale.float().to(device)
            points *= 0.9
            points = points / 2.0 + 0.5
    else:  #! initialize our source point cloud from a sphere
        sphere_radius = cfg["model"]["sphere_radius"]
        sphere_mesh = trimesh.creation.uv_sphere(radius=sphere_radius, count=[256, 256])
        points, idx = sphere_mesh.sample(cfg["data"]["num_points"], return_index=True)
        points += 0.5  # make sure the points are within the range of [0, 1)
        normals = sphere_mesh.face_normals[idx]
        points = torch.from_numpy(points).unsqueeze(0).to(device)
        normals = torch.from_numpy(normals).unsqueeze(0).to(device)

    points = torch.log(points / (1 - points))  # inverse sigmoid
    inputs = torch.cat([points, normals], axis=-1).float()
    inputs.requires_grad = True

    model = None  # no network

    # initialize optimizer
    cfg["train"]["schedule"]["pcl"]["initial"] = cfg["train"]["lr_pcl"]
    print("Initial learning rate:", cfg["train"]["schedule"]["pcl"]["initial"])
    if "schedule" in cfg["train"]:
        lr_schedules = get_learning_rate_schedules(cfg["train"]["schedule"])
    else:
        lr_schedules = None

    optimizer = update_optimizer(inputs, cfg, epoch=0, model=model, schedule=lr_schedules)

    try:
        # load model
        state_dict = torch.load(os.path.join(cfg["train"]["out_dir"], "model.pt"))
        if ("pcl" in state_dict.keys()) & (state_dict["pcl"] is not None):
            inputs = state_dict["pcl"].to(device)
            inputs.requires_grad = True

        optimizer = update_optimizer(inputs, cfg, epoch=state_dict.get("epoch"), schedule=lr_schedules)

        out = "Load model from epoch %d" % state_dict.get("epoch", 0)
        print(out)
        logger.info(out)
    except:
        state_dict = dict()

    start_epoch = state_dict.get("epoch", -1)

    trainer = Trainer(cfg, optimizer, device=device)
    runtime = {}
    runtime["all"] = AverageMeter()

    # training loop
    for epoch in range(start_epoch + 1, cfg["train"]["total_epochs"] + 1):
        # schedule the learning rate
        if (epoch > 0) & (lr_schedules is not None):
            if epoch % lr_schedules[0].interval == 0:
                adjust_learning_rate(lr_schedules, optimizer, epoch)
                if len(lr_schedules) > 1:
                    print(
                        "[epoch {}] net_lr: {}, pcl_lr: {}".format(
                            epoch, lr_schedules[0].get_learning_rate(epoch), lr_schedules[1].get_learning_rate(epoch),
                        ),
                    )
                else:
                    print(f"[epoch {epoch}] adjust pcl_lr to: {lr_schedules[0].get_learning_rate(epoch)}")

        start = time.time()
        loss, loss_each = trainer.train_step(data, inputs, model, epoch)
        runtime["all"].update(time.time() - start)

        if epoch % cfg["train"]["print_every"] == 0:
            log_text = ("[Epoch %02d] loss=%.5f") % (epoch, loss)
            if loss_each is not None:
                for k, l in loss_each.items():
                    if l.item() != 0.0:
                        log_text += f" loss_{k}={l.item():.5f}"

            log_text += (" time={:.3f} / {:.3f}").format(runtime["all"].val, runtime["all"].sum)
            logger.info(log_text)
            print(log_text)

        # visualize point clouds and meshes
        if (epoch % cfg["train"]["visualize_every"] == 0) & (vis is not None):
            trainer.visualize(data, inputs, model, epoch, o3d_vis=vis)

        # save outputs
        if epoch % cfg["train"]["save_every"] == 0:
            trainer.save_mesh_pointclouds(inputs, epoch, center.cpu().numpy(), scale.cpu().numpy() * (1 / 0.9))

        # save checkpoints
        if (epoch > 0) & (epoch % cfg["train"]["checkpoint_every"] == 0):
            state = {"epoch": epoch}
            if isinstance(inputs, torch.Tensor):
                state["pcl"] = inputs.detach().cpu()

            torch.save(state, os.path.join(cfg["train"]["dir_model"], "%04d" % epoch + ".pt"))
            print("Save new model at epoch %d" % epoch)
            logger.info("Save new model at epoch %d" % epoch)
            torch.save(state, os.path.join(cfg["train"]["out_dir"], "model.pt"))

        # resample and gradually add new points to the source pcl
        if (
            (epoch > 0)
            & (cfg["train"]["resample_every"] != 0)
            & (epoch % cfg["train"]["resample_every"] == 0)
            & (epoch < cfg["train"]["total_epochs"])
        ):
            inputs = trainer.point_resampling(inputs)
            optimizer = update_optimizer(inputs, cfg, epoch=epoch, model=model, schedule=lr_schedules)
            trainer = Trainer(cfg, optimizer, device=device)

    # visualize the Open3D outputs
    if cfg["train"]["o3d_show"]:
        out_video_dir = os.path.join(cfg["train"]["out_dir"], "vis/o3d/video.mp4")
        if os.path.isfile(out_video_dir):
            os.system(f"rm {out_video_dir}")
        os.system(
            "ffmpeg -framerate 30 \
                    -start_number 0 \
                    -i {}/vis/o3d/%04d.jpg \
                    -pix_fmt yuv420p \
                    -crf 17 {}".format(
                cfg["train"]["out_dir"], out_video_dir,
            ),
        )
        out_video_dir = os.path.join(cfg["train"]["out_dir"], "vis/o3d/video_pcd.mp4")
        if os.path.isfile(out_video_dir):
            os.system(f"rm {out_video_dir}")
        os.system(
            "ffmpeg -framerate 30 \
                    -start_number 0 \
                    -i {}/vis/o3d/%04d_pcd.jpg \
                    -pix_fmt yuv420p \
                    -crf 17 {}".format(
                cfg["train"]["out_dir"], out_video_dir,
            ),
        )
        print("Video saved.")


if __name__ == "__main__":
    main()