import lxml.etree as ET
import pandas as pd
from train_data import TrainData

CSV_TRAINSET = "..\\csv_tables\\train_extended.csv"
XML_TRAINSET = "..\\csv_tables\\trainset.xml"

def write_to_xml(filename, samples, features):
    """Trainset xml serialization

    Arguments:
    filename(str): output filename
    samples(dict) : dictionary sampels[class_name] = 
                    [retv, color, class_ID,  [(sample_ID, [feature_arr]), 
                                              (sample_ID, [feature_arr])]]

    features(dict) : features['feature_name'] = (Index, Weight)
    """

    root = ET.Element("Klasifikace")
    common = ET.SubElement(root, "Common")
    ET.SubElement(common, "AutoAddStrangeClasses", Enabled="1", 
                                                   ThreshValue="800").text = ""

    priznaky = ET.SubElement(root, "Priznaky", Count=str(len(features)))

    for feature in sorted(features.items(), key=lambda x: x[1][0]): # sort by idx

        ET.SubElement(priznaky, "P", Name=feature[0], 
                                     Index=str(feature[1][0]),
                                     Vaha=str(feature[1][1])).text = ""

    tridy = ET.SubElement(root, "Tridy") 
    for cl in samples:
        retv = str(samples[cl][0])
        color = str(samples[cl][1])
        ID = str(samples[cl][2])
        t = ET.SubElement(tridy, "T", Name=cl, Retv=retv, Color=color, ID=ID)
        for s in samples[cl][3]:
            ET.SubElement(t, "V", ID=str(s[0])).text = ";".join(map(str, s[1]))

    tree = ET.ElementTree(root)
    tree.write(filename, pretty_print=True)



if __name__ == "__main__":
    # read trainset from csv and write to xml
    df = pd.read_csv(CSV_TRAINSET, sep = ';')
    df['filename'] = map(lambda x: x[x.index('FRQ'):],  df['filename'])
    df = TrainData.include_classes(df, TrainData.MAP_CLASSES_FINAL.keys())
    df.loc[:,'class_name'] = df.loc[:,'class_name'].map( \
                         lambda x: TrainData.MAP_CLASSES_FINAL[x][0] + " -" + x)

    classes = sorted(df['class_name'].unique())
    samples = dict()
    features = dict()
    for idx, cl in enumerate(classes):
        classID = idx
        retv = int(cl[0])
        color = "ffff"
        samples[cl] = [retv, color, classID, []]
        for _, row in df[df['class_name'] == cl].iterrows():
            sampleID = row['filename']
            feature_arr = row[TrainData.FEATURES].tolist()
            feature_arr = map(lambda x: int(1000*x), feature_arr)
            samples[cl][3].append((sampleID, feature_arr))

    for idx, f in enumerate(TrainData.FEATURES):
        features[f] = (idx, 1000)

    write_to_xml(XML_TRAINSET, samples, features)
