#!/usr/bin/env python
# Katherine Lake, UW-Madison, klake2@wisc.edu
# Last modified: 7/8/2013
# Code to run linear regression on C output files for simple plant model
# Requires scipy, numpy, matplotlib to run. No GUI

from scipy import polyfit
import numpy as np
import matplotlib.pyplot as plt
from tkFileDialog import askopenfilenames
import Tkinter
import sys
from subprocess import Popen, PIPE
import time

# Calls McKay C code using Popen to open a command line pipe
def call_mckay_code(timeString):

    # Get a list of file names
    files = askopenfilenames()
    if not files:
	print '---> Please select at least one file\n'
	return

    # In some OSes askopenfilenames() will return a string
    # This converts string of filenames to a tuple of filenames
    if type(files) is not tuple:
    	files = string2list(files)

    i = 0
    
    # Loop over the files
    while (i < len(files)):
        inputFile = files[i]

	# Create the string that will become the name of the McKay model output file
	# Example: inputfile.txt_out_YYYYMMDD_HH:MM.txt
	outputFile = inputFile[:-4] + "_out_" + timeString + ".txt"

	# Feed correct command to command shell using Popen
	# Command: ./McKay_akc_dev inputfile.txt outputfile.txt
	cmd="./Plant_main_akc " + inputFile + " " + outputFile
	sub = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE)
	
	# Get command line output from the previous command
	subout = sub.communicate()
	
	# Print the command line output so it is visible from the Python code
	# Input is returned as tuple, print the first item in tuple
	print subout[0]

        i = i + 1

def clean_input_file():
  
    # Get a list of file names
    files = askopenfilenames()
    if not files:
	print '---> Please select at least one file\n'
	return

    # In some OSes askopenfilenames() will return a string
    # This converts string of filenames to a tuple of filenames
    if type(files) is not tuple:
    	files = string2list(files)

    i = 0
    
    # Loop over the files
    while (i < len(files)):
        inputFile = files[i]

	clean = open(inputFile).read().replace('\"', '')
	newFileString = inputFile[:-4] + "_cleaned.csv"
	newFile = open(newFileString, 'w')
	newFile.write(clean)
	newFile.close()

	i = i + 1

# Reads the header portion of the LinReg and Graphing input file (McKay output file)
def read_ROI_summary(filename):

    i = 0

    # Open file
    fo = open(filename, 'r')

    # Start reading file
    str = fo.readline()

    # Read until getting to ROI name, MAT, Xi, sv
    while (not "cor" in str):

        str = fo.readline()
        
        # Record scan subject and MAT
        if "subject" in str:
            str = fo.readline()
            subject = str.strip()

        if "calculated at" in str:
            str = fo.readline()
            MAT_bound = int(str.strip())

    j = 0
    ROI = []
    MAT = []
    Xi = []    
    sv = []

    # Read MAT and corresponding Xi for each ROI
    # Based on ROI name being 'corXX'; if ROI naming convention changes, change string matching
    while ("cor" in str):
        split = str.split()
        ROI.append(split[0])
        Xi.append(float(split[1]))
        sv.append(float(split[2]))
        MAT.append(float(split[3]))
        str = fo.readline()
        j = j + 1

    # Close file
    fo.close()

    return (ROI, MAT, Xi, subject, MAT_bound, sv)

# Plot both MAT vs Xi and linear regression on same graph
# Automatically save graphs to location of python code
def plot_linreg_data(arrival_data, slope, intercept, equation, fileName):

    xVal = []
    yVal = []
    # Minimum x value is 10 s before minimum MAT
    x = int(min(arrival_data[1])) - 10

    # Calculate linear regression for 10 s before and after MAT values
    for num in range((int(min(arrival_data[1]))-10), (int(max(arrival_data[1]))+10)):
        xVal.append(x)
        yVal.append((x*slope) + intercept)
        x = x + 1

    # Create a new figure for each plot
    plt.figure()

    # Plot the Xi vs MAT and the linear regression
    plt.plot(arrival_data[1], arrival_data[2], '.')
    plt.plot(xVal, yVal)
    # Add a legend (upper right), title, x and y axis labels
    plt.legend(('Xi vs MAT', equation), loc=2)
    plt.title(arrival_data[3])
    plt.xlabel('Mean Arrival Time (s)')
    plt.ylabel('Position (mm)')

    # Save the figure, do not display it
    save_name = fileName[:-4] + '.png'
    plt.savefig(save_name)

    # Display figure
    # plt.show()

# Read the time, observed, bound, free, bound+free (M) modeled curves
# Return a tuple of (time, bound, observed, M)
def read_model_curves(filename, i):

    # Create arrays to store model data
    time = []
    observed = []
    trapped = []
    free = []
    M = []

    # Retrieve a list of ROI names and the scan subject
    arrival_data = read_ROI_summary(filename)
    ROIs = arrival_data[0]

    # Open file
    fo = open(filename)

    str = fo.readline()

    # Read the header, up until the first occurrence of the ROI name
    # Skip this first occurrence because it lists the MAT, Xi, and sv
    while (not ROIs[i] in str):
        str = fo.readline()
        if ("nROIs" in str):
            nROIs = int(fo.readline())
            
    str = fo.readline()

    # Find the next occurrence of the ROI name
    # This occurrence is where the actual model data is stored
    while (not ROIs[i] in str):
        str = fo.readline()

    # At this point str = the name of the ROI whose data we want
    # The next line is the ROI data header. Throw away that line
    fo.readline()
    
    # Read the first line of model data
    str = fo.readline()

    # A blank line separates each set of ROI data. Read in model data until we reach a blank line. Store the data in arrays
    while (str.strip() != ''):
        split = str.split()
        time.append(split[0])
        observed.append(split[1])
        trapped.append(split[2])
        free.append(split[3])
        M.append(split[4])
        str = fo.readline()

    fo.close()

    # Return the arrays of model data, the subject, the number of ROIs, and the name of the ROI whose data we have
    return (time, observed, trapped, free, M, nROIs, arrival_data)

def plot_model_curves(model_data, k, fileName):
    
    # Give our input with usable names
    time = model_data[0]
    observed = model_data[1]
    trapped = model_data[2]
    free = model_data[3]
    M = model_data[4]
    arrival_data = model_data[6]
    MAT_times = arrival_data[1]
    MAT_bound = arrival_data[4]
    sv = arrival_data[5]

    # Based on the maximum observed data point calculate the head and tail location of the MAT-indicating arrow.
    # This may need to change between plants
    # Arrows are drawn using an absolute location for the tail and dx, dy amounts for the location of the head
    MAT_arrow_param = float(max(observed))
    arrow_yTail = MAT_arrow_param * 1.1
    arrow_yHead = -1 * MAT_arrow_param * 0.1

    # Store the height of the data at the MAT upper bound; to be used to draw an arrow to mark the MAT upper bound
    UB_arrow_param = float(observed[MAT_bound])

    # Create a new figure for each plot
    plt.figure()

    # Plot the data
    plt.plot(time, observed, '.', label = 'Observed')
    plt.plot(time, trapped, label = 'Trapped', color = 'r')
    plt.plot(time, free, label = 'Free', color = 'g')
    plt.plot(time, M, label = 'Trapped + Free', color = 'k')

    # Make an arrow to point to the MAT as calculated by the model
    plt.arrow(MAT_times[k], arrow_yTail, 0, arrow_yHead, length_includes_head = "TRUE",
              fc="k", ec="k", head_width = 65,
              head_length = float(max(observed)) * .05)
    plt.annotate('MAT', xy =(MAT_times[k], arrow_yTail*1.005))

    # Make an arrow to point to the MAT upper bound
    plt.arrow(float(time[MAT_bound]), UB_arrow_param * 1.6, 0, UB_arrow_param * - 0.5,
              length_includes_head = "TRUE", fc="k", ec="k",
              head_width = 65, head_length = float(max(observed)) * .05)
    plt.annotate('UB', xy = (time[MAT_bound], float(observed[MAT_bound]) * 1.605))

    # Add a legend, title, x and y axis labels, and set axis max/min
    plt.legend(loc=1)
    plt.ylim(0, MAT_arrow_param * 1.3)
    plt.xlim(0,3600)
    plt.xlabel('Time (s)')
    plt.ylabel('Activity (Bq/mm) / Administered Activity (Bq)')
    # Title: TAC for RBO0## - cor_##
    plt.title('TAC for: ' + arrival_data[3] + " - " + arrival_data[0][k])
    plt.suptitle('sv: ' + str(sv[k]) + ' MAT: ' + str(MAT_times[k]))

    print "graphed"

    # Saved as: RBO0## cor_##.png
    # Saved in the same location as the python code
    plt.savefig(fileName[:-4] + "_" + arrival_data[0][k] + '.png')

# Used to change output of askopenfilenames from string to tuple
def string2list(input_string):
    input_string = input_string.lstrip('{')
    input_string = input_string.rstrip('}')
    output = input_string.split('} {')
    return output

# Perform modeling when button is clicked
def linear_regression():
    
    # Get a list of file names
    files = askopenfilenames()
    if not files:
	print '---> Please select at least one file\n'
	return

    if type(files) is not tuple:
    	files = string2list(files)
    i = 0

    while (i < len(files)):
        # Read and store the (ROI, MAT, Xi, subject, MAT_bound)
        arrival_data = read_ROI_summary(files[i])
        # Fit the data to a curve
        (slope, intercept) = polyfit(arrival_data[1], arrival_data[2], 1)
        # Create a label for the legend
        equation = 'Linear Regression\ny ={0}x + {1}'.format(round(slope, 4), round(intercept, 4))
        # Plot the data
        plot_linreg_data(arrival_data, slope, intercept, equation, files[i])
        
        i = i + 1

def graph_model_curves():

    # Get a list of file names
    files = askopenfilenames()
    if not files:
	print '---> Please select at least one file\n'
	return

    if type(files) is not tuple:
    	files = string2list(files)

    i = 0
    j = 0
    k = 0
    
    while (i < len(files)):
        # nROIs is needed so perform the graphing once
	# Read the header information and model data
        model_data = read_model_curves(files[i], k)

	# Plot and save the model data
        plot_model_curves(model_data, k, files[i])
        nROIs = model_data[5]

        while (j < nROIs):
	    # Read the header information and model data
            model_data = read_model_curves(files[i], k)
	    
	    # Plot and save the model data
            plot_model_curves(model_data, k, files[i])
            j = j + 1
            k = k + 1
        i = i + 1

def validate_percent_height():
    percent = float(raw_input('Enter percentage as a decimal (0.0 - 1.0) '))

    # Get a list of file names
    files = askopenfilenames()
    if not files:
	print '---> Please select at least one file\n'
	return

    if type(files) is not tuple:
    	files = string2list(files)
    i = 0

    model_data = read_model_curves(files[i], i)
    nROIs = model_data[5]

    while (i < nROIs):
        model_data = read_model_curves(files[0], i)
        time = model_data[0]
        observed = model_data[1]
        trapped = model_data[2]
        free = model_data[3]
        M = model_data[4]
        nROIs = model_data[5]
        arrival_data = model_data[6]
	ROIs = arrival_data[0]
        MAT_times = arrival_data[1]
        MAT_bound = arrival_data[4]
        maxFree = max(free)
    
	print ROIs[i]
        percentFree = percent * float(maxFree)

        diff = 1000
        percentIndex = -1

        j = free.index(maxFree)
    
        while (j < (28 - free.index(maxFree))):
            currDiff = abs(percentFree - float(free[j]))
            if (currDiff < diff):
                diff = currDiff
	        percentIndex = j
	    j = j + 1

        print percentIndex
        i = i + 1

### Main executing part of program ###

# Don't show an empty Tkinter frame when the program starts
root = Tkinter.Tk()
root.withdraw()

# Print the menu
print '0 - Perform McKay Modeling'
print '1 - Linear Regression of MAT and Position'
print '2 - Graph TACs'
print '3 - Clean ASIPro File'
print '4 - Exit'
print '5 - validate maxheight'
answer = raw_input('What would you like to do?\n')

# Get the timestamp
localtime = time.localtime()
timeString = time.strftime("%Y%m%d_%H%M", localtime)

# Act based on response
while (answer != "4"):
    if (answer == "0"):
        call_mckay_code(timeString)
	print '---> McKay model completed\n'
    elif (answer == "1"):
        linear_regression()
	print '---> Linear regression completed\n'
    elif (answer == "2"):
        graph_model_curves()
	print '---> TACs graphed\n'
    elif (answer == "3"):
	clean_input_file()
	print '---> Input file cleaned'
    elif (answer == "5"):
	validate_percent_height()
	print 'validate height displayed'
    else:
        print '---> Please choose one of the above responses\n'

    print '0 - Perform McKay Modeling'
    print '1 - Linear Regression of MAT and Position'
    print '2 - Graph TACs'
    print '3 - Clean ASIPro File'
    print '4 - Exit'
    print '5 - validate output'
    answer = raw_input('What would you like to do?\n')
    
print 'Exiting now'