#!/usr/bin/env python

# Copyright (c) 2021, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration.
#
# All rights reserved.
#
# The "ISAAC - Integrated System for Autonomous and Adaptive Caretaking
# platform" software is licensed under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

"""
A wrapper around the tools that when run together produce a textured mesh.
"""
import argparse, os, re, shutil, subprocess, sys, glob
import numpy as np

# Set up the path to Python modules about to load
basepath    = os.path.abspath(sys.path[0])
pythonpath  = os.path.abspath(basepath + '/../Python')  # for dev ASP
libexecpath = os.path.abspath(basepath + '/../libexec') # for packaged ASP
sys.path.insert(0, basepath) # prepend to Python path
sys.path.insert(0, pythonpath)
sys.path.insert(0, libexecpath)

import asp_rig_utils, asp_system_utils

def process_args(args):
    """
    Set up the parser and parse the args.
    """

    parser = argparse.ArgumentParser(description = "Parameters for texrecon.")
    
    parser.add_argument("--rig_config",  dest="rig_config", default="",
                        help = "Rig configuration file.")
    
    parser.add_argument("--rig_sensor", dest="rig_sensor", default="",
                        help="Which rig sensor images to texture. Must be among the " + \
                        "sensors specified via --rig_config. To use images from " +
                        "several sensors, pass in a quoted list of them, " + \
                        "separated by a space.")

    parser.add_argument("--camera_poses", dest="camera_poses",
                        default="", help= "Read images and camera poses from this list.")

    parser.add_argument("--subset", dest="subset",
                        default="", help= "Use only the subset of images from this list.")
    
    parser.add_argument("--mesh", dest="mesh",
                        default="", help="The mesh to use for texturing, in .ply format.")

    parser.add_argument("--max_texture_size", dest="max_texture_size",
                        default="2048",
                        help="The maximum size (in pixels) of each texture file " + \
                        "created for the produced textured mesh.")

    parser.add_argument("--texture_alg", dest="texture_alg",
                        default="center", help= "Use one of the two texture creation modes: 'center' (for a surface patch choose the image in which the patch shows up closer to the image center), 'area' (for a surface patch choose the image whose camera view direction is most aligned with the surface normal).")
    parser.add_argument('--skip_local_seam_leveling', dest = "skip_local_seam_leveling",
                        action='store_true')
    
    parser.add_argument("--out_dir", dest="out_dir",
                        default="", help="The directory where to write the textured mesh " + \
                        "and other data.")
    
    # Note how the percent sign below is escaped, by writing: %%
    parser.add_argument("--undistorted_crop_win", dest="undistorted_crop_win", default = "",
                        help = "The dimensions of the central image region to keep after "   + \
                        "undistorting an image and before using it in texturing. Normally "  + \
                        "85%% - 90%% of distorted (actual) image dimensions would do. "      + \
                        "Suggested the Astrobee images: "        + \
                        "sci_cam: '1250 1000' nav_cam: '1100 776'. haz_cam: '250 200'.")
    
    args = parser.parse_args()

    return args

def sanity_checks(args):

    if args.camera_poses == "":
        raise Exception("The path to the list having input images and poses was not specified.")

    if args.mesh == "":
        raise Exception("The mesh to use for texturing was not specified.")

    if args.rig_config == "":
        raise Exception("The path to the rig configuration was not specified.")

    if args.rig_sensor == "":
        raise Exception("The rig sensor to use for texturing was not specified.")

    if args.out_dir == "":
        raise Exception("The path to the output directory was not specified.")

    if args.undistorted_crop_win == "":
        raise Exception("The undistorted crop win was not specified.")

def convert_intrinsics_to_texrecon(undist_intrinsics_file):

    (widx, widy, f, cx, cy) = asp_rig_utils.read_intrinsics(undist_intrinsics_file)

    max_wid = widx
    if widy > max_wid:
        max_wid = widy

    # normalize
    nf = f / max_wid
    ncx = cx / widx
    ncy = cy / widy
    d0 = 0.0
    d1 = 0.0
    paspect = 1.0

    return (nf, d0, d1, paspect, ncx, ncy)

def create_texrecon_cameras(undistorted_images, world_to_cam, nf, d0, d1, paspect, ncx, ncy):
    if len(undistorted_images) != len(world_to_cam):
        raise Exception("Expecting as many images as cameras.")

    cams = []
    for it in range(len(undistorted_images)):
        path, ext = os.path.splitext(undistorted_images[it])
        cam = path + ".cam"
        cams.append(cam)

        print("Writing: " + cam)
        with open(cam, "w") as g:
            M = world_to_cam[it]
            # write translation
            g.write("%0.17g %0.17g %0.17g " % (M[0][3], M[1][3], M[2][3]))

            # write rotation
            g.write(
                "%0.17g %0.17g %0.17g %0.17g %0.17g %0.17g %0.17g %0.17g %0.17g\n"
                % (M[0][0], M[0][1], M[0][2], M[1][0], M[1][1], M[1][2],
                   M[2][0], M[2][1], M[2][2]))

            # normalized inrinsics
            g.write("%0.17g %0.17g %0.17g %0.17g %0.17g %0.17g\n"
                % (nf, d0, d1, paspect, ncx, ncy))
    return cams
    
def run_texrecon(tools_base_dir, undistorted_images, cams, mesh,
                 skip_local_seam_leveling, texture_dir):

    # That is one long path
    texrecon_path = tools_base_dir + "/bin/texrecon_bin"
    if not os.path.exists(texrecon_path):
        # Handle the case when the dev version of ASP is used
        print("Warning: Cannot find texrecon at: " + texrecon_path)
        texrecon_path = asp_system_utils.which('texrecon_bin')
        if texrecon_path is not None and os.path.exists(texrecon_path):
            print("Will use: " + texrecon_path)
        else:
            raise Exception("Cannot find: " + texrecon_path)
    
    asp_rig_utils.mkdir_p(texture_dir)

    # Form the list of images to pass in
    index = texture_dir + "/index.txt"
    if len(undistorted_images) != len(cams):
        print("Must have as many images as cameras.")
        sys.exit(1)
    print("Writing: " + index)
    with open(index, "w") as fh:
        for it in range(len(undistorted_images)):
            # texrecon likes the cameras first
            fh.write(cams[it] + "\n")
            fh.write(undistorted_images[it] + "\n")
    
    cmd = [texrecon_path, index, mesh, texture_dir,
           "-o", "gauss_clamping",
           "-d", args.texture_alg,
           "--keep_unseen_faces",
           # TODO(oalexan1): Need to understand why texrecon insists on an equal sign
           # for the option below while it does not need that for other options.
           "--max_texture_size=" + str(args.max_texture_size)]

    if skip_local_seam_leveling:
        cmd += ['--skip_local_seam_leveling']

    log_file = os.path.join(texture_dir, "texrecon_log.txt")
    print("Running texrecon.\n")
    asp_rig_utils.run_cmd(cmd)

    textured_mesh = texture_dir + ".obj"

    print("\nWrote: " + textured_mesh)
    
    return textured_mesh

if __name__ == "__main__":

    args = process_args(sys.argv)
    tools_base_dir = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))

    sanity_checks(args)

    asp_rig_utils.mkdir_p(args.out_dir)
    all_undistorted_images = []
    all_cams   = []

    # Can handle multiple sensors
    sensors = args.rig_sensor.split()
    for sensor in sensors:
        (images, world_to_cam) \
                 = asp_rig_utils.parse_cameras(args.camera_poses, args.subset,
                                                           sensor)
        
        extension = '.jpg' # extension of undistorted images, as preferred by texrecon
        extra_opts = ["--save_bgr"] # create color images
        suff = "_texrecon"
        (undist_intrinsics_file, undistorted_images, undist_dir) \
                                 = asp_rig_utils.undistort_images(args, sensor,
                                                              images, tools_base_dir,
                                                              extension, extra_opts, suff)
        (nf, d0, d1, paspect, ncx, ncy) = convert_intrinsics_to_texrecon(undist_intrinsics_file)
        cams = create_texrecon_cameras(undistorted_images, world_to_cam,
                                       nf, d0, d1, paspect, ncx, ncy)

        all_undistorted_images += undistorted_images
        all_cams += cams

    texture_dir = args.out_dir + "/" + "_".join(sensors) + "/texture"
    run_texrecon(tools_base_dir, all_undistorted_images, all_cams, args.mesh,
                 args.skip_local_seam_leveling, texture_dir)

