#!/usr/bin/env python3

'''
Copyright (c) 2016 Computational Biomechanics (CoBi) Core, Department of
Biomedical Engineering, Cleveland Clinic

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to permit
persons to whom the Software is furnished to do so, subject to the
following conditions:

The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
-------------------

Tekscan_contact_metrics.py

DESCRIPTION:

Python script to read Tekscan pressure csv file(s) and output contact metrics
(see OUTPUT below for metrics).

REQUIREMENTS:

Python 3 (http://www.python.org)
NumPy (http://www.numpy.org)

ORIGINALLY WRITTEN BY:

Diana Suciu
Summer Intern, 2015
Undergraduate student at Case Western Reserve University

REWRITTEN BY:

Craig Bennetts
Computational Biomodeling (CoBi) Core
Department of Biomedical Engineering
Lerner Research Institute
Cleveland Clinic
Cleveland, OH
bennetc2@ccf.org

INITIAL CONTIBUTION BY:

Ahmet Erdemir
Computational Biomodeling (CoBi) Core
Department of Biomedical Engineering
Lerner Research Institute
Cleveland Clinic
Cleveland, OH
erdemira@ccf.org

INPUT(S):

csv file(s) output from Tekscan binary data file, which include:
- number of sensels per row and column
- sensel area (mm^2)
- pressure data (MPa) for an arbitrary number of 'frames' (i.e. time points)

OUTPUT(S):

Text file(s) (name(s) based on input file(s)) containing the following contact metrics:
- input csv filename
- senspr dimensions (sensel rows x columns)
- contact force (N)
- contact area (mm^2)
- peak pressure (MPa)
- peak pressure location (sensel indices (row,col), indexed from 1)

ASSUMPTIONS:

- Input file contains Tekscan pressure data with an arbitrary
number of sampled time points.
- All time points were collected at a static loading condition on
the Tekscan.  Therefore, average all time points together will
remove noise and reduce potential variations (due to control, etc.).

'''

import sys
from numpy import array,zeros,sum,unravel_index,amax,arange
from numpy import set_printoptions,nan  # for repr
from numpy import spacing
try:
    import xml.etree.cElementTree as ET  # preferrably load cElementTree is much faster than ElementTree
except ImportError:
    import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap


def USAGE(argv):
	print()
	print('USAGE: ', argv[0], ' <csv file(s)>')
	print()


def Tekscan_contact_metrics(argv):

	# Check appropriate usage
	if len(argv)<2:
		USAGE(argv)
		sys.exit(1)

	# set options for repr if desired
	# threshold option allows full 2D array to be converted to a string without abbreviating with ellipses
	# formatter option enforces not to use exponential notation
	# WARNING: this assumes specific pressure sensor values only have 4 decimals of precision
	float_formatter = lambda x: "%0.4f" % x
	set_printoptions(threshold=sys.maxsize, formatter={'float_kind':float_formatter})

	# Set filenames
	input_filenames = argv[1:]

	# PROCESS INPUT TEKSCAN FILE(S)
	for input_filename in input_filenames:

		# CONVERT CSV TO XML:
		# 1) reads desired Tekscan csv parameters
		# 2) averages all temporal 'frames'
		# 3) computes contact metrics (contact area, contact force, peak pressure)
		# 4) builds XML tree, including: parameters, averaged frame, contact metrics
		# 5) outputs XML tree to file
		# RETURNS:
		tekscan_array, sensel_list, metrics_list = tekscan_csv_2_AVG_METRICS_XML(input_filename)


		# CREATE AVERAGE PRESSURE PLOT

		fig = plt.figure()

		#cmap = plt.get_cmap('spectral')
		#print(cmap._segmentdata)

		spectral_mod = modified_spectral_colormap()

		# OPTION 1: Use integer axes dimension ticks
		plt.imshow(tekscan_array, cmap=spectral_mod, interpolation='none', extent=[0.0,sensel_list[2]*float(tekscan_array.shape[1]),sensel_list[0]*float(tekscan_array.shape[0]),0.0])

#		# OPTION 2: Manually set uniformly divided axes dimension ticks
#		plt.imshow(tekscan_array, cmap=spectral_mod, interpolation='none')
#
#		NUM_XTICKS = 4
#		xmax = float(tekscan_array.shape[1]-1)  # NOTE: center-to center of sensels => -1!
#		dx = xmax/float(NUM_XTICKS)
#		xtick_values = arange(0.0,xmax+dx,dx)
#		xtick_labels = list( map(str, sensel_list[2]*xtick_values) ) # multiply sensel indices by sensel row dimension
#		print(xtick_values)
#		print(xtick_labels)
#		plt.xticks(xtick_values, xtick_labels)
#
#		NUM_YTICKS = 4
#		ymax = float(tekscan_array.shape[0]-1)
#		dy = ymax/float(NUM_YTICKS)
#		ytick_values = arange(0.0,ymax+dy,dy)
#		ytick_labels = list( map(str, sensel_list[0]*ytick_values) ) # multiply sensel indices by sensel row dimension
#		print(ytick_values)
#		print(ytick_labels)
#		plt.yticks(ytick_values, ytick_labels)

		# Set figure options
		xml_filename = input_filename[:-4] + '_AVG_METRICS.xml'
		title_string = xml_filename + '\nCONTACT: force = ' + '{:.3f}'.format(metrics_list[0]) + ' (' + metrics_list[1] + '), area = ' + '{:.3f}'.format(metrics_list[2]) + ' (' + metrics_list[3] + ')\nPEAK PRESSURE: ' + str(metrics_list[4]) + ' (' + str(metrics_list[5]) + ')'
		plt.title(title_string, fontsize=12)

		# show colorbar
		colorbar_label = 'Pressure (' + str(metrics_list[5]) + ')'
		plt.colorbar(label=colorbar_label)

		# label axes (default: 'mm')
		plt.xlabel(sensel_list[1])
		plt.ylabel(sensel_list[3])

		# pads are specified in fraction of fontsize
		plt.tight_layout(pad=0.75) # , w_pad=0.5, h_pad=1.0)
		image_filename = input_filename[:-4] + '_AVG_METRICS.png'
		plt.savefig(image_filename, bbox_inches='tight', dpi=fig.dpi)

		plt.close()
		
	# end of file loop

# END OF Tekscan_contact_metrics()


# Function to reformat XML ElementTree
# INPUTS:
# elem = ElementTree root => default level=0
# indent_string: e.g. '\t', '   ', etc.
# NOTE: modifies tree string in place
def indent_elements(elem, indent_string, level=0):
    i = "\n" + level*indent_string
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + indent_string
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent_elements(elem, indent_string, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

# END OF indent_elements()

# Function to redifine the colormap
# removes the upper white end of the colormap, stopping at pure red
def modified_spectral_colormap():
	
	cmap = plt.get_cmap('nipy_spectral')
	cmap_dict = cmap._segmentdata

	# unpack each color channel
	r = cmap_dict['red']
	g = cmap_dict['green']
	b = cmap_dict['blue']
	
	# remove upper white end of spectrum to stop at pure red
	r = r[:-3]
	g = g[:-3]
	b = b[:-3]

	# Normalize colormap range from 0.0 to 1.0
	# find max, e.g. 0.85
	norm_max = r[-1][0]

	# Normalize red colormap channel
	r_mod = []
	for p in r:
		r_mod.append( (p[0]/norm_max, p[1], p[2]) )

	# Normalize green colormap channel
	g_mod = []
	for p in g:
		g_mod.append( (p[0]/norm_max, p[1], p[2]) )

	# Normalize blue colormap channel
	b_mod = []
	for p in b:
		b_mod.append( (p[0]/norm_max, p[1], p[2]) )
	cmap_mod_dict = {}
	cmap_mod_dict['red'] = r_mod
	cmap_mod_dict['green'] = g_mod
	cmap_mod_dict['blue'] = b_mod

	cmap_mod = LinearSegmentedColormap('spectral_mod', cmap_mod_dict)
	
	return cmap_mod

# END OF modified_spectrum_colormap()


# FUNCTION TO:
# 1) read desired Tekscan csv parameters
# 2) average all temporal 'frames'
# 3) compute contact metrics (contact area, contact force, peak pressure)
# 4) build XML tree, including: parameters, averaged frame, contact metrics
# 5) output XML tree to file
# INPUT:
# Tekscan csv filename
# OUTPUT:
# Tekscan data and metrics in XML file
# RETURNS:
# - averaged Tekscan pressure array (MPa)
# - sensel dimensions [row_spacing, row_spacing_units, col_spacing, col_spacing_units]
# - desired metrics to add to average pressure sensor image [contact_force, calib_force_units, contact_area, area_units, peak_pressure, pressure_units]
def tekscan_csv_2_AVG_METRICS_XML(input_filename):

	# Input csv file
	file = open(input_filename,'r')
	lines = file.readlines()
	file.close()

	# INITIALIZE VALUES

	line_index = 0
	num_rows = 0  # number of sensel rows
	num_cols = 0  # number of sensel columns
	#sensel_area = 0.0  # area (mm^2) of each sensel
	#scale_factor = 0.0  # MPa / normalized value (i.e. raw)


	# EXTRACT TEKSCAN PARAMETERS

	calib_force = []
	calib_sum = []
	calib_cells = []

	# Find desired parameters in csv file
	#while (scale_factor==0.0 or sensel_area==0.0 or num_rows==0 or num_cols==0) and line_index<len(lines):
	while line_index<len(lines):

		line = lines[line_index]

		# NOTE: conditional order is the same order as in the csv file!
		# TODO: more efficient to organize in descending order of frequency in the csv file (i.e. put 'CALIBRATION_POINT_' first)

		# Set number of sensel rows, if found
		if line[:4]=='ROWS':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['ROWS', '44']
			num_rows = int(line_split[-1])

		# Set number of sensel columns, if found
		elif line[:4]=='COLS':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['COLS', '44']
			num_cols = int(line_split[-1])

		# Set sensel row dimension
		elif line[:11]=='ROW_SPACING':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['ROW_SPACING', '1.27', 'millimeters']
			row_spacing = float(line_split[1])
			row_spacing_units = line_split[-1]
			if row_spacing_units=='millimeters':
				row_spacing_units = 'mm'

		# Set sensel column dimension
		elif line[:11]=='COL_SPACING':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['COL_SPACING', '1.27', 'millimeters']
			col_spacing = float(line_split[1])
			col_spacing_units = line_split[-1]
			if col_spacing_units=='millimeters':
				col_spacing_units = 'mm'

		# Set sensel area, if found
		elif line[:11]=='SENSEL_AREA':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['SENSEL_AREA', '1.6129', 'mm2']
			sensel_area = float(line_split[1])
			area_units = line_split[-1]
			if area_units=='mm2':
				area_units = 'mm^2'

		# Set the noise threshold
		elif line[:15]=='NOISE_THRESHOLD':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['NOISE_THRESHOLD', '3']
			noise_threshold = int(line_split[-1])  # raw (type=integer)
		
		# Set scale factor, if found
		elif line[:12]=='SCALE_FACTOR':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['SCALE_FACTOR', '', '0.0011882', 'MPa', '/', 'raw']
			scale_factor = float(line_split[2])
			scale_units = ''.join(line_split[-3:])

		# Set exponent for conversion from raw to pressure data
		elif line[:8]=='EXPONENT':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['EXPONENT', '1.5754']
			exponent = float(line_split[-1])

		# Set exponent for conversion from raw to pressure data
		elif line[:19]=='SATURATION_PRESSURE':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['SATURATION_PRESSURE', '7.34773', 'MPa', '(Exponential', 'Extrapolation)']
			saturation_pressure = float(line_split[1])
			saturation_units = line_split[2]
			saturation_type = ' '.join(line_split[-2:])[1:-1]

		# Accumulate calibration data in lists
		# works for an arbitrary number of calibration points
		elif line[:18]=='CALIBRATION_POINT_':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['CALIBRATION_POINT_1', '49.9998', '(Newtons)', '3173', '(Raw', 'Sum)', '109', '(Number', 'of', 'Loaded', 'Cells)']
			calib_force.append( float(line_split[1]) )
			calib_force_units = line_split[2][1]  # grab 'N' from Newtons
			calib_sum.append( int(line_split[3]) )
			calib_sum_units = line_split[4][1:].lower()  # 'raw'
			calib_cells.append( int(line_split[6]) )
			calib_cells_tag = '_'.join(line_split[-2:])[:-1].lower()  # 'loaded_cells
			calib_cells_units = 'count'
			
		# Set pressure units for sensor frame data
		elif line[:5]=='UNITS':
			line_split = line.rstrip('\r\n').split(' ')  # result: e.g. ['UNITS', 'MPa']
			pressure_units = line_split[-1]
			if pressure_units=='KPa':
				pressure_units = 'kPa'  # use standardized unit notation for kPa

		# increment line index
		line_index += 1

	# end of tekscan parameter extraction loop

	#print(num_rows)
	#print(num_cols)
	#print(scale_factor)
	#print(sensel_area)


	# VERIFY TEKSCAN PARAMETERS

#	# Ensure scale factor was found
#	if scale_factor==0.0:
#		print('\n')
#		print('ERROR: Tekscan scale factor not found in csv file: %s\n' % filename)
#		print('\n')
#		sys.exit(1)

	# Ensure sensel dimensions were found

	# Initialize accumulation array
	if num_rows!=0 and num_cols!=0:
		accum_array = zeros(shape=(num_rows,num_cols))
	else:
		print('\n')
		print('ERROR: Tekscan sensel dimensions not found in csv file: %s\n' % input_filename)
		print('\n')
		sys.exit(1)


	# ACCUMULATE AND AVERAGE SENSEL PRESSURES FOR ALL FRAMES
	#
	# ASSUMPTION: frames come after Tekscan parameters in csv file!
	#             IF NOT: uncomment line below to reset line_index to zero here
	#                     (safer, but potentially unnecessary/slower/redundant)!
	line_index = 0

	frame_count = 0

	# Read to end of input csv file
	while line_index<len(lines):

		# get current line
		line = lines[line_index]

		# Accumulate current frame
		if line[:5]=='Frame':

			# Advance to next line
			# ASSUMPTION: frame data starts on line following the 'Frame' tag!
			line_index += 1
			# reset the sensel row counter
			row_count = 0

			# Read/accumulate current frame
			while row_count<num_rows:

				# Split csv line by comma
				row_strings = lines[line_index].strip('\n').split(',')
				# Convert list of strings to a 1D array(row) of floating values
				row_array = array(list(map(float,row_strings)))

				# Add current row to accumulation array
				# NOTE: could do sanity check to ensure the same number of columns!
				accum_array[row_count,:] += row_array

				# Increment line index and sensel row count
				line_index += 1
				row_count += 1
				
			# end of frame accumulating loop

			# Increment frame counter
			frame_count += 1

		# If frame not found, advance to next line
		else:
			line_index += 1
			
	# end of frame finding loop


	# Average the accumulated pressure sensor array data:
	# - divide by the number of frames, and
	# - multiply by the scaling factor  # WRONG! frame units are actually in MPa, not MPa/raw
	#avg_array = scale_factor*(accum_array/frame_count)
	tekscan_array = accum_array/frame_count


	# COMPUTE CONTACT METRICS

	# Set default pressure threshold for sensel inclusion in contact metrics
	THRESHOLD = 0.0  # MPa

	# Determine contact force (N)
	# sum(sensel_pressure*sensel_area) = sensel_area*sum(sensel_pressure)
	contact_force = sensel_area * sum( tekscan_array[ tekscan_array > THRESHOLD] )

	# Determine contact area (mm^2)
	contact_area = sensel_area * len( tekscan_array[ tekscan_array > THRESHOLD] )

	# Determine peak pressure (MPa) and sensel location indices (i,j)
	#peak_pressure = amax(tekscan_array)
	# NOTE: not sure what will happen if multiple max values exist in image?!
	peak_row, peak_col = unravel_index( tekscan_array.argmax(), tekscan_array.shape )
	peak_pressure = tekscan_array[peak_row,peak_col]
	print(peak_pressure)

	#-----#
	# XML #
	#-----#

	# Create XML tree
	root = ET.Element('Pressure_Sensor')

	# INPUT FILE

	# Define input file in XML tree
	element = ET.SubElement(root,'Input_File')
	element.set('type','csv')
	element.text = input_filename

	# SENSOR

	# Define Sensor parameters in XML tree
	sensor_element = ET.SubElement(root,'Dimensions')

	element = ET.SubElement(sensor_element,'rows')
	element.set('units','sensels')
	element.text = str(tekscan_array.shape[0])

	element = ET.SubElement(sensor_element,'cols')
	element.set('units','sensels')
	element.text = str(tekscan_array.shape[1])

	element = ET.SubElement(sensor_element,'sensel_area')
	element.set('units',area_units)
	element.text = str(sensel_area)

	# CALIBRATION

	# Define Calibration parameters in XML tree
	calibration_element = ET.SubElement(root,'Calibration')

	element = ET.SubElement(calibration_element,'noise_threshold')
	element.set('units','raw')
	element.text = str(noise_threshold)

	element = ET.SubElement(calibration_element,'scale_factor')
	element.set('units',scale_units)
	element.text = str(scale_factor)

	element = ET.SubElement(calibration_element,'exponent')
	element.text = str(exponent)

	element = ET.SubElement(calibration_element,'saturation_pressure')
	element.set('units',saturation_units)
	element.set('type',saturation_type)
	element.text = str(saturation_pressure)

	# Define Calibration Points in XML tree
	points_element = ET.SubElement(calibration_element,'Points')

	for i in range(len(calib_force)):
		
		point_element = ET.SubElement(points_element,'Point')
		point_element.set('num',str(i+1))
		
		element = ET.SubElement(point_element,'force')
		element.set('units',calib_force_units)
		element.text = str(calib_force[i])

		element = ET.SubElement(point_element,'sum')
		element.set('units',calib_sum_units)
		element.text = str(calib_sum[i])

		element = ET.SubElement(point_element,calib_cells_tag)
		element.set('units',calib_cells_units)
		element.text = str(calib_cells[i])

	# CONTACT METRICS

	# Define Contact Metrics element in XML tree
	metrics_element = ET.SubElement(root,'Contact_Metrics')

	# Define Contact Area in XML tree
	element = ET.SubElement(metrics_element,'area')
	element.set('units',area_units)
	element.text = str(contact_area)
	
	# Define Contact Force in XML tree
	element = ET.SubElement(metrics_element,'force')
	element.set('units',calib_force_units)
	element.text = str(contact_force)
	
	# Define Peak Pressure element in XML tree
	peak_element = ET.SubElement(metrics_element,'Peak_Pressure')
	
	# Define peak pressure element in XML tree
	element = ET.SubElement(peak_element,'value')
	element.set('units',pressure_units)
	element.text = str(peak_pressure)
	
	# Define Peak Pressure Location in XML tree
	location_element = ET.SubElement(peak_element,'Location')

	# define peak pressure sensel row
	element = ET.SubElement(location_element,'row')
	element.set('units','sensels')
	element.text = str(peak_row)

	# define peak pressure sensel column
	element = ET.SubElement(location_element,'col')
	element.set('units','sensels')
	element.text = str(peak_col)

	# AVERAGED TEKSCAN ARRAY

	# Build average Tekscan array string
	tekscan_string = repr(tekscan_array).replace(' ','').replace('\n','')

	element = ET.SubElement(root,'Pressure_Array')
	element.set('type','numpy array')
	element.set('processing','averaged')
	element.set('units',pressure_units)
	element.text = tekscan_string

	# OUTPUT XML

	XML_INDENT_STRING = "\t"
	
	indent_elements(root,XML_INDENT_STRING)
	
	output_filename = input_filename[:-4] + '_AVG_METRICS.xml'

	# Write XML tree to file
	tree = ET.ElementTree(root)
	tree.write(output_filename)

	# RETURN DESIRED INFO
	# build lists to return desired Tekscan parameters for image labels
	sensel_list = [row_spacing, row_spacing_units, col_spacing, col_spacing_units]
	metrics_list = [contact_force, calib_force_units, contact_area, area_units, peak_pressure, pressure_units]

	return (tekscan_array, sensel_list, metrics_list)

# END OF tekscan_csv_2_AVG_METRICS_XML()


# make this script callable from the command line
if __name__ == "__main__":
	Tekscan_contact_metrics(sys.argv)