#!/usr/bin/env python

#######################################################################
#######################################################################
## Created on January 23rd 2019 to get autosomes from samtools FAI file
#######################################################################
#######################################################################

import os
import errno
import argparse

############################################
############################################
## PARSE ARGUMENTS
############################################
############################################

Description = "Get a list of autosomes from FAI file for assembly and write to file."
Epilog = """Example usage: python get_autosomes.py <FAI_FILE> <OUT_FILE>"""

argParser = argparse.ArgumentParser(description=Description, epilog=Epilog)

## REQUIRED PARAMETERS
argParser.add_argument("FAI_FILE", help="FAI input file.")
argParser.add_argument("OUT_FILE", help="Output file containing one chromosome per line.")
args = argParser.parse_args()

############################################
############################################
## HELPER FUNCTIONS
############################################
############################################


def makedir(path):
    if not len(path) == 0:
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise


############################################
############################################
## MAIN FUNCTION
############################################
############################################

## TESTED WITH IGENOMES ASSEMBLIES FOR:
# Arabidopsis_thaliana
# Caenorhabditis_elegans
# Danio_rerio
# Drosophila_melanogaster
# Enterobacteriophage_lambda
# Escherichia_coli_K_12_DH10B
# Escherichia_coli_K_12_MG1655
# Gallus_gallus
# Homo_sapiens
# Mus_musculus
# Mycobacterium_tuberculosis_H37RV
# PhiX
# Rattus_norvegicus
# Saccharomyces_cerevisiae
# Schizosaccharomyces_pombe
# Staphylococcus_aureus_NCTC_8325


def get_autosomes(FAIFile, OutFile):
    makedir(os.path.dirname(OutFile))

    ## READ IN CHROMOSOME IDS
    chrList = []
    fin = open(FAIFile, "r")
    while True:
        line = fin.readline()
        if line:
            chrList.append(line.strip().split("\t")[0])
        else:
            fin.close()
            break

    ## REMOVE EXACT MATCHES TO MITOCHONDRIAL, SEX AND OTHER NON-AUTOSOMAL CHROMOSOMES
    exactMatch = ["y", "w", "z", "m", "mt", "mtdna", "pt", "mtr", "2-micron", "ebv"]
    if (
        "VI" not in chrList and "chrVI" not in chrList
    ):  ## Caenorhabditis elegans has roman numerals and chrX but only 5 chromosomes!
        exactMatch += ["x"]
    filteredList = [x for x in chrList if x.lower() not in exactMatch and x.lower().lstrip("chr") not in exactMatch]

    ## REMOVE INEXACT MATCHS TO SMALLER/RANDOM HUMAN CONTIGS
    fuzzyMatch = ["chrUn", "random", "v1", "v2", "kn707", "jtfh"]
    filteredList = [x for x in filteredList if not any(y in x.lower() for y in fuzzyMatch)]

    ## WRITE TO FILE
    fout = open(OutFile, "w")
    for chrom in filteredList:
        fout.write("%s\n" % (chrom))
    fout.close()


############################################
############################################
## RUN FUNCTION
############################################
############################################

get_autosomes(FAIFile=args.FAI_FILE, OutFile=args.OUT_FILE)

############################################
############################################
############################################
############################################
