import re
import xml.sax
import argparse
from time import time
from concurrent.futures import ThreadPoolExecutor, as_completed


class BiosampleHandler(xml.sax.ContentHandler):
    # 创建一个BiosampleHandler类，继承xml.sax.ContentHandler类
    def __init__(self, output):
        self.CurrentData = ""
        self.n = ""
        self.output = output
        self.bio_samples = []  # 存储每个BioSample的字典

    # 文档启动的时候调用
    def startDocument(self):
        self.__title = "Biosample_accession\tBioProject\tSRA\tobserved biotic relationship\tTitle\tPackage\tHost\ttaxonomy_id\ttaxonomy_name\tisolation_source\tgeo_loc_name\tcollection_date\tBacterial Vaginosis\tsample_type\thost_disease\tlatitude and longitude\tpH\tTemperature Optimum\tTemperature Range\tgrowth condition\tisol_growth_condt\tGram Staining\tgrowth medium\trelationship to oxygen\tenvironmental package\tenvironment\tsalinity\tPhenotypes\n"
        with open(self.output, "w", encoding="utf-8", errors="ignore") as biosample_extract:
            biosample_extract.writelines(self.__title)

    # 元素开始事件处理
    def startElement(self, tag, attributes):
        self.CurrentData = tag
        if self.CurrentData == "BioSample":
            self.dic = {"Biosample_accession": "", "BioProject": "", "SRA": "", "observed biotic relationship": "", "Title": "", "Package": "", "Host": "",
                    "taxonomy_id": "", "taxonomy_name": "", "geo_loc_name": "",
                    "latitude and longitude": "", "pH": "", "Temperature Optimum": "", "Temperature Range": "",
                    "collection_date": "", "isolation_source": "", "Bacterial Vaginosis": "",
                    "sample_type": "", "host_disease": "", "growth condition":"", "isol_growth_condt":"", "Gram Staining":"", "growth medium":"", 
                    "relationship to oxygen":"", "environmental package":"", "environment":"", "salinity":"", "Phenotypes":""}
            if "accession" in attributes:
                self.dic["Biosample_accession"] = attributes["accession"]
            else:
                self.dic["Biosample_accession"] = "not_exist"
        elif self.CurrentData == "Attribute":
            if "display_name" in attributes and attributes["display_name"] == "host":
                self.n = "host"
            if "display_name" in attributes and attributes["display_name"] == "isolation source":
                self.n = "isolation source"
            if "display_name" in attributes and attributes["display_name"] == "geographic location":
                self.n = "geographic location"
            if "display_name" in attributes and attributes["display_name"] == "collection date":
                self.n = "collection date"
            if "display_name" in attributes and attributes["attribute_name"] == "Bacterial Vaginosis":
                self.n = "Bacterial Vaginosis"
            if "display_name" in attributes and attributes["attribute_name"] == "sample_type":
                self.n = "sample_type"
            if "display_name" in attributes and attributes["display_name"] == "host disease":
                self.n = "host disease"
            if "display_name" in attributes and attributes["display_name"] == "observed biotic relationship":
                self.n = "observed biotic relationship"
            if "display_name" in attributes and attributes["display_name"] == "latitude and longitude":
                self.n = "latitude and longitude"
            if "attribute_name" in attributes and attributes["attribute_name"] == "pH":
                self.n = "pH"
            if "attribute_name" in attributes and attributes["attribute_name"] == "Temperature Optimum":
                self.n = "Temperature Optimum"
            if "attribute_name" in attributes and attributes["attribute_name"] == "Temperature Range":
                self.n = "Temperature Range"
            if "attribute_name" in attributes and attributes["attribute_name"] == "growth condition":
                self.n = "growth condition"
            if "attribute_name" in attributes and attributes["attribute_name"] == "isol_growth_condt":
                self.n = "isol_growth_condt"
            if "display_name" in attributes and attributes["display_name"] == "relationship to oxygen":
                self.n = "relationship to oxygen"
            if "attribute_name" in attributes and attributes["attribute_name"] == "Gram Staining":
                self.n = "Gram Staining"
            if "attribute_name" in attributes and attributes["attribute_name"] == "growth medium":
                self.n = "growth medium"
            if "attribute_name" in attributes and attributes["attribute_name"] == "environmental package":
                self.n = "environmental package"
            if "attribute_name" in attributes and attributes["attribute_name"] == "environment":
                self.n = "environment"
            if "display_name" in attributes and attributes["display_name"] == "salinity":
                self.n = "salinity"
            if "attribute_name" in attributes and attributes["attribute_name"] == "Phenotypes":
                self.n = "Phenotypes"
        elif self.CurrentData == "Organism":
            if "taxonomy_id" in attributes:
                self.dic["taxonomy_id"] = attributes["taxonomy_id"]
            else:
                self.dic["taxonomy_id"] = "not_exist"
            if "taxonomy_name" in attributes:
                self.dic["taxonomy_name"] = attributes["taxonomy_name"]
        elif self.CurrentData == "Link":
            if "type" in attributes and attributes["type"] == "entrez" and "target" in attributes and attributes["target"] == "bioproject":
                if "label" in attributes:
                    self.dic["BioProject"] = attributes["label"]
        elif self.CurrentData == "Id":
            if "db" in attributes and attributes["db"] == "SRA":
                self.n = "SRA"

    # 元素结束事件处理
    def endElement(self, tag):
        self.n = ""
        if tag == "BioSample":
            self.bio_samples.append(self.dic)  # 保存当前BioSample的数据
        self.CurrentData = ""

    # 内容事件处理
    def characters(self, content):
        if self.CurrentData == "Title":
            self.dic["Title"] = content
        elif self.CurrentData == "Package":
            self.dic["Package"] = content
        elif self.n == "host":
            self.dic["Host"] = content
        elif self.n == "isolation source":
            self.dic["isolation_source"] = content
        elif self.n == "geographic location":
            self.dic["geo_loc_name"] = content
        elif self.n == "collection date":
            self.dic["collection_date"] = content
        elif self.n == "Bacterial Vaginosis":
            self.dic["Bacterial Vaginosis"] = content
        elif self.n == "sample_type":
            self.dic["sample_type"] = content
        elif self.n == "host disease":
            self.dic["host_disease"] = content
        elif self.n == "observed biotic relationship":
            self.dic["observed biotic relationship"] = content
        elif self.n == "latitude and longitude":
            self.dic["latitude and longitude"] = content
        elif self.n == "pH":
            self.dic["pH"] = content
        elif self.n == "Temperature Optimum":
            self.dic["Temperature Optimum"] = content
        elif self.n == "Temperature Range":
            self.dic["Temperature Range"] = content
        elif self.n == "growth condition":
            self.dic["growth condition"] = content
        elif self.n == "isol_growth_condt":
            self.dic["isol_growth_condt"] = content
        elif self.n == "Gram Staining":
            self.dic["Gram Staining"] = content
        elif self.n == "growth medium":
            self.dic["growth medium"] = content
        elif self.n == "relationship to oxygen":
            self.dic["relationship to oxygen"] = content
        elif self.n == "environmental package":
            self.dic["environmental package"] = content
        elif self.n == "environment":
            self.dic["environment"] = content
        elif self.n == "salinity":
            self.dic["salinity"] = content
        elif self.n == "Phenotypes":
            self.dic["Phenotypes"] = content
        elif self.n == "SRA":
            self.dic["SRA"] = content

    # 处理每个BioSample并写入文件
    def process_biosample(self, biosample):
        with open(self.output, "a", encoding="utf-8", errors="ignore") as biosample_extract:
            biosample_extract.writelines(
                re.sub("[\t\n]", "", biosample["Biosample_accession"]) + "\t" +
                re.sub("[\t\n]", "", biosample["BioProject"]) + "\t" +
                re.sub("[\t\n]", "", biosample["SRA"]) + "\t" +
                re.sub("[\t\n]", "", biosample["observed biotic relationship"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Title"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Package"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Host"]) + "\t" +
                re.sub("[\t\n]", "", biosample["taxonomy_id"]) + "\t" +
                re.sub("[\t\n]", "", biosample["taxonomy_name"]) + "\t" +
                re.sub("[\t\n]", "", biosample["isolation_source"]) + "\t" +
                re.sub("[\t\n]", "", biosample["geo_loc_name"]) + "\t" +
                re.sub("[\t\n]", "", "'" + biosample["collection_date"]) + "'" + "\t" +
                re.sub("[\t\n]", "", biosample["Bacterial Vaginosis"]) + "\t" +
                re.sub("[\t\n]", "", biosample["sample_type"]) + "\t" +
                re.sub("[\t\n]", "", biosample["host_disease"]) + "\t" +
                re.sub("[\t\n]", "", biosample["latitude and longitude"]) + "\t" +
                re.sub("[\t\n]", "", biosample["pH"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Temperature Optimum"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Temperature Range"]) + "\t" +
                re.sub("[\t\n]", "", biosample["growth condition"]) + "\t" +
                re.sub("[\t\n]", "", biosample["isol_growth_condt"]) + "\t" +
                re.sub("[\t\n]", "", biosample["Gram Staining"]) + "\t" +
                re.sub("[\t\n]", "", biosample["growth medium"]) + "\t" +
                re.sub("[\t\n]", "", biosample["relationship to oxygen"]) + "\t" +
                re.sub("[\t\n]", "", biosample["environmental package"]) + "\t" +
                re.sub("[\t\n]", "", biosample["environment"]) +"\t"+
                re.sub("[\t\n]","",self.dic["salinity"]) +"\t"+
                re.sub("[\t\n]","",self.dic["Phenotypes"]) +"\n")


def main():
    parser = argparse.ArgumentParser(description="A script that can extract information of interest from NCBI BioSample database")
    parser.add_argument('--input', help="eg, biosample_set.xml", required=True)
    parser.add_argument('--output', help="eg, biosample_extract.tsv", required=True)
    argv = parser.parse_args()

    # 创建一个 XMLReader
    parser = xml.sax.make_parser()
    # turn off namespaces
    parser.setFeature(xml.sax.handler.feature_namespaces, 0)
    # 重写 ContextHandler
    Handler = BiosampleHandler(argv.output)
    # 设置当前的ContentHandler为我们自己写的handler实例
    parser.setContentHandler(Handler)
    # 开始解析 xml文件。
    parser.parse(argv.input)

    # 并行处理每个BioSample
    with ThreadPoolExecutor() as executor:
        futures = [executor.submit(Handler.process_biosample, biosample) for biosample in Handler.bio_samples]
        
        # 等待所有任务完成
        for future in as_completed(futures):
            future.result()


if (__name__ == "__main__"):
    begin_time = time()
    main()
    end_time = time()
    print("RUN TIME:", end_time - begin_time)
