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

#===============================================================================
# PROGRAM METADATA
#===============================================================================
__author__ = 'xinwei <xwei97@zju.edu.cn>'
__copyright__ = ''
__license__ = ''
__date__ = 'Tue Mar 23 14:30:10 2021'
__version__ = '0.1'

USAGE = \
"""
    %prog [options] -i [path of input.kreport] -o [path of output.kreport]
    
    Description: 
        Order by the rank of classification, put the rows with same rank together.

    Dependency: 
        pandas, os

    Examples:

    AUTHOR: Xin Wei <xwei97@zju.edu.cn>
    SINCE: Tue Mar 23 14:30:10 2021
    VERSION: 0.0.0

    Kraken2 output format:
    1. Percentage of fragments covered by the clade rooted at this taxon
    2. Number of fragments covered by the clade rooted at this taxon
    3. Number of fragments assigned directly to this taxon
    4. A rank code, indicating (U)nclassified, (R)oot, (D)omain, (K)ingdom, (P)hylum, (C)lass, (O)rder, (F)amily, (G)enus, or (S)pecies. 
    Taxa that are not at any of these 10 ranks have a rank code that is formed by using the rank code of the closest ancestor rank with a number indicating the distance from that rank. 
    E.g., "G2" is a rank code indicating a taxon is between genus and species and the grandparent taxon is at the genus rank.
    5. NCBI taxonomic ID number
    6. Indented scientific name
"""

#===============================================================================
# IMPORT STATEMENTS
#===============================================================================
import pandas as pd
import os
import numpy as np
import re
import sys
import time
from optparse import OptionParser, OptionGroup

#===============================================================================
# MAIN
#===============================================================================
def main(argv = sys.argv[1:]):
    parser = OptionParser( usage=USAGE, version = "Version: " + __version__)
	# Compulsory arguments
    compOptions = OptionGroup(parser, "Compulsory parameters", "There options are compulsory, and may be supplied in any order.")
    compOptions.add_option("-i", "--indir", type = "string", metavar = "DIR", help = "the directory of .kreport file generated by kraken2")
    compOptions.add_option("-o", "--outdir", type = "string", metavar = "DIR", help = "the directory of reordered .kreport file")
    parser.add_option_group(compOptions)
    try:
        options, x = parser.parse_args(argv)
    except:
        parser.error('[FATAL]: parsing failure.\n')
        exit(1)
    
    if not options.indir:
        parser.error('[FATAL]: input directory is not accessible.\n')
        exit(1)
    if not options.outdir:
        parser.error('[FATAL]: output directory is not accessible.\n')
        exit(1)
    inDir = options.indir
    outDir = options.outdir
    fileList = os.listdir(inDir)
    for i in fileList:
        if ('.report' in i) & ('_sorted.report' not in i):
             df = pd.read_csv(inDir + '/' + i, sep = '\t', header = None, names = ['percentage', 'reads_covered', 'reads_assigned', 'rank', 'taxid', 'name'])
             df['name'] = [s.strip() for s in df['name']]
             uniRank = set(df['rank'])             
             uniRank_list = [re.split(r'(\d+)', u, maxsplit=1) for u in uniRank]
             uni = list()
             for j in uniRank_list:
                 if len(j) == 1:
                     t = (j[0], 0)
                 if len(j) == 3:
                     t = (j[0], int(j[1]))
                 uni.append(t)
             d = {'U':1, 'R':2, 'S':3, 'G':4, 'F':5, 'O':6, 'C':7, 'P':8, 'K':9, 'D':10}
             uni.sort(key=lambda x:(d[x[0]], x[1]))
             custom_ord = []
             ten = []
             for x,y in uni:
                 if y == 0:
                     z = x
                     ten.append(z)
                 else:
                     z = x + str(y)
                     custom_ord.append(z)
             custom_ord = ten + custom_ord
                 
             df['rank'] = pd.Categorical(df['rank'], custom_ord)
             df.sort_values(by = ['rank', 'percentage', 'reads_assigned', 'reads_covered', 'taxid'], ascending=[True] + [False]*4, inplace = True)
             df.to_csv(outDir + '/' + re.sub(r'\.report$', '', i) + '_sorted.report', sep='\t', index=False, header=True)
        else:
            pass


if __name__ == '__main__':
    t0 = time.time()
    main()
    sys.stderr.write('Elapsed time to run reorder_kraken2_report.py: {} s\n'.format( (time.time()-t0) ) )

