import xml.etree.ElementTree as ET
from pylab import *
import numpy as np
import os
import tkMessageBox

def PlotSum(XMLS):

    F_mag = []
    skin = []
    fat = []
    muscle = []

    # Pares xml files to extract data from tissue thickness analysis
    for xmlname in XMLS:

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

        subj = root.find('Subject')
        src = subj.find('Source')

        for loc in src.findall("Frame"):
            forces = loc.find("Forces")
            F_mag.append(math.sqrt((float(forces.find("Fx").text))**2+(float(forces.find("Fy").text))**2+(float(forces.find("Fz").text))**2))
            thick = loc.find("Thickness")
            skin.append(float(thick.find("Skin").text))
            fat.append(float(thick.find("Fat").text))
            muscle.append(float(thick.find("Muscle").text))


    total = [sum(x) for x in zip(skin, fat, muscle)] # Total combined thickness of all layers

    # Plotting and fitting a line to the data
    fig = plt.figure(figsize=(16,10))
    ax = fig.add_subplot(111)
    ax1 = fig.add_subplot(411)
    ax2 = fig.add_subplot(412)
    ax3 = fig.add_subplot(413)
    ax4 = fig.add_subplot(414)

    if len(F_mag) > 1:
        p1 = np.polyfit(F_mag, total, 1)
        p2 = np.polyfit(F_mag, muscle, 1)
        p3 = np.polyfit(F_mag, fat, 1)
        p4 = np.polyfit(F_mag, skin, 1)
    else:
        tkMessageBox.showerror("Error", "The plots were not formed because only one data point was measured")
        return

    fit1 = np.poly1d(p1)
    fit2 = np.poly1d(p2)
    fit3 = np.poly1d(p3)
    fit4 = np.poly1d(p4)

    # Turn off axis lines and ticks of the big subplot
    ax.spines['top'].set_color('none')
    ax.spines['bottom'].set_color('none')
    ax.spines['left'].set_color('none')
    ax.spines['right'].set_color('none')
    ax.tick_params(labelcolor='none', top='off', bottom='off', left='off', right='off')
    ax.set_axis_bgcolor('none')

    ax1.scatter(F_mag, total, color='y', label="Total (m = %.4f)" %(p1[0]))
    ax1.plot(F_mag, fit1(F_mag), color='y')
    ax2.scatter(F_mag, muscle, color='g', label="Muscle (m = %.4f)" % (p2[0]))
    ax2.plot(F_mag, fit2(F_mag), color='g')
    ax3.scatter(F_mag, fat, color='b', label="Fat (m = %.4f)" % (p3[0]))
    ax3.plot(F_mag, fit3(F_mag), color='b')
    ax4.scatter(F_mag, skin, color='r', label="Skin (m = %.4f)" % (p4[0]))
    ax4.plot(F_mag, fit4(F_mag), color='r')

    ax.set_xlabel("Force Magnitude (N)")
    ax.set_ylabel("Thickness (mm)")
    ax.yaxis.set_label_coords(-0.2, 0.5)

    if xmlname[-28:-27] is 'I':
        start, end = ax1.get_ylim()
        ax1.yaxis.set_ticks(np.linspace(start, end, 3))
        ax1.ticklabel_format(useOffset=False)
        start, end = ax2.get_ylim()
        ax2.yaxis.set_ticks(np.linspace(start, end, 3))
        ax2.ticklabel_format(useOffset=False)
        start, end = ax3.get_ylim()
        ax3.yaxis.set_ticks(np.linspace(start, end, 3))
        ax3.ticklabel_format(useOffset=False)
        start, end = ax4.get_ylim()
        ax4.yaxis.set_ticks(np.linspace(start, end, 3))
        ax4.ticklabel_format(useOffset=False)
    else:
        start, end = ax1.get_ylim()
        ax1.yaxis.set_ticks(np.linspace(start, end, 3))
        ax1.ticklabel_format(useOffset=False)
        start, end = ax2.get_ylim()
        ax2.yaxis.set_ticks(np.linspace(start, end, 3))
        ax2.ticklabel_format(useOffset=False)
        start, end = ax3.get_ylim()
        ax3.yaxis.set_ticks(np.linspace(start, end, 3))
        ax3.ticklabel_format(useOffset=False)
        start, end = ax4.get_ylim()
        ax4.yaxis.set_ticks(np.linspace(start, end, 3))
        ax4.ticklabel_format(useOffset=False)


    ax1.grid("on")
    ax2.grid("on")
    ax3.grid("on")
    ax4.grid("on")

    h1, l1 = ax1.get_legend_handles_labels()
    h2, l2 = ax2.get_legend_handles_labels()
    h3, l3 = ax3.get_legend_handles_labels()
    h4, l4 = ax4.get_legend_handles_labels()

    lgd = ax1.legend(h1+h2+h3+h4, l1+l2+l3+l4, bbox_to_anchor = (1.05,1),loc="upper left")
    fig.tight_layout()
    fig.subplots_adjust(right=0.75)
    plt.show()

if __name__ == "__main__":
    xml_list = ['/home/morrile2/Documents/MULTIS Data/MULTIS_test/MULTIS008-1/TissueThickness/UltrasoundManual/005_MULTIS008-1_UA_AC_I-2_manThick201612211012.xml']
    PlotSum(xml_list)