#!/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/>.
'''
Generate cluster ID's according to a template
Pipeline usage: -
'''
from __future__ import print_function
#
# This tool reads a cluster file and creates a unique id based on user input.
#
# The mapping between the simplified id and real id is saved in a file.
#
import os
import sys
import time
import argparse
import string
import orthologer_tools as OT
#
# Convert integer to a nr with the base equal to 36 (digts+uppcase ascii letters)
# Limits are:
#  1 - 36
#  2 - 1296
#  3 - 46656
#  4 - 1679616
#  5 - 60466176
#
digs = string.digits+string.ascii_uppercase
def int2basemax(ndigs):
    base=len(digs)
    return pow(base,ndigs)

def int2base(x,pad=0):
    base=len(digs)
    if x < 0: sign = -1
    elif x == 0: return digs[0].zfill(pad)
    else: sign = 1
    x *= sign
    digits = []
    while x>0:
        digits.append(digs[x % base])
        x = int(x/base)
    if sign < 0:
        digits.append('-')
    digits.reverse()
    rval=''.join(digits)
    if pad > 0:
        return rval.zfill(pad)
    return rval

def uname():
    ud=os.uname()
    rval=ud.sysname+' '+ud.nodename+' '+ud.release+' '+ud.version+' '+ud.machine
    return rval

def process( fname, fmap, prefix, rootpad, nodepad, clidpad, rootid, nodeid, nodecnt ):
    if fname is None:
        OT.print_error('no filename given')
        return 255
    if fmap is None:
        OT.print_error('no map filename given')
        return 255

    rootidmax  = int2basemax(rootpad)
    nodecntmax = int2basemax(nodepad)
    clidmax    = int2basemax(clidpad)

    if nodecntmax<nodecnt:
        OT.print_error('node count (%d) is greater than max allowed (%d) for given padding (%d)'%(nodecnt,nodecntmax,nodepad))
        return 255
    if rootidmax<rootid:
        OT.print_error('root id (%d) is greater than max allowed (%d) for given padding (%d)'%(rootid,rootidmax,rootpad))
        return 255

    # open output
    try:
        fsh  = open(fname,mode='r')
    except:
        OT.print_error('unable to open for read the file (%s)'%(fname))
        return 255
    try:
        mfsh = open(fmap,mode='w')
    except:
        fsh.close()
        OT.print_error('unable to open for write the map file (%s)'%(fmap))
        return 255

    # write header to output cluster file
    print(('# Date                : %s')%(time.strftime('%c')))
    print(('# System              : %s')%(uname()))
    print(('# User                : %s')%(os.getlogin()))
    print(('# Input cluster file  : %s')%(fname))
    print(('# Associated map file : %s')%(fmap))
    
    # write header to map file
    mfsh.write(('# Date                : %s\n')%(time.strftime('%c')))
    mfsh.write(('# System              : %s\n')%(uname()))
    mfsh.write(('# User                : %s\n')%(os.getlogin()))
    mfsh.write(('# Input cluster file  : %s\n')%(fname))
    mfsh.write(('# %s %s %s\n')%('ODB','Node','Clid'))

    # scan cluster file
    lastclid=None
    idstr=None
    
    xnodecnt = int2base(nodecnt,nodepad)
    xrootid  = int2base(rootid,rootpad)

    xclid=''
    warned=False
    nwarn=0
    for rawline in fsh.readlines():
        line = rawline.strip().rstrip('\n')
        if len(line)==0: continue
        if line[0] in '#': continue
            
        # split the line
        elems=line.split()
        clid = int(elems[0]) # cluster id
        rest=elems[1:]  # rest of the line

        if (clid>clidmax):
            if not warned:
                print('#WARNING: max cluster id reached - the following ids will be unique but with more characters...')
            warned=True
            nwarn=nwarn+1
        # use nodecnt to encode the new id
        if idstr is None or (clid != lastclid):                
            xclid=int2base(clid,clidpad)
            idstr = ('%s%s%s%s') % (prefix,xrootid,xnodecnt,xclid)
            mfsh.write(('%s %d %d\n')%(idstr,nodeid,clid))

        print('%s %s'%(idstr,' '.join(rest)))

        lastclid=clid

    # files closing
    fsh.close()
    mfsh.close()
    if nwarn>0:
        print('# Warning - got %d longer ids'%nwarn)
        return 1
    return 0
    
#-----------------------------------------------------------------
# MAIN
#

parser = argparse.ArgumentParser(description='Create cluster ids - ids are created as <prefix><rootid><nodecnt><clid>')

parser.add_argument('-f', '--filename', dest='filename',  help='cluster input file',          metavar='FILE',   required=False, default=None)
parser.add_argument('-l', '--prefix',   dest='prefix',    help='prefix of the ids',           metavar='string', required=False, default=None)
parser.add_argument('-r', '--rootid',   dest='rootid',    help='root node id (NOT taxid, but a simple nr)',   metavar='int',    required=False, default=None, type=int)
parser.add_argument('-x', '--nodecnt',  dest='nodecnt',   help='node count - unique index used in the OG id', metavar='int',    required=False, default=None, type=int)
parser.add_argument('-n', '--nodeid',   dest='nodeid',    help='node taxon id - only used in the output map file',               metavar='int',    required=False, default=None, type=int)

parser.add_argument('-R', '--rootpad',  dest='rootpad',   help='zero padding in root node id, default=1',  type=int, metavar='int', default=2)
parser.add_argument('-N', '--nodepad',  dest='nodepad',   help='zero padding in node count, default=2',  type=int, metavar='int', default=2)
parser.add_argument('-C', '--clidpad',  dest='clidpad',   help='zero padding in cluster id, default=4',  type=int, metavar='int', default=4)
parser.add_argument('-m', '--map',      dest='mapfile',   help='output map file (default <filename>.maptxt)', metavar='FILE', default=None, required=False)
parser.add_argument('--showrange',   dest='showrange',    help='print range of allowed ids',  action='store_true', default=False)


options = parser.parse_args()

if options.showrange:
    rootcntmax=int2basemax(options.rootpad)
    nodecntmax=int2basemax(options.nodepad)
    clidmax=int2basemax(options.clidpad)
    print('Root node padding  = %d'%options.rootpad)
    print('Node padding       = %d'%options.nodepad)
    print('Cluster id padding = %d'%options.clidpad)
    print('Max node count     = %d'%nodecntmax)
    print('Max cluster id     = %d'%clidmax)
    sys.exit(0)
    
if options.filename is None:
    print('error: missing --filename')
    sys.exit(255)
        
if options.prefix is None:
    print('error: missing --prefix')
    sys.exit(255)

if options.rootid is None:
    print('error: missing --rootid')
    sys.exit(255)

if options.nodeid is None:
    print('error: missing --nodeid')
    sys.exit(255)
        
if options.nodecnt is None:
    print('error: missing --nodecnt')
    sys.exit(255)
    
fmap=options.mapfile
if fmap is None:
    fmap='%s.maptxt'%options.filename
    

rval=process( options.filename, fmap, options.prefix, options.rootpad, options.nodepad, options.clidpad, options.rootid, options.nodeid, options.nodecnt )
if rval>1: sys.exit(rval)
sys.exit(0)
