#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
input: prohane data
output: data for ktImportText
"""
import argparse
import os
import pandas as pd
import sys
import xml.etree.ElementTree as ET


def check_file_existence(lst_files):
    lst_missing_files = list()
    b_missing_files = False
    for f in lst_files:
        if not os.path.isfile(f):
            b_missing_files = True
            lst_missing_files.append(f)
    if b_missing_files:
        raise IOError("File(s) do not exist: {}".format(lst_missing_files))


def create_sub_nodes(node_in, df_in, col_amount, col_desc, lst_samples):
    c_level = col_amount.pop(0)

    groups = df_in.groupby(c_level)
    for classifier, df in groups:
        node = ET.SubElement(node_in, "node", attrib={"name": classifier})
        xml_add_amounts(df, lst_samples, node)

        c_desc = get_description_column_name(c_level, col_desc)
        if c_desc:
            xml_add_desc(df, lst_samples, node, c_desc)

        if col_amount:
            create_sub_nodes(node, df, list(col_amount), col_desc, lst_samples)


def get_description_column_name(c_level, col_desc):
    c_desc = None
    lst_desc = [c_desc for c_desc in col_desc if (c_level in c_desc) and ("_desc" in c_desc)]
    if len(lst_desc) == 1:
        c_desc = lst_desc[0]
    elif len(lst_desc) > 1:
        raise ValueError(f"'{c_desc}' in {lst_desc} more than once. I do not know which to choose.")
    return c_desc


def convert_to_krona_xml(in_file):
    df = pd.read_csv(in_file, sep='\t')
    df = df.drop(df.filter(regex="sd").columns, axis=1)
    lst_samples = df["sample"].unique().tolist()
    col_desc = list(df.filter(regex="descr").columns)
    col_amount = [c for c in df.columns if c not in (["sample", "quant"])]

    df = remove_unnecessary_vals(df)

    tree = ET.ElementTree()
    krona = ET.Element("krona")
    tree._setroot(krona)

    krona.set("collapse", "False")
    attrs = ET.SubElement(krona, "attributes")
    attrs.set("magnitude", "amount")
    a_amount = ET.SubElement(attrs, "attribute", attrib={"display": "Total"})
    a_amount.text = "amount"
    a_desc = ET.SubElement(attrs, "attribute", attrib={"display": "Description"})
    a_desc.text = "desc"

    # datasets
    datasets = ET.SubElement(krona, "datasets")
    for s in lst_samples:
        dataset = ET.SubElement(datasets, "dataset")
        dataset.text = s

    node = ET.SubElement(krona, "node", attrib={"name": "all"})
    xml_add_amounts(df, lst_samples, node)

    create_sub_nodes(node, df, col_amount, col_desc, lst_samples)

    f_out = ".".join([
        os.path.splitext(in_file)[0],
        "xml"
    ])
    print("Writing output to: {}".format(f_out))
    tree.write(f_out, encoding="utf-8")
    return f_out


def xml_add_amounts(df, lst_samples, node):
    node_amount = ET.SubElement(node, "amount")
    for s in lst_samples:
        node_val = ET.SubElement(node_amount, "val")
        df_s = df[df["sample"] == s]
        node_val.text = str(df_s["quant"].sum())


def xml_add_desc(df, lst_samples, node, c_desc):
    node_desc = ET.SubElement(node, "desc")
    for s in lst_samples:
        node_val = ET.SubElement(node_desc, "val")
        df_s = df[df["sample"] == s]
        desc_vals = df_s[c_desc].values
        if len(desc_vals) >= 1:
            node_val.text = df_s[c_desc].values[0]


def convert_to_krona_txt(in_file):
    df_raw = pd.read_csv(in_file, sep='\t')
    # split by sample
    groups = df_raw.groupby("sample")
    lst_f_out = list()
    for sample, df in groups:
        df = reorder_columns(df)
        df = remove_unnecessary_cols_and_vals(df)
        f_out = ".".join([
            os.path.splitext(in_file)[0],
            sample,
            "txt"
        ])
        print("Writing output to: {}".format(f_out))
        df.to_csv(f_out, sep="\t", header=False, index=False)
        lst_f_out.append(f_out)
    return lst_f_out


def remove_unnecessary_cols_and_vals(df):
    cols_re_to_drop = ["sample", r".*desc", "sd"]

    for r in cols_re_to_drop:
        df = df.drop(df.filter(regex=r).columns, axis=1)

    df = remove_unnecessary_vals(df)
    return df


def remove_unnecessary_vals(df):
    vals_to_drop = ["None"]
    df = df.replace(vals_to_drop, "")
    return df


def reorder_columns(df):
    quant = df.quant
    df = df.drop("quant", axis=1)
    df.insert(0, "quant", quant)
    return df


def parse_args():
    parser = argparse.ArgumentParser(
        description="Input file(s) from Prophane in arguments are split up by sample and converted to ktImportTxt" +
                    "compatible text file(s)."
    )
    parser.add_argument('input_files', type=str, nargs='+',
                        help='Prophane quant files to be processed.')
    parser.add_argument('-t', '--type', type=str, default="xml",
                        help="type of output file. Can be 'tsv' for ktImportText or 'xml' for ktImportXml.")

    if len(sys.argv) < 2:
        parser.print_help()
        sys.exit(1)
    return parser.parse_args()


def main():
    args = parse_args()
    lst_f_out = list()

    lst_f_in = args.input_files
    type_out = args.type

    if type_out.lower() == "tsv":
        fun_convert = convert_to_krona_txt
    elif type_out.lower() == "xml":
        fun_convert = convert_to_krona_xml
    else:
        raise AttributeError("Specified output filetype '{}' is unknown.".format(type_out))

    check_file_existence(lst_f_in)
    for f in lst_f_in:
        lst_f_out.append(fun_convert(f))


if __name__ == "__main__":
    main()
