import os
import subprocess
from extract_sfm import extract_all_to_dir
from normalize_cam_dict import normalize_cam_dict

#########################################################################
# Note: configure the colmap_bin to the colmap executable on your machine
#########################################################################


def bash_run(cmd):
    colmap_bin = "/usr/local/bin/colmap"
    cmd = colmap_bin + " " + cmd
    print("\nRunning cmd: ", cmd)

    subprocess.check_call(["/bin/bash", "-c", cmd])


gpu_index = "-1"


def run_sift_matching(img_dir, db_file, remove_exist=False):
    print("Running sift matching...")

    if remove_exist and os.path.exists(db_file):
        os.remove(db_file)  # otherwise colmap will skip sift matching

    # feature extraction
    # if there's no attached display, cannot use feature extractor with GPU
    cmd = f" feature_extractor --database_path {db_file}"\
          f" --image_path {img_dir}"\
          f" --ImageReader.single_camera 1"\
          f" --ImageReader.camera_model SIMPLE_RADIAL"\
          f" --SiftExtraction.max_image_size 5000 "\
          f" --SiftExtraction.estimate_affine_shape 0"\
          f" --SiftExtraction.domain_size_pooling 1"\
          f" --SiftExtraction.use_gpu 1"\
          f" --SiftExtraction.max_num_features 16384"\
          f" --SiftExtraction.gpu_index {gpu_index}"
    bash_run(cmd)

    # feature matching
    cmd = f" exhaustive_matcher --database_path {db_file}"\
          f" --SiftMatching.guided_matching 1"\
          f" --SiftMatching.use_gpu 1"\
          f" --SiftMatching.max_num_matches 65536"\
          f" --SiftMatching.max_error 3"\
          f" --SiftMatching.gpu_index {gpu_index}"
    bash_run(cmd)


def run_sfm(img_dir, db_file, out_dir):
    print("Running SfM...")

    cmd = f" mapper"\
          f" --database_path {db_file}"\
          f" --image_path {img_dir}"\
          f" --output_path {out_dir}"\
          f" --Mapper.tri_min_angle 3.0"\
          f" --Mapper.filter_min_tri_angle 3.0"
    bash_run(cmd)


def prepare_mvs(img_dir, sparse_dir, mvs_dir):
    print("Preparing for MVS...")

    cmd = f" image_undistorter"\
          f" --image_path {img_dir}"\
          f" --input_path {sparse_dir}"\
          f" --output_path {mvs_dir}"\
          f" --output_type COLMAP"\
          f" --max_image_size 2000"
    bash_run(cmd)


def run_photometric_mvs(mvs_dir, window_radius):
    print("Running photometric MVS...")

    cmd = f" patch_match_stereo --workspace_path {mvs_dir}"\
          f" --PatchMatchStereo.window_radius {window_radius}"\
          f" --PatchMatchStereo.min_triangulation_angle 3.0"\
          f" --PatchMatchStereo.filter 1"\
          f" --PatchMatchStereo.geom_consistency 1"\
          f" --PatchMatchStereo.gpu_index={gpu_index}"\
          f" --PatchMatchStereo.num_samples 15"\
          f" --PatchMatchStereo.num_iterations 12"
    bash_run(cmd)


def run_fuse(mvs_dir, out_ply):

    print("Running depth fusion...")

    cmd = f"stereo_fusion"\
          f" --workspace_path {mvs_dir}"\
          f" --output_path {out_ply}"\
          f" --input_type geometric"
    bash_run(cmd)


def run_possion_mesher(in_ply, out_ply, trim):
    print("Running possion mesher...")

    cmd = f" poisson_mesher"\
          f" --input_path {in_ply}"\
          f" --output_path {out_ply}"\
          f" --PoissonMeshing.trim {trim}"
    bash_run(cmd)


def main(img_dir, out_dir, run_mvs=False):
    os.makedirs(out_dir, exist_ok=True)

    #### run sfm
    sfm_dir = os.path.join(out_dir, "sfm")
    os.makedirs(sfm_dir, exist_ok=True)

    img_dir_link = os.path.join(sfm_dir, "images")
    if os.path.exists(img_dir_link):
        os.remove(img_dir_link)
    os.symlink(img_dir, img_dir_link)

    db_file = os.path.join(sfm_dir, "database.db")
    run_sift_matching(img_dir, db_file, remove_exist=False)
    sparse_dir = os.path.join(sfm_dir, "sparse/0")
    os.makedirs(sparse_dir, exist_ok=True)
    run_sfm(img_dir, db_file, sparse_dir)

    # undistort images
    mvs_dir = os.path.join(out_dir, "mvs")
    os.makedirs(mvs_dir, exist_ok=True)
    prepare_mvs(img_dir, sparse_dir, mvs_dir)

    # extract camera parameters and undistorted images
    os.makedirs(os.path.join(out_dir, "posed_images"), exist_ok=True)
    extract_all_to_dir(
        os.path.join(mvs_dir, "sparse"), os.path.join(out_dir, "posed_images")
    )
    undistorted_img_dir = os.path.join(mvs_dir, "images")
    posed_img_dir_link = os.path.join(out_dir, "posed_images/images")
    if os.path.exists(posed_img_dir_link):
        os.remove(posed_img_dir_link)
    os.symlink(undistorted_img_dir, posed_img_dir_link)
    # normalize average camera center to origin, and put all cameras inside the unit sphere
    normalize_cam_dict(
        os.path.join(out_dir, "posed_images/kai_cameras.json"),
        os.path.join(out_dir, "posed_images/kai_cameras_normalized.json"),
    )

    if run_mvs:
        # run mvs
        run_photometric_mvs(mvs_dir, window_radius=7)

        out_ply = os.path.join(out_dir, "mvs/fused.ply")
        run_fuse(mvs_dir, out_ply)

        out_mesh_ply = os.path.join(out_dir, "mvs/meshed_trim_3.ply")
        run_possion_mesher(out_ply, out_mesh_ply, trim=3)


if __name__ == "__main__":
    ### note: this script is intended for the case where all images are taken by the same camera, i.e., intrinisics are shared.

    # CHANGE THIS TWO PATHS
    img_dir = "~/code/my_awesome_dataset/"
    out_dir = "~/home/code/nerfplusplus/outdir/"
    run_mvs = False
    main(img_dir, out_dir, run_mvs=run_mvs)