import PhotoScan
import random
import math
import csv
NaN = float('NaN')

# For use with PhotoScan Pro v.1.3, with projects saved as .psz archives.
#
# Python script associated with James et al (2017):
# Optimising UAV topographic surveys processed with structure-from-motion: Ground control quality, 
# quantity and bundle adjustment, Geomorphology, doi: 10.1016/j.geomorph.2016.11.021
# This updated version of the script has been tested on PhotoScan v.1.3.0
# The script parameters below should be edited, and it can then be run in PhotoScan via the Tools->Run script menu command.
# 
# This script uses a Monte Carlo approach to assess GCP performance in SfM surveys.
# Output files from the script can be processed using sfm_georef (http://tinyurl.com/sfmgeoref) 
# Designed for use on projects containing a single chunk, with all photos 
# taken with the same camera.

# Author: Mike James, Lancaster University, U.K.
# Contact: m.james at lancaster.ac.uk
# Updates: Check http://tinyurl.com/sfmgeoref
# 
# Tested and used in PhotoScan Pro v.1.3, with projects saved as .psz archives.
# 11/03/17 Updated the camera parameter optimisation options to exploit the greater flexibility now offered.
# 11/03/17 Multiple changes to export function parameters to accommodate PhotoScan updates.

########################################################################################
######################################   SETUP    ######################################
########################################################################################
# Edit parameters up to the Initialisation section
# Directory where output will be stored and active control file is saved.
# The directory should exist before the script is run.
# Note use of '/' in the path (not '\'); end the path with '/'
dir_path = 'E:/SfM/Monte_Carlo/'

# Output text file which describes the active markers used as control.
# Use 'active_ctrl_indices.txt' for compatibility with sfm_georef for analysis of results.
# Each line in the file allocates the markers as active (1) or inactive (0).
# Thus, the number of elements in a line equals the number of markers in the PhotoScan project.
# The number of lines in the file gives the number of times each unique combination of the accuracy
# settings will be used.
# [RECOMMENDED] - don't change the name of act_ctrl_file
act_ctrl_file = 'active_ctrl_indices.txt'

# Accuracy settings values to use.
# For Stage II analyses, use a range for marker_accuracies:
# Note that any marker-specific accuracies in the PhotoScan project will be overwritten, and X, Y and
# Z components of marker accuracy will be identical.
# e.g. marker_accuracies = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.10, 0.20, 0.50, 1.00, 2.00]
# Projection and tie point accuracies should be set to reflect the RMS residual values on markers and tie points respectively.
marker_location_accuracies = [0.02]
marker_projection_accuracies = [0.03]
tiepoint_accuracies = [0.74]

# Define the percent of markers that should be active as control (all = 100)
# A list of one or more values:
# e.g. perc_act_markers = [10, 20, 30, 40, 50, 60, 70, 80, 90]
perc_act_markers = [50]

# Define the number of random selections of markers to use (i.e. how many times bundle adjustment)
# will be carried out for each different value of perc_act_markers.
num_randomisations = 50

# Define the camera parameter set to optimise in the bundle adjustment.
# v.1.3 of PhotoScan enables individual selection/deselection of all parameters.
# Note - b1 was previously 'aspect ratio' (i.e. the difference between fx and fy)
#        b2 was previously 'skew'
optimise_f=True
optimise_cx=True
optimise_cy=True
optimise_b1=False
optimise_b2=False
optimise_k1=True
optimise_k2=True
optimise_k3=True
optimise_k4=False
optimise_p1=False
optimise_p2=False
optimise_p3=False
optimise_p4=False
optimise_shutter=False


# Points are exported as floats in binary ply files for speed and size, and thus cannot represent very small changes
# in large geographic coordinates. Thus, the offset below can be set to form a local origin and ensure numerical
# precision is not lost when coordinates are saved as floats. The offset will be subtracted from point coordinates.
# [RECOMMENDED] - Leave as NaN; the script will automatically calculate and apply a suitable offset, which will be saved
# as a text file for further processing, OR edit the line to impose a specific offset of your choice - 
# e.g.  pts_offset = PhotoScan.Vector( [266000, 4702000, 0] )
pts_offset = PhotoScan.Vector( [NaN, NaN, NaN] )

########################################################################################
# Initialisation
chunk = PhotoScan.app.document.chunk

# Reset the random seed, so that all equivalent runs are started identically
random.seed(1)

# Determine the total number of markers and the number of active markers 
num_markers = len(chunk.markers)
num_act_markers = [round(num_markers*x/100) for x in perc_act_markers]

# Construct a list of lists representing the randomised selection of active markers
act_markers = []
for repID in range(0, len(perc_act_markers)):
	nam = num_act_markers[repID]
	for rand_num in range(0, num_randomisations):
		act_marker_line = [True] * nam + [False] * (num_markers-nam)
		random.shuffle( act_marker_line )
		act_markers.append( act_marker_line )

# Write the active marker flags to a text file - one line per BA iteration		
with open(dir_path + act_ctrl_file, 'w') as f:
	fwriter = csv.writer(f, delimiter=' ', lineterminator='\n')
	for line_ID in range(0, len(act_markers)):
		fwriter.writerow( [int(elem) for elem in act_markers[line_ID]] ) 
	f.close()
	
# Need CoordinateSystem object for camera export, but PS only returns 'None' if an arbitrary coordinate system is being used
# thus need to set manually in this case; otherwise use the Chunk coordinate system.
if chunk.crs == None:
	crs = PhotoScan.CoordinateSystem('LOCAL_CS["Local CS",LOCAL_DATUM["Local Datum",0],UNIT["metre",1]]')			
else:
	crs = chunk.crs

# Counter for the number of bundle adjustments carried out, to prepend to files
file_idx = 1  	

# If required, calculate the mean point coordinate to use as an offset
if math.isnan(pts_offset[0]):
	points = chunk.point_cloud.points
	npoints = 0
	pts_offset = PhotoScan.Vector( [0, 0, 0] )
	for point in points:
		if not point.valid:
			continue
		npoints += 1
		pts_offset[0] += point.coord[0]
		pts_offset[1] += point.coord[1]
		pts_offset[2] += point.coord[2]
	
	pts_offset = crs.project(chunk.transform.matrix.mulp(pts_offset/npoints))
	pts_offset[0] = round(pts_offset[0], -2)
	pts_offset[1] = round(pts_offset[1], -2)
	pts_offset[2] = round(pts_offset[2], -2)
	
# Save the used offset to text file
with open(dir_path + '_coordinate_local_origin.txt', "w") as f:
	fwriter = csv.writer(f, dialect='excel-tab', lineterminator='\n')
	fwriter.writerow( pts_offset )
	f.close()


########################################################################################
# Main set of nested loops which control the repeated bundle adjustment

# Repeat all inner loops of accuracy settings values for each permutation of control points ('active markers')
for line_ID in range(0, len(act_markers) ):

	# Update which markers are enabled for use as control points in the bundle adjustment
	act_marker_flags = act_markers[line_ID]
	num_act_markers = sum(act_marker_flags)
	for i, marker in enumerate(chunk.markers):
		marker.reference.enabled = act_marker_flags[i]
						
	# For this arrangement of control points, loop through all settings permutations
	for marker_location_accuracy in marker_location_accuracies:
		if marker_location_accuracy >= 0:					# Update the marker accuracy setting value if required
			try:
				for i, marker in enumerate(chunk.markers):
					marker.reference.accuracy = PhotoScan.Vector( [marker_location_accuracy, marker_location_accuracy, marker_location_accuracy])			# Update the marker accuracy setting value
			except TypeError:
				chunk.marker_location_accuracy = marker_location_accuracy	# Single value for legacy versions		
				
		for marker_projection_accuracy in marker_projection_accuracies:
			chunk.marker_projection_accuracy = marker_projection_accuracy	# Update the projection accuracy setting value
			
			for tiepoint_accuracy in tiepoint_accuracies:
				chunk.tiepoint_accuracy = tiepoint_accuracy	# Update the projection accuracy setting value
			
				# Construct the output file names
				out_file = '{0:04d}'.format(file_idx) + '_MA' + '{0:0.5f}'.format(marker_location_accuracy) + '_PA' + '{0:0.5f}'.format(marker_projection_accuracy) + '_TA' + '{0:0.5f}'.format(tiepoint_accuracy)+ '_NAM' + '{0:03d}'.format(num_act_markers) + '_LID' + '{0:03d}'.format(line_ID+1)
				out_gc_file = out_file + '_GC.txt'
				out_cam_file = out_file + '_cams.xml'
				print( out_gc_file )
				
				# Bundle adjustment
				chunk.optimizeCameras(fit_f=optimise_f, fit_cx=optimise_cx, fit_cy=optimise_cy, fit_b1=optimise_b1, fit_b2=optimise_b2, fit_k1=optimise_k1, fit_k2=optimise_k2, fit_k3=optimise_k3, fit_k4=optimise_k4, fit_p1=optimise_p1, fit_p2=optimise_p2, fit_p3=optimise_p3, fit_p4=optimise_p4, fit_shutter=optimise_shutter)

				# Export the ground control
				chunk.saveReference(dir_path + out_gc_file, PhotoScan.ReferenceFormatCSV, items=PhotoScan.ReferenceItemsMarkers,)

				# Export the cameras
				chunk.exportCameras(dir_path + out_cam_file, format=PhotoScan.CamerasFormatXML, projection=crs, rotation_order=PhotoScan.RotationOrderXYZ)
				
				# Export the calibrations [NOTE - only one camera implemented in export here]
				chunk.cameras[0].sensor.calibration.save(dir_path + out_file + '_cal1.xml')
				
				# Export the sparse point cloud
				chunk.exportPoints(dir_path + out_file + '_pts.ply', normals=False, colors=False, format=PhotoScan.PointsFormatPLY, projection=crs, shift=pts_offset)			
				
				# Increment the file counter
				file_idx = file_idx+1