#!/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 removes all singleton clusters.
As an extra option, it can also relabel the taxids.
'''

import sys
import argparse

def process_celems( clidline, elements, options ):
    if len(elements)<2:
        return 0
    print(clidline)
    if options.prefix is None:
        for e in elements:
            print(e)
    else:
        for e in elements:
            print(e.replace(' >',' >{}_'.format(options.prefix)))

    return 1
    
def process_clstr( options ):
    """
    Read the cluster file output from cd-hit

    >Cluster 728
    0       307aa, >74969_1:000311... *
    1       307aa, >74969_1:00063c... at 1:307:1:307/97.39%

    """
    try:
        inpf = open(options.clstr,"r")
    except IOError:
        print(( 'ERROR: file open failed for %s' % (options.clstr) ))
        return 255

    #    print '# read clstr file: '+inpfile
    # cluster id
    clidline=""
    
    # list of lines in the current cluster
    elements=[]

    # loop over all lines in input file
    # expects a line like '1       59nt, >Hsap:71763444... at 1:58:28:85/+/98%'
    n=0
    for line in inpf:
        # cleanup line
        line = line.strip().rstrip('\n')
        # skip empty or commented lines
        if (len(line)==0) or (line[0]=='#'):
            print(line)
            continue
        
        if line.startswith('>Cluster'):
            if len(clidline)==0:
                clidline = line
            else:
                n += process_celems( clidline, elements, options )
                clidline = line
            elements = []
        else:
            elements.append(line)
            
    if len(elements)>1:
        n += process_celems( clidline, elements, options )

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

parser = argparse.ArgumentParser(description='Clean cd-hit clstr file')

parser.add_argument('-c', '--clstr',  dest='clstr',  help='cd-hit cluster file', metavar='FILE', required=True)
parser.add_argument('-p', '--prefix', dest='prefix', help='prefix to be added to the taxids', default=None, required=False)

options = parser.parse_args()

# open fasta file and read data
process_clstr( options )

            
