import matplotlib

matplotlib.use("TkAgg")
import os
import xml.etree.ElementTree as ET
import pandas

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

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

    for th_xml in thick_xml_list:
        loc_files.append(os.path.join(directory, os.path.split(th_xml)[1]))

    return loc_files


def parseXMLinclusion(location, inclusion):
    doc = ET.parse(inclusion)
    root = doc.getroot()
    thickness_xmls = []

    for child in root:
        for subChild in child:
            XML = subChild.attrib["Anatomical"]
            if XML != "None" and XML[-28:-27] == "A" and XML[-34:-31] == location:
                thickness_xmls.append(XML)

    return thickness_xmls

def get_thicknesses(fname):

    tree = ET.parse(fname)
    root = tree.getroot()
    segment = ''.join([c for c in root.getchildren()[0].tag if c.isupper()])

    analysis_type = ["UltrasoundManual", "CTManual_zProbe", "MRManual_zProbe"]
    df_new = pandas.DataFrame(data=None)

    for type in analysis_type:
        print type
        df = pandas.DataFrame(data=None)
        US_thick_XMLS = getThickFiles(segment+'_', os.path.join(os.path.split(fname)[0], os.path.join('TissueThickness', type)))

        for loc in list(list(root)[0]):

            LOC = ''.join([c for c in loc.tag if c.isupper()])
            LOC_f = segment+'_'+LOC[-1]+LOC[0]
            print LOC_f
            tree_US = ET.parse([s for s in US_thick_XMLS if LOC_f in s][0])
            root_US = tree_US.getroot()
            subj = root_US.find('Subject')
            src = subj.find('Source')
            if type == "UltrasoundManual":
                loc = src.find("Frame")
            else:
                loc = src.find(type[0:2]+"Frame")
                # loc = src.find("CTFrame")
            thick = loc.find("Thickness")

            skin = (float(thick.find("Skin").text))
            fat = (float(thick.find("Fat").text))
            muscle = (float(thick.find("Muscle").text))

            # print LOC_f, skin, fat, muscle, skin_US, fat_US, muscle_US
            if type == "UltrasoundManual":
                df = df.append(pandas.DataFrame([[LOC_f, skin, fat, muscle]], columns = ['Location', type+'_Skin', type+'_Fat', type+'Muscle']), ignore_index=True)
            else:
                df = df.append(pandas.DataFrame([[skin, fat, muscle]], columns = [type+'_Skin', type+'_Fat', type+'Muscle']), ignore_index=True)

        df_new = pandas.concat([df_new, df], axis = 1)

    df_new.to_csv(os.path.join(os.path.split(fname)[0], 'THICKNESScomparison_'+segment+'.csv'))


fname = '/home/morrile2/Documents/MULTIS Data/MULTIS_invitro/CMULTIS002-3/CMULTIS002-3_UA_US_CT.xml'

get_thicknesses(fname)