#!/usr/bin/python3
import sys
import argparse
import io

import cluster_tools as CT
"""
Subsample fasta files based on species cluster coverage.

"""

def process_cluster( cluster, spcov, N, byHighScore, doStep ):
    cltype = 0
    # ignore type 2,4,5 and inparalogs (7,9)
    # check if any element is a triangle
    if 0 not in cluster.cltypes:
        # no triangle, check if any element is a chain
        if 3 in cluster.cltypes:
            cltype = 3
        else:
            # nope, its a pair
            cltype = 1

    class CelemData:
        def __init__(self,clid,celem,facc):
            self.clid  = clid
            self.facc  = facc
            self.celem = celem

        # the below makes sure that whatever the direction of the sort
        # the keys with lowest facc come first
        # only the sorting within a fixed facc is sorted according to direction
        def __lt__(self,other):
            if self.facc == other.facc:
                return self.celem.score < other.celem.score
            return self.facc < other.facc
        
        def __gt__(self,other):
            if self.facc == other.facc:
                return self.celem.score > other.celem.score
            return self.facc < other.facc
            
        def __le__(self,other):
            if self.facc == other.facc:
                return self.celem.score <= other.celem.score
            return self.facc < other.facc
        
        def __ge__(self,other):
            if self.facc == other.facc:
                return self.celem.score >= other.celem.score
            return self.facc < other.facc

            
    # find all celem with cltype
    celems=[]
    for sp,data in cluster.data.items():
        # sp   == species
        # data == all cluster elements containing this species
        facc = spcov.get( sp, 0 )
            
        for ce in data:
            if ce.cltype == cltype:
                celems.append(CelemData(cluster.clid, ce,facc))
            
    # sort according to score while keeping the species priority order (given by spcov)
    celems = sorted( celems, reverse = byHighScore )
    allsp = set()
    result=[]
    
    # retain only best element per species
    celemsBest = []
        
    for c in celems:
        #print('TTT: {} {} {}'.format(c.celem.key, c.facc, c.celem.score))
        if c.celem.species in allsp:
            continue
        allsp.add(c.celem.species)
        celemsBest.append(c)
        
    # establish step size -select from the top half
    Ncelems = len(celemsBest)/2
    step = 1
    if doStep:
        step = 1 if Ncelems <= N else int(Ncelems/N)
    
    # loop over elements, note any species is only represented once
    istep = 0
    for c in celemsBest:
        if (istep%step == 0):
            result.append(c)
            if len(result)==N:
                break
        istep += 1
        
    return result
        
def readSpeciesCoverage( options ):
    if options.spcov is None:
        return {}

    spcov={}
    with open( options.spcov, mode='r' ) as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            # remove empty/commented lines
            if len(line)==0 or line[0]=='#':
                continue
            elems = line.split()
            if len(elems)<9:
                continue
            
            if not elems[1].isnumeric():
                continue

            facc = float(elems[8])
            species = elems[0]
            spcov[species] = facc

    return spcov
                
def readClusterFile( options ):
    '''
    Reads clusters from input file.
    '''
    spcov = readSpeciesCoverage( options )
    
    nelemsTot=0
    nelemsSel=0
    speciesList=set()
    speciesAdded=set()
    with open( options.cluster, mode='r' ) as fh:
        print('# process file      {}'.format(options.cluster))
        print('# spcov file        {}'.format(options.spcov))
        print('# N(subsample)      {}'.format(options.Ngenes))
        print('# Select high score {}'.format(options.highscore))
        print('# Select by step    {}'.format(options.dostep))

        done = False
        while not done:
            done,cluster = CT.get_next( fh, options.taxsep )
            if cluster.nelems == 0:
                continue
            
            # process cluster
            genes = process_cluster(cluster, spcov, options.Ngenes, options.highscore, options.dostep)
            if options.details:
                for g in genes:
                    speciesAdded.add( g.celem.species )
                    print('{} {} {} {}'.format(cluster.clid, g.celem.key, g.facc, g.celem.score))
            else:
                for g in genes:
                    speciesAdded.add( g.celem.species )
                    print(g.celem.key)
            #print(*genes,sep='\n')
            nelemsSel += len(genes)
            nelemsTot += cluster.nelems
            
    print('#N(selected) {}'.format(nelemsSel))
    print('#N(total)    {}'.format(nelemsTot))
            

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Simple plot tool')

    parser.add_argument('-o', '--og', dest='cluster',
                        help='cluster input file', metavar='FILE', required=True )
    parser.add_argument('-c', '--spcov', dest='spcov',
                        help='cluster coverage file', metavar='FILE', default=None, required=False )
    parser.add_argument('-n', '--ngenes', dest='Ngenes',
                        help='cluster input file', type = int, default=8 )
    parser.add_argument('-s', '--sep', dest='taxsep',
                        help='species separator', default = ':' )
    parser.add_argument('-r', '--highscore', dest='highscore',
                        help='select from highest scores instead of lowest', action = 'store_true', default = False )
    parser.add_argument('-d', '--dostep', dest='dostep',
                        help='sample by steps', action = 'store_true', default = False )
    parser.add_argument('--details', dest='details',
                        help='print details in output', action = 'store_true', default = False )
    
    options = parser.parse_args()
    if options.Ngenes<1:
        print("error: --ngenes must be a non-zero positive number",file=sys.stderr)
    else:
        readClusterFile( options )
