#!/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/>.

import sys
import argparse
import os

def process(filename,trgdir,force):
    '''
    Input:
     filename          alignment file

    '''
    nline=0
    ndiag=0
    data={}
    txidset=set()
    with open(filename) as fh:
        for line in fh:
            nline=nline+1
            line = line.strip().rstrip('\n')
            # skip empty line
            if len(line)==0: continue

            # skip commented or empty lines
            if len(line)==0 or line[0] in '#':
                continue

            # split line in elements
            elems = line.split()

            #
            A = elems[0]
            B = elems[1]
            eA = A.split(':')
            eB = B.split(':')
            label = eA[0]+";"+eB[0]
            txidset.add(eA[0])
            txidset.add(eB[0])
            if label not in data:
                data[label]=[]
                if eA[0]==eB[0]:
                    ndiag=ndiag+1
            data[label].append(line)
    print('Nr of diagonal elements found: {}'.format(ndiag))
    print('Nr of taxids found           : {}'.format(len(txidset)))

    for label in data:
        txids = label.split(';')
        dirA  = trgdir+'/'+txids[0]
        if not os.path.exists(dirA):
            print('error: missing path %s'%dirA)
            continue
        fname = dirA +'/'+txids[0]+'_'+txids[1]+'.align'
        if os.path.exists(fname) and not force:
            print('warning: skip - file exists %s'%fname)
            continue
        with open(fname,'w') as fwh:
            print('writing file %s'%fname)
            prevLine=""
            for line in data[label]:
                # remove identical lines
                if line==prevLine:
                    continue
                print(line,file=fwh)
                prevLine=line
            print('#END',file=fwh)

parser = argparse.ArgumentParser(description='split global alignment file according to taxids')
parser.add_argument('-f', '--file', dest='filename',
                    help='read data from FILE', metavar='FILE', required=True)
parser.add_argument('-d', '--trgdir', dest='trgdir',
                    help='target dir', metavar='DIR', required=True)
parser.add_argument('--force',    dest='force',      help='overwrite existing target files', action='store_true', default=False)



options = parser.parse_args()

process(options.filename,options.trgdir,options.force)
