"""

Description:



Getting started:


    Original Author:
        Erica Morrill
        Department of Biomedical Engineering
        Lerner Research Institute
        Cleveland Clinic
        Cleveland, OH
        morrile2@ccf.org

"""

import xml.etree.ElementTree as ET
import os
import Tkinter as tk
import pandas
import tkFileDialog
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np
import time
import math

class FileSelectionApp(tk.Tk):
    """Application to display all of the trials in a list"""

    def __init__(self):
        tk.Tk.__init__(self)

        home = os.path.expanduser('~')
        for dirname, subdirList, fileList in os.walk(home):
            for dir in subdirList:
                if "MULTIS_test" in dir and "studies" not in dirname:
                    multis_dir = dirname + '/' + dir

        try:
            multis_dir
        except NameError:
            multis_dir = tkFileDialog.askdirectory(title="Open MULTIS trials directory")

        self.directory = multis_dir

        self.getSubjects()
        self.title('Select Files')

        self.masterdf = pandas.DataFrame(data=None)

        self.columnconfigure(0, weight=1)

        self.var = []
        self.i = 0
        for item in self.subFiles:
            self.var.append(tk.IntVar())
            c = tk.Checkbutton(self, text=item, variable=self.var[self.i])
            c.grid(column = 0, row=self.i, sticky='w')
            self.i +=1

        tk.Button(self, text="Okay", command=self.checkBoxes).grid(row=1, column =1, sticky='ens')
        tk.Button(self, text="Select All", command=self.SelectAll).grid(row=0, column=1, sticky='ens')

    def yview(self, *args):
        apply(self.yview, args)

    def checkBoxes(self):

        s_time = time.time()
        locations  = ['LA_A', 'LA_P', 'UA_A', 'UA_P', 'LL_A','LL_P', 'UL_A', 'UL_P']
        # locations = ['LA_A'] #Used for testing

        for loc in locations:
            count = 0
            self.df = pandas.DataFrame(data=None)
            self.location = loc
            for bb in self.var:
                if bb.get() == 1:
                    xml = self.dir[count] + '/Configuration/' + os.path.split(self.dir[count])[1] + '.xml'
                    # self.saveDemographicData(xml)
                    self.saveData(count)
                count += 1

            #Build Function to create report
            self.avg = self.df.mean()
            self.stdDev = self.df.std()

            self.plot_data()
        self.masterdf.to_csv(self.directory + '/ThickSumFigsNew/' + '001_MasterList_indentation.csv')

        self.plot_SegLoc_bar()
        # self.plot_average_hist()
        # self.plot_SegLoc_hist()

        print("Elapsed Time: %f seconds" %float((time.time()-s_time)))
        plt.show()
        self.quit()

    def plot_SegLoc_bar(self):
        """Plots bar charts showing mean stiffness of each tissue for each segment/location pair. Error bars are standard deviation"""
        MEAN = self.masterdf.groupby(["Segment", "Location"]).mean()
        STD = self.masterdf.groupby(["Segment", "Location"]).std()

        print(MEAN)

        plt.figure()
        MEAN['Total_Stiff'].plot.bar(yerr=STD['Total_Stiff'])
        plt.axhline(0, color='k')
        plt.suptitle("Total Stiffness", fontsize=14)

        plt.figure()
        MEAN['Muscle_Stiff'].plot.bar(yerr=STD['Muscle_Stiff'])
        plt.axhline(0, color='k')
        plt.suptitle("Muscle Stiffness", fontsize=14)

        plt.figure()
        MEAN['Fat_Stiff'].plot.bar(yerr=STD['Fat_Stiff'])
        plt.axhline(0, color='k')
        plt.suptitle("Fat Stiffness", fontsize=14)

        plt.figure()
        MEAN['Skin_Stiff'].plot.bar(yerr=STD['Skin_Stiff'])
        plt.axhline(0, color='k')
        plt.suptitle("Skin Stiffness", fontsize=14)

    def plot_average_hist(self):
        """Plots histograms of average stiffness values for each tissue type."""
        fig = plt.figure(figsize=(16, 8))
        fig.add_subplot(141)
        plt.title("Skin_Stiffness \nAvg = %.2f +/- %.2f" % (
        self.masterdf["Skin_Stiff"].mean(), self.masterdf["Skin_Stiff"].std()))
        plt.hist(self.masterdf["Skin_Stiff"][~np.isnan(self.masterdf["Skin_Stiff"])])

        fig.add_subplot(142)
        plt.title("Fat_Stiffness \nAvg = %.2f +/- %.2f" % (
        self.masterdf["Fat_Stiff"].mean(), self.masterdf["Fat_Stiff"].std()))
        plt.hist(self.masterdf["Fat_Stiff"][~np.isnan(self.masterdf["Fat_Stiff"])])

        fig.add_subplot(143)
        plt.title("Muscle_Stiffness \nAvg = %.2f +/- %.2f" % (
        self.masterdf["Muscle_Stiff"].mean(), self.masterdf["Muscle_Stiff"].std()))
        plt.hist(self.masterdf["Muscle_Stiff"][~np.isnan(self.masterdf["Muscle_Stiff"])])

        fig.add_subplot(144)
        plt.title("Total_Stiffness \nAvg = %.2f +/- %.2f" % (
            self.masterdf["Total_Stiff"].mean(), self.masterdf["Total_Stiff"].std()))
        plt.hist(self.masterdf["Total_Stiff"][~np.isnan(self.masterdf["Total_Stiff"])])

    def plot_SegLoc_hist(self):
        """Plots histogram of each tissue for each segment/location pair."""
        self.masterdf["Total_Stiff"].hist(by=[self.masterdf["Segment"], self.masterdf["Location"]], figsize=(12, 8))
        plt.suptitle("Total Stiffness", fontsize=14)

        self.masterdf["Muscle_Stiff"].hist(by=[self.masterdf["Segment"], self.masterdf["Location"]], figsize=(12, 8))
        plt.suptitle("Muscle Stiffness", fontsize=14)

        self.masterdf["Fat_Stiff"].hist(by=[self.masterdf["Segment"], self.masterdf["Location"]], figsize=(12, 8))
        plt.suptitle("Fat Stiffness", fontsize=14)

        self.masterdf["Skin_Stiff"].hist(by=[self.masterdf["Segment"], self.masterdf["Location"]], figsize=(12, 8))
        plt.suptitle("Skin Stiffness", fontsize=14)

    def saveDemographicData(self, xml_name):

        doc = ET.parse(xml_name)
        root = doc.getroot()

        subjData = root.find("Subject_Data")
        Anatomical = subjData.find("Anatomical_Measurements")

        for child in Anatomical:
            if child.attrib['type'] == 'Cluster':
                for cl in child.findall("Cluster"):
                    if cl._children[0].text == 'Length':
                        length = float(cl._children[1].text)
                    elif cl._children[0].text == 'Distal Circumference':
                        circ_d = float(cl._children[1].text)
                    elif cl._children[0].text == 'Central Circumference':
                        circ_c = float(cl._children[1].text)
                    elif cl._children[0].text == 'Proximal Circumference':
                        circ_p = float(cl._children[1].text)

                if child.tag == 'Upper_Arm':
                    column_names = ['dc (cm)', 'cc (cm)', 'pc (cm)']
                    self.df_UA = pandas.DataFrame([[circ_d, circ_c, circ_p]], columns=column_names)
                elif child.tag == 'Lower_Arm':
                    column_names = ['dc (cm)', 'cc (cm)', 'pc (cm)']
                    self.df_LA = pandas.DataFrame([[circ_d, circ_c, circ_p]], columns=column_names)
                elif child.tag == 'Upper_Leg':
                    column_names = ['dc (cm)', 'cc (cm)', 'pc (cm)']
                    self.df_UL = pandas.DataFrame([[circ_d, circ_c, circ_p]], columns=column_names)
                elif child.tag == 'Lower_Leg':
                    column_names = ['dc (cm)', 'cc (cm)', 'pc (cm)']
                    self.df_LL = pandas.DataFrame([[circ_d, circ_c, circ_p]], columns=column_names)


    def saveData(self, count):
        # print(self.dir[count])

        files = self.getThickFiles(self.location, self.dir[count])
        # print(len(files))

        if len(files) > 0:

            for xml_name in files:
                # print(xml_name)
                # xml_name = self.dir[count]+'/'+self.subFiles[count]
                skin = []
                fat = []
                muscle = []
                force_mag = []
                total = []

                doc = ET.parse(xml_name)
                root = doc.getroot()

                subj = root.find('Subject')
                src = subj.find('Source')
                locs = src.findall("Frame")
                for loc in locs:
                    thick = loc.find("Thickness")
                    force = loc.find("Forces")
                    force_mag.append(math.sqrt((float(force.find("Fx").text))**2 + (float(force.find("Fy").text))**2 + (float(force.find("Fz").text))**2))
                    skin.append(float(thick.find("Skin").text))
                    fat.append(float(thick.find("Fat").text))
                    muscle.append(float(thick.find("Muscle").text))
                    total.append(skin[-1] + fat[-1] + muscle[-1])

                if 'C_' in os.path.split(xml_name)[1]:
                    column_names = ['Force_Magnitude', 'Skin_Central', 'Fat_Central', 'Muscle_Central']
                    df_cent = pandas.DataFrame([[force_mag, skin, fat, muscle]], columns=column_names)

                    column_names = ['MaxForce', 'IndentDepth(%)', 'SkinIndent', 'FatIndent', 'MuscleIndent']
                    df_data = pandas.DataFrame([[np.max(force_mag), (total[0]-total[-1])/total[0], (skin[0]-skin[-1])/skin[0], (fat[0]-fat[-1])/fat[0], (muscle[0]-muscle[-1])/muscle[0]]], columns=column_names)

                    column_names = ['Skin_Stiff', 'Fat_Stiff', 'Muscle_Stiff', 'Total_Stiff']

                    if min(force_mag) > 2:
                        print(os.path.split(xml_name), min(force_mag))

                    try:
                        p1 = np.polyfit(force_mag, skin, 1, full=True)
                        fit1 = np.poly1d(p1[0])
                        # r2 = np.corrcoef(skin, fit1(force_mag))[0,1]**2
                        # if r2<0.5:
                        #     print(xml_name,"skin", r2)
                        stiff_s = p1[0][0]

                    except:
                        stiff_s = np.nan
                    try:
                        p1 = np.polyfit(force_mag, fat, 1, full=True)
                        fit1 = np.poly1d(p1[0])
                        # r2 = np.corrcoef(fat, fit1(force_mag))[0,1]**2
                        # if r2<0.5:
                        #     print(xml_name,"fat", r2)
                        stiff_f = p1[0][0]
                    except:
                        stiff_f = np.nan
                    try:
                        p1 = np.polyfit(force_mag, muscle, 1, full=True)
                        fit1 = np.poly1d(p1[0])
                        # r2 = np.corrcoef(muscle, fit1(force_mag))[0,1]**2
                        # if r2<0.5:
                        #     print(xml_name,"muscle", r2)
                        stiff_m = p1[0][0]
                    except:
                        stiff_m = np.nan
                    try:
                        p1 = np.polyfit(force_mag, total, 1, full=True)
                        fit1 = np.poly1d(p1[0])
                        # r2 = np.corrcoef(total, fit1(force_mag))[0,1]**2
                        # if r2<0.5:
                        #     print(xml_name, "total", r2)
                        stiff_t = p1[0][0]
                    except:
                        stiff_t = np.nan

                    df_cent_stiff = pandas.DataFrame([[stiff_s, stiff_f, stiff_m, stiff_t]], columns=column_names)
                else:
                    print("Error", os.path.split(xml_name)[1])

                column_names = ['SubID', 'Segment', 'Location']
                df_demo = pandas.DataFrame([[xml_name[-40:-37], self.location[0:2], self.location[3]]], columns=column_names)

                try:
                    df_cent
                except NameError:
                    column_names = ['Force_Magnitude', 'Skin_Central', 'Fat_Central', 'Muscle_Central']
                    df_cent = pandas.DataFrame([[np.nan, np.nan, np.nan, np.nan]], columns=column_names)


                self.df = self.df.append(pandas.concat([df_cent], axis=1), ignore_index=True)
                self.masterdf = self.masterdf.append(pandas.concat([df_demo, df_data, df_cent_stiff], axis=1), ignore_index=True)



    def plot_data(self):

        fig, axes = plt.subplots(3, figsize=(12, 12))
        ii = 0
        # print(len(self.df["Force_Magnitude"]))
        for v in range(len(self.df["Force_Magnitude"])):
            # print("max", max(self.df["Muscle_Central"][v]))
            axes[0].plot(np.asarray(self.df["Muscle_Central"].values[v])-self.df["Muscle_Central"][v][0], self.df["Force_Magnitude"].values[v])
            axes[0].set_title("Muscle")
            axes[0].set_ylabel("Force (N)")
            axes[1].plot(np.asarray(self.df["Fat_Central"].values[v])-self.df["Fat_Central"][v][0], self.df["Force_Magnitude"].values[v])
            axes[1].set_title("Fat")
            axes[1].set_ylabel("Force (N)")
            axes[2].plot(np.asarray(self.df["Skin_Central"].values[v])-self.df["Skin_Central"][v][0], self.df["Force_Magnitude"].values[v])
            axes[2].set_title("Skin")
            axes[2].set_xlabel("Deformation (mm)")
            axes[2].set_ylabel("Force (N)")

        if self.location[0] == 'U':
            title = 'Upper'
        elif self.location[0] == 'L':
            title = 'Lower'

        if self.location[1] == 'A':
            title += ' Arm'
        elif self.location[1] == 'L':
            title += ' Leg'

        if self.location[3] == 'A':
            title += ' Anterior'
        elif self.location[3] == 'P':
            title += ' Posterior'

        fig.suptitle(title, fontsize = 20)
        plt.tight_layout()
        fig.subplots_adjust(top = 0.9)
        fig.savefig(self.directory + '/ThickSumFigsNew/Indentation_' + self.location + '.png')
        plt.close()


    def SelectAll(self):
        for bb in self.var:
            bb.set(1)


    def getSubjects(self):
        self.subFiles = []
        self.dir = []
        for dirname,subdirList,fileList in os.walk(self.directory):
            for file in fileList:
                if "TA_inclusion.xml" in file:
                    self.dir.append(os.path.split(os.path.split(dirname)[0])[0])
                    self.subFiles.append(os.path.split(os.path.split(os.path.split(dirname)[0])[0])[1])
        sortedTrials = sorted(zip(self.subFiles, self.dir))
        self.subFiles, self.dir = zip(*sortedTrials)

    def getThickFiles(self, location, directory):
        loc_files = []

        thick_xml_list = self.parseXMLinclusion(location, os.path.join(os.path.join(os.path.join(directory,'TissueThickness'), 'UltrasoundManual'), os.path.split(directory)[1]+"_TA_inclusion.xml"))

        for th_xml in thick_xml_list:
            loc_files.append(os.path.join(directory, th_xml))

        return loc_files

    def parseXMLinclusion(self, location, inclusion):

        doc = ET.parse(inclusion)
        root = doc.getroot()
        thickness_xmls = []

        for child in root:
            for subChild in child:
                try:
                    XML = subChild.attrib["Indentation"]
                    if XML != "None" and XML[-28:-27] == "I" and XML[-34:-30] == location:
                        thickness_xmls.append(XML)
                except:
                    continue

        return thickness_xmls

if __name__ == "__main__":

    app = FileSelectionApp()
    app.mainloop()