#!/usr/bin/python3
'''
Remaps id's in cluster/fasta/pairwise files.
Pipeline usage: POSTPROC:REMAP
'''

import os
import argparse
import sys
import cluster_tools as CT
import fasta_tools as FT

def print_error(s):
    print(('ERROR   : %s')%s,file=sys.stderr)

def process_cluster( options, cluster ):
    if cluster.nelems == 0:
        print('skip cluster {} - empty'.format(cluster.clid))
        return 0

    print('create fasta file: {}'.format(options.output))
    
    with open(options.output,'w') as fhout:
    
        # loop over cluster elements
        # they are ordered by species id
        for spid,vcelems in cluster.data.items():

            # get all gene id labels in cluster
            geneids = set()
            for celem in vcelems:
                geneids.add( celem.key )
                ngeneids = len(geneids)
            if ngeneids == 0:
                continue

            # get source fasta file in Sequences dir
            ffasta = os.path.join( options.sequences, '{}.{}'.format(spid,options.fsext) )

            # open fasta for given species id and extract all gene ids
            nmissing = ngeneids
            with open( ffasta, 'r' ) as fh:
                done = False
                while not done:
                    # get next valid fasta
                    eof,label,description,sequence = FT.getNext( fh )

                    # check if label is in our set
                    if label in geneids:

                        # it's in our set - remove and print data
                        geneids.remove(label)
                        print('>{}\n{}'.format(label,sequence),file=fhout)
                        nmissing = nmissing-1

                    # check if done (end of file or gene set empty)
                    done =  eof or (nmissing==0)
            if nmissing>0:
                print('warning: missing {}/{} genes of species {} from cluster {}'.format(nmissing,ngeneids,spid,cluster.clid))

        print('created fasta file with {} gene ids'.format(ngeneids))
        return 0
    print('failed opening for output: {}'.format(options.output))
    return 255

def process( options ):
    speciesList=set()
    with open( options.cluster, mode='r' ) as fh:
        print('process cluster file {}'.format(options.cluster))
        
        # loop over cluster file
        done = False
        while not done:
            done,cluster = CT.get_next( fh, options.taxsep, CT.ClusterElementBase )
            if cluster.clid == options.clid:
                # cluster found -> process
                return process_cluster(options, cluster)
        
        print('failed to find cluster id {} in file {}'.format(options.clid,options.cluster))
        return 255
    
    print('unable to open cluster file {}'.format(options.cluster))
    return 255
    

#-----------------------------------------------------------------
# MAIN
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract fasta for a given cluster id in an orthologer environment.')

    parser.add_argument( '-c', '--cluster',   dest='cluster',      help='cluster file',                      metavar='FILE', required=True )
    parser.add_argument( '-C', '--clid',      dest='clid',         help='cluster id',                        metavar='clid', required=True )
    parser.add_argument( '-o', '--output',    dest='output',       help='fasta output file',                 metavar='dir',  required=True )
    parser.add_argument( '-x', '--sep',       dest='taxsep',       help='taxid separator',                   metavar='char', required=False, default = ':' )
    parser.add_argument( '-s', '--sequences', dest='sequences',    help='directory containing sequences',    metavar='dir',  required=False, default='Sequences' )
    parser.add_argument( '-e', '--fsext',     dest='fsext',        help='fasta file extension',              metavar='char', required=False, default = 'fs_masked' )
                        
    options = parser.parse_args()
    process( options )
