#!/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 tool reads a fasta file and inserts a unique ID for each sequence.
 A line like:
   >ENSRNOP00000063300 pep:known ...
 will become
   >MYRAT:00001 ENSRNOP00000063300 pep:known ...

 The unique ID has the format <LABEL>:<GENE ID>
 where <LABEL> is given in the command line.
 The gene id is a hexadecimal number, increased by 1 for each gene
 and zero-padded to a given total length.

 The result is output to the standard output.

 The mapping between the simplified id and real id is saved in a file.

 Pipeline usage: import of fasta files
'''

from __future__ import print_function

import fasta_tools as FT

import argparse

label2taxid={}
class LabelMap:
    def __init__(self):
        self.name2id={}
        self.id2name=[]

    def getId(self,name):
        if name in self.name2id:
            return self.name2id[name]
        else:
            return None

    def getName(self,idnum):
        if len(self.id)>idnum:
            return self.id2name[idnum]
        else:
            return None

    def addNew(self,name):
        newid=None
        if name not in self.name2id:
            newid=len(self.id2name)
            self.name2id[name]=newid
            self.id2name.append(name)
        return newid


def readTaxid2Label( fname ):
    try:
        fsh  = open(fname,mode='r')
    except:
        raise Exception('given label file does not exist %s'%fname)

    for rawline in fsh.readlines():
        line = rawline.strip().rstrip('\n')
        if len(line)==0: continue
        if line[0] in '#': continue
        elems=line.split(';')
        if len(elems)<2: continue
        taxid=elems[0]
        label=elems[1]
        if label in label2taxid:
            print('# error: duplicate label to taxid mapping : %s'%label)
            continue
        label2taxid[label]=taxid



def processFasta( fname, fmap, deflabel, padding, allinfo ):
    fsh  = open(fname,mode='r')
    if fmap is None:
        fmap='%s.maptxt'%fname

    mfsh = open(fmap,mode='w')

    # label map
    labelMap = LabelMap()

    # nr of duplicates
    ndup = 0

    # scan fasta file
    done=False
    nskip=0
    while not done:
        prevpos = fsh.tell()
        line = fsh.readline()
        if len(line)==0:
            done=True
            continue
        # print comment line
        if line[0]=='#':
            print(line.strip())
            continue

        # seek back to beginning of line
        fsh.seek(prevpos)
        
        # get next fasta
        (done,orglabel,orgdescr,sequence)=FT.getNext(fsh)
        if sequence is None:
            nskip=nskip+1
            continue
        if orglabel is None:
            continue
        
        # give original label and it will return a number
        newid=labelMap.addNew(orglabel)
        # if none, it's a duplicate - skip it
        if newid is None:
            nskip=nskip+1
            continue


        # check what label to use
        uselabel=None
        if orglabel in label2taxid:
            uselabel=label2taxid[orglabel]
        else:
            uselabel=deflabel

        # if no label found -> skip!
        if uselabel is None:
            nskip = nskip+1
            continue

        # numerical ID, zero padded
        padfmt = ('%%0%dx') % (padding)
        gidstr = (padfmt) % (newid)

        extras=orglabel
        if allinfo and orgdescr is not None:
            extras = extras+' '+orgdescr

        # dont print extras in output fasta file
        newline = '>'+uselabel+':'+gidstr
        print(newline)
        print(sequence)
        # print extras in mapping
        mfsh.write(uselabel+':'+gidstr+' '+extras+'\n')

#-----------------------------------------------------------------
# MAIN



parser = argparse.ArgumentParser(description='Prepare fasta file')

#parser.add_argument('-m', '--map',   dest='map',       help='gene map output file - a Python pickle file with a bi-directional map (LabelMap)',    metavar='FILE')
parser.add_argument('-f', '--fasta',   dest='filename', help='fasta input file',       metavar='FILE', required=True)
parser.add_argument('-l', '--label',   dest='deflabel', help='default taxid label to be used',    metavar='string', required=False, default=None)
parser.add_argument('-L', '--flbl',    dest='flbl',     help='map taxid labels to be used for a given gene',    metavar='string', required=False, default=None)
parser.add_argument('-p', '--padding', dest='padding',  help='zero padding in gene number id',  type=int, metavar='int', default=6)
parser.add_argument('-m', '--map',     dest='mapfile',  help='output map file', metavar='FILE', default=None, required=False)
parser.add_argument('-a', '--all',     dest='allinfo',  help='add all info from the original gene id line to map', action='store_true', required=False)


options = parser.parse_args()

# open fasta file and read data
if options.flbl is not None:
    readTaxid2Label(options.flbl)

processFasta( options.filename, options.mapfile, options.deflabel, options.padding, options.allinfo )
