import xml.etree.ElementTree as ET
from lxml import etree
import XMLparser
import tkFileDialog
import os
import zipfile
import numpy as np

def prettyPrintXml(xmlFilePathToPrettyPrint):
    """Pretty print the xml file after all frames have been analyzed"""
    assert xmlFilePathToPrettyPrint is not None
    parser = etree.XMLParser(resolve_entities=False, remove_blank_text=True)
    document = etree.parse(xmlFilePathToPrettyPrint, parser)
    document.write(xmlFilePathToPrettyPrint, pretty_print=True, encoding='ISO-8859-1', xml_declaration=True)

# def newCluster(subj, tdms_path, root_incl, xml_inclusion, thick_file):
#     cluster_new = ET.SubElement(subj, 'Cluster')
#     # src_tdms = ET.SubElement(cluster_new, 'TDMSSource', attrib={"Filename": tdms_path[-30:]})
#     # src_TA = ET.SubElement(cluster_new, 'ThicknessSource', attrib={"Filename":thick_file})
#     # tree = ET.ElementTree(root_incl)
#     # tree.write(xml_inclusion, xml_declaration=True)

def TAxmlFormat(subj, seg, mod):
    root = ET.Element("UltrasoundPositions")
    root.set("subject", subj)
    root.set("modality", mod)
    SegLocations(ET.SubElement(root, seg), mod)
    return root

def SegLocations(seg, m):
    USPositions(ET.SubElement(seg, "ProximalLateral"), m)
    USPositions(ET.SubElement(seg, "ProximalAnterior"), m)
    USPositions(ET.SubElement(seg, "ProximalMedial"), m)
    USPositions(ET.SubElement(seg, "ProximalPosterior"), m)
    USPositions(ET.SubElement(seg, "CentralLateral"), m)
    USPositions(ET.SubElement(seg, "CentralAnterior"), m)
    USPositions(ET.SubElement(seg, "CentralMedial"), m)
    USPositions(ET.SubElement(seg, "CentralPosterior"), m)
    USPositions(ET.SubElement(seg, "DistalLateral"), m)
    USPositions(ET.SubElement(seg, "DistalAnterior"), m)
    USPositions(ET.SubElement(seg, "DistalMedial"), m)
    USPositions(ET.SubElement(seg, "DistalPosterior"), m)

def USPositions(parent, modality):
    # parent.set('Anatomical', 'None')
    A = ET.SubElement(parent, "Anatomical", file = 'None', CoordinateSys = modality)

    if parent.tag == 'CentralAnterior' or parent.tag == 'CentralPosterior':
        # parent.set('Indentation', 'None')
        A = ET.SubElement(parent, "Indentation", file='None', CoordinateSys = modality)

def add_US_idx(A, f, f_US, position):
    P = ET.SubElement(A, "USPosition", attrib={"OptotrakIndex": str(f[0]), "USFrame": str(f_US)})
    ET.SubElement(P, 'x', units='m', value=str(position[0]))
    ET.SubElement(P, 'y', units='m', value=str(position[1]))
    ET.SubElement(P, 'z', units='m', value=str(position[2]))
    ET.SubElement(P, 'roll', units='rad', value=str(position[3]))
    ET.SubElement(P, 'pitch', units='rad', value=str(position[4]))
    ET.SubElement(P, 'yaw', units='rad', value=str(position[5]))

def edit_values(seg, location, testtype, root, position, tdmsName, f, f_US):

    position = np.fromstring(position[1:-1], sep=" ")
    loc1 = location[1]
    loc2 = location[0]

    if loc1 == 'P':
        loc = 'Proximal'
    elif loc1 == 'C':
        loc = 'Central'
    elif loc1 == 'D':
        loc = 'Distal'

    if loc2 == 'L':
        loc = loc + 'Lateral'
    elif loc2 == 'M':
        loc = loc + 'Medial'
    elif loc2 == 'A':
        loc = loc + 'Anterior'
    elif loc2 == 'P':
        loc = loc + 'Posterior'

    if testtype == 'A':
        a = 'Anatomical'
    elif testtype == 'I':
        a = 'Indentation'

    A = root.find(seg).find(loc).find(a)
    A.set('file', tdmsName)

    add_US_idx(A, f, f_US, position)

# def create_zip(path, ziph, xmlPath):
#     path = os.path.join(path, 'Analysis_final')
#     for root, dirs, files in os.walk(path):
#         for file in files:
#             ziph.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), start=os.path.split(path)[0]))
#     ziph.write(xmlPath, os.path.split(xmlPath)[1])

def create_xml(subj_path, seg, segment, modality, reg_id):
    """Create xml file with forces, moments, and thicknesses"""

    if not os.path.exists(os.path.join(subj_path, "Registration", reg_id, "USPositions", modality)):
        os.makedirs(os.path.join(subj_path, "Registration", reg_id, "USPositions", modality))

    # Create inclusion xml
    xml_inclusion = os.path.join(subj_path, "Registration", reg_id, "USPositions", modality, reg_id + '_'+ os.path.split(subj_path)[1]+seg+'_US_'+modality[0:2]+'.xml')
    subj = os.path.split(subj_path)[1]
    root_incl = TAxmlFormat(subj, segment, modality[0:2])

    return xml_inclusion, root_incl


def writeXML(xml_inclusion, root_incl, subj_path):
    tree = ET.ElementTree(root_incl)
    tree.write(xml_inclusion, xml_declaration=True)
    prettyPrintXml(os.path.join(subj_path, xml_inclusion))