import PhotoScan
import random
import csv

# This script enables a Monte Carlo bundle adjustment approach to be run in PhotoScan Pro, as described in 
# James et al (2016) 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.2.6
# The script parameters below should be edited, and it can then be run in PhotoScan via the Tools->Run script menu command.

########################################################################################
# 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:/Data/Monte Carlo GCP/'

# 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_accuracies = [0.005]
projection_accuracies = [1.0]
tiepoint_accuracies = [1.0]

# 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
optimise_f=False
optimise_cxcy=False
optimise_b1=False
optimise_b2=False
optimise_k1k2k3=False
optimise_p1p2=False
optimise_k4=False
optimise_p3=False
optimise_p4=False

########################################################################################
# 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  	

########################################################################################
# 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_accuracy in marker_accuracies:
		if marker_accuracy >= 0:					# Update the marker accuracy setting value if required
			try:
				chunk.accuracy_markers = PhotoScan.Vector( [marker_accuracy, marker_accuracy, marker_accuracy])			# Update the marker accuracy setting value
			except TypeError:
				chunk.accuracy_markers = marker_accuracy	# Single value for legacy versions		
				
		for projection_accuracy in projection_accuracies:
			chunk.accuracy_projections = projection_accuracy	# Update the projection accuracy setting value
			
			for tiepoint_accuracy in tiepoint_accuracies:
				chunk.accuracy_tiepoints = 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_accuracy) + '_PA' + '{0:0.5f}'.format(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_cxcy=optimise_cxcy, fit_b1=optimise_b1, fit_b2=optimise_b2, fit_k1k2k3=optimise_k1k2k3, fit_p1p2=optimise_p1p2, fit_k4=optimise_k4, fit_p3=optimise_p3,fit_p4=optimise_p4)

				# Export the ground control
				chunk.saveReference(dir_path + out_gc_file, 'csv', items='markers')

				# Export the cameras
				chunk.exportCameras(dir_path + out_cam_file, format='xml', projection=crs, rotation_order='xyz')
				
				# 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='ply')			
				
				# Increment the file counter
				file_idx = file_idx+1