import xml.etree.ElementTree as ET
from pylab import *
import numpy as np
import os
import tkMessageBox
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.ticker import FormatStrFormatter

def createSubplot(AX, X, Y, C, L):

    try:
        p1 = np.polyfit(X, Y, 1)
        fit1 = np.poly1d(p1)
        AX.plot(X, fit1(X), color=C)
        AX.scatter(X, Y, color=C, label=L % (p1[0]))
    except:
        AX.scatter(X, Y, color=C, label=L % (np.nan))

    start, end = AX.get_ylim()
    AX.yaxis.set_ticks(np.linspace(start, end, 3))
    AX.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
    AX.xaxis.set_major_formatter(FormatStrFormatter('%.2f'))

    AX.grid(True)

def PlotSum(xmlname):

    path = os.path.split(xmlname)

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

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

    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))
        if thick.find("Muscle").text == 'nan':
            muscle.append(np.nan)
        else:
            muscle.append(float(thick.find("Muscle").text))

    if True in isnan(muscle):
        total = list(np.empty(len(muscle))*np.nan)
    else:
        total = [sum(x) for x in zip(skin, fat, muscle)]

    fig = Figure(figsize=(9.5,4.5))
    canvas = FigureCanvas(fig)
    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)

    createSubplot(ax1, F_mag, total, 'y', "Total (m = %.4f)")
    createSubplot(ax2, F_mag, muscle, 'g', "Muscle (m = %.4f)")
    createSubplot(ax3, F_mag, fat, 'b', "Fat (m = %.4f)")
    createSubplot(ax4, F_mag, skin, 'r', "Skin (m = %.4f)")

    # 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=False, bottom=False, left=False, right=False)
    ax.set_facecolor('none')

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

    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(left = 0.15, right=0.6)

    fig.savefig(path[0]+'/ThicknessPNG/'+path[1][0:-4]+'Graph.png', bbox_extra_artists = (lgd, ), bbox_inches = 'tight')
    plt.close(fig)