import os
from argparse import ArgumentParser
import shutil

# This Python script is based on the shell converter script provided in the MipNerF 360 repository.
parser = ArgumentParser("Colmap converter")
parser.add_argument("--no_gpu", action='store_true')
parser.add_argument("--source_path", "-s", required=True, type=str)
parser.add_argument("--camera", default="OPENCV", type=str)
parser.add_argument("--colmap_executable", default="", type=str)
parser.add_argument("--resize", action="store_true")
parser.add_argument("--magick_executable", default="", type=str)
args = parser.parse_args()
colmap_command = '"{}"'.format(args.colmap_executable) if len(args.colmap_executable) > 0 else "colmap"
magick_command = '"{}"'.format(args.magick_executable) if len(args.magick_executable) > 0 else "magick"
use_gpu = 1 if not args.no_gpu else 0

os.makedirs(args.source_path + "/distorted/sparse", exist_ok=True)

## Feature extraction
os.system(colmap_command + " feature_extractor "\
    "--database_path " + args.source_path + "/distorted/database.db \
    --image_path " + args.source_path + "/input \
    --ImageReader.single_camera 1 \
    --ImageReader.camera_model " + args.camera + " \
    --SiftExtraction.use_gpu " + str(use_gpu))

## Feature matching
os.system(colmap_command + " exhaustive_matcher \
    --database_path " + args.source_path + "/distorted/database.db \
    --SiftMatching.use_gpu " + str(use_gpu))

### Bundle adjustment
# The default Mapper tolerance is unnecessarily large,
# decreasing it speeds up bundle adjustment steps.
os.system(colmap_command + " mapper \
    --database_path " + args.source_path + "/distorted/database.db \
    --image_path "  + args.source_path + "/input \
    --output_path "  + args.source_path + "/distorted/sparse \
    --Mapper.ba_global_function_tolerance=0.000001")

### Image undistortion
## We need to undistort our images into ideal pinhole intrinsics.
os.system(colmap_command + " image_undistorter \
    --image_path " + args.source_path + "/input \
    --input_path " + args.source_path + "/distorted/sparse/0 \
    --output_path " + args.source_path + "\
    --output_type COLMAP")

files = os.listdir(args.source_path + "/sparse")
os.makedirs(args.source_path + "/sparse/0", exist_ok=True)
# Copy each file from the source directory to the destination directory
for file in files:
    if file == '0':
        continue
    source_file = os.path.join(args.source_path, "sparse", file)
    destination_file = os.path.join(args.source_path, "sparse", "0", file)
    shutil.move(source_file, destination_file)

if(args.resize):
    print("Copying and resizing...")

    # Resize images.
    os.makedirs(args.source_path + "/images_2", exist_ok=True)
    os.makedirs(args.source_path + "/images_4", exist_ok=True)
    os.makedirs(args.source_path + "/images_8", exist_ok=True)
    # Get the list of files in the source directory
    files = os.listdir(args.source_path + "/images")
    # Copy each file from the source directory to the destination directory
    for file in files:
        source_file = os.path.join(args.source_path, "images", file)

        destination_file = os.path.join(args.source_path, "images_2", file)
        shutil.copy2(source_file, destination_file)
        os.system(magick_command + " mogrify -resize 50% " + destination_file)

        destination_file = os.path.join(args.source_path, "images_4", file)
        shutil.copy2(source_file, destination_file)
        os.system(magick_command + " mogrify -resize 25% " + destination_file)

        destination_file = os.path.join(args.source_path, "images_8", file)
        shutil.copy2(source_file, destination_file)
        os.system(magick_command + " mogrify -resize 12.5% " + destination_file)

print("Done.")