#!/usr/bin/python3

# orthologer - Pipeline for deducing orthology from fasta files.
# Copyright (C) 2005-2023 E. Zdobnov lab: F. Tegenfeldt
# <fredrik.tegenfeldt@unige.ch> based on code by E. Kriventseva
#
# This file is part of orthologer.
#
# orthologer is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orthologer is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orthologer.  If not, see <http://www.gnu.org/licenses/>.

'''
This is generic tool for getting basic statistics of columnwise data.
Pipeline usage: STATS
'''
import os
import argparse

def getElem(elems,col):
    if col < len(elems):
        return elems[col]
    return None

class Element:
    def __init__(self,label,taxid,assembly):
        self.label    = label
        self.taxid    = taxid
        self.assembly = assembly
        
    def __gt__(self,other):
        return self.assembly > other.assembly
    
def process( options ):
    with open(options.table) as fh:
        allElems={}
        relabel = '' if options.norelabel else '+'
        first = True
        for line in fh.readlines():
            if first and options.ignoreHeader:
                first = False
                continue
            # ignore empty or commented lines (starting with '#')
            line = line.strip().rstrip('\n')
            if len(line)==0 or line[0] in '#':
                continue
            
            elems = line.split(options.delim)
            
            label    = getElem(elems,options.label)
            if label is None:
                continue
            taxid    = getElem(elems,options.taxid)
            assembly = getElem(elems,options.assembly)
            
            e = Element(label,taxid,assembly)
            if e.taxid not in allElems:
                allElems[e.taxid] = []
            allElems[e.taxid].append(e)                

        prefix = relabel
        for x,elems in allElems.items():

            if len(elems) > 1:
                elems = sorted(elems,reverse=True)
            if options.ignoreDuplicates or options.noversion:
                elems = elems[:1]
            count=0
            for e in elems:
                f = os.path.join(options.fastaDir,e.label)
                if options.checkFiles:
                    if not os.path.exists(f):
                        prefix = '#'+relabel
                    else:
                        prefix = relabel
                if options.noversion:
                    print("{}{} {}".format(prefix,e.taxid,f))
                else:
                    print("{}{}_{} {}".format(prefix,e.taxid,count,f))
                count += 1
            
parser = argparse.ArgumentParser(description='Generate fasta import file based on a table')

parser.add_argument("-t", "--table", dest="table",
                    help="file with labels, taxids and assembly id (or more)", metavar="FILE", required=True)
parser.add_argument("-f", "--dir",    dest="fastaDir",   default='', help="directory where to find the files" )

parser.add_argument("-l", "--label",    dest="label",    type=int, default=0, help="column for label" )
parser.add_argument("-x", "--taxid",    dest="taxid",    type=int, default=8, help="column for taxid" )
parser.add_argument("-a", "--assembly", dest="assembly", type=int, default=9, help="column for assembly id" )
parser.add_argument("-d", "--delim",    dest="delim", default=',', help="delimiter" )

parser.add_argument("--checkFiles", dest="checkFiles", action="store_true",
                    help="check files and comment out those not found (prefix with '#')")

parser.add_argument("--norelabel", dest="norelabel", action="store_true",
                    help="do not add '+' before taxid in file")

parser.add_argument("--ignoreHeader", dest="ignoreHeader", action="store_true",
                    help="ignore first row - assumes its a header")

parser.add_argument("-i", "--ignoreDuplicates", dest="ignoreDuplicates", action="store_true",
                    help="ignore duplicates, take the one with the latest assembly id (by sorting)")

parser.add_argument("-n", "--noversion", dest="noversion", action="store_true",
                    help="with -i, exclude version extension of taxid (do not append '_0')")

options = parser.parse_args()

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

process( options )
