#!python
from argparse import ArgumentParser
import pandas as pd
from Bio import SeqIO
import importlib_resources


#get the default kraken2 output and the file to save translated names file to
parser = ArgumentParser()
parser.add_argument("-f", "--file", help="Kraken output file to be translated")
parser.add_argument("-c", "--contig_file", default=None, help="Contig file to filter")
parser.add_argument("-r", "--remove_viral", action='store_true', help="Remove viral contigs from fasta")
parser.add_argument("-l", "--length", default=250, help="Threshold to remove contigs from fasta")
parser.add_argument('-o', "--output", default='Viral_Contigs',help="Output CSV File Name")

args = parser.parse_args()


## Load Database from resources
my_resources = importlib_resources.files("virkraken")
data = (my_resources / "data" / "TaxID.csv")

## Read data
df2 = pd.read_csv(data, sep=",", low_memory=False)
df = pd.read_csv(args.file, sep="\t", header=None)

df = df[[0,1,2]]
df['Kingdom'] = [x.split("_")[1] for x in df[1]]

## Parse Metaspades output (Coverage/Contig_Length)
df['Contig_Length'] = [x.split("_")[3] for x in df[1]]
df['Coverage'] = [x.split("_")[5] for x in df[1]]
df.columns = ["Classified", "ContigID", "TaxID", "Kingdom", "Contig_Length", "Coverage"]


df_Classified = df[df["Classified"]=="C"]
q = df_Classified.merge(df2, left_on="TaxID", right_on='tax_id')
q2=q[q["superkingdom"]=="Viruses"]
q2 = q2[["ContigID", 'TaxID', "Contig_Length", "Coverage"]]
q2['Contig_Length'] = q2['Contig_Length'].astype(int)

## Filter by length
length = int(args.length)
q2 = q2[q2['Contig_Length']>=length]

print("Writing csv file")
outfile = args.output+".csv"
q2.to_csv(outfile, index=False)

contigs = list(q2['ContigID'])

## Filter fasta file by list of headers
def filter_fasta(fasta_file, list_of_headers, outputs):
    fasta_sequences = SeqIO.parse(open(fasta_file),'fasta')
    if args.remove_viral != False:
        with open(outputs, "w") as out_file:
            for fasta in fasta_sequences:
                if fasta.id in list_of_headers:
                    continue
                else:
                    SeqIO.write(fasta, out_file, "fasta")
    else:
        with open(outputs, "w") as out_file:
            for fasta in fasta_sequences:
                if fasta.id in list_of_headers:
                    SeqIO.write(fasta, out_file, "fasta")
                else:
                    continue
                    
                    
if args.contig_file != None:
    print("Writing filtered fasta")
    outfile = args.output+".fasta"
    filter_fasta(args.contig_file, contigs, outfile)

print("Done! Thanks for using virkraken!")
