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

def printerr( v ):
    print( v, file=sys.stderr )

class PairData:
    def __init__(self, options, labelsA=[], labelsMapA={}, labelsB=None, labelsMapB=None, ):
        self.options    = options         # full options
        self.labelsA    = labelsA         # labels from todo
        self.labelsMapA = labelsMapA
        self.labelsB    = labelsB         # labels from todo2 if defined, otherwise same as A
        self.labelsMapB = labelsMapB
        self.rangeA = None
        self.rangeB = None
        self.fullA  = False
        self.fullB  = False
        self.requiredA = None
        self.requiredB = None
        self.dopwcA = False # flag if pairs are to be created for labelsA (when labelsB exists)
        if self.options.todo is not None:
            self.nameA = options.todo
        elif self.options.pairs is not None:
            self.nameA = options.pairs
            
    # return index of labels in csvLabels [ind1,ind2,...]
    # labels is empty -> returns empty list
    # csvLabels is empty -> returns [0,len(labels)-1]
    #
    def hasRequired(self):
        if self.requiredA is None:
            return False
        if self.requiredB is None:
            return False
        return len(self.requiredA)+len(self.requiredB)>0
    
    def getIndexes( self, labels, csvLabels, isrange ):
        if (labels is None) or (len(labels)==0):
            if isrange:
                return [],0,0
            else:
                return set(),0,0

        if csvLabels is None:
            if isrange:
                return [0,len(labels)-1],0,0
            else:
                return set(),0,0

        # rangeLabels is a CSV of labels
        elems = csvLabels.split(',')

        # it's a range given by first and last label
        if isrange:
            rval=[]
            if len(elems[0])==0:
                rval.append(0)
            else:
                rval.append( labels.index(elems[0]) )
            if len(elems)==1:
                rval.append(rval[0])
            else:
                if len(elems[1])==0:
                    rval.append(len(labels)-1)
                else:
                    rval.append( labels.index(elems[1]) )
            # make sure it's a valid range
            if rval[0]>rval[1]:
                raise ValueError('csvLabels expects first label to be BEFORE the second - {}'.format(csvLabels))
            return rval,2,2

        #----------------------------------------
        # it's a collection of indices
        rval = set()
        nok = 0
        for e in elems:
            ind = labels.index(e)
            rval.add(ind)
            nok += 1

        # return indices as a set()
        return rval,nok,len(elems)

    # will always return a list of size 2 with index [first,last]
    def getRange( self, labels, rangeLabels ):
        if labels is None:
            return None,False
        
        # no range -> return full range
        if rangeLabels is None:
            return [0,len(labels)-1],True

        # rangeLabels is a CSV of labels
        rval,nok,nelems = self.getIndexes( labels, rangeLabels, True )
        if (nok!=2):
            raise ValueError('Illegal label range {}'.format(rangeLabels))

        # check if full range
        if (rval[1] == len(labels)-1) and (rval[0] == 0):
            return rval,True

        # not full range
        return rval,False

    # will return an empty set of no required labels
    def getRequired( self, labels, requiredLabels ):
        if labels is None:
            return set(),0,0
        
        if requiredLabels is None:
            return set(),0,0
        # rangeLabels is a CSV of labels
        return self.getIndexes( labels, requiredLabels, False )

    def fillLabelsMap( self, labels ):
        if labels is None:
            return None
        index=0
        labelsMap = {}
        for L in labels:
            labelsMap[L] = index
            index += 1
        return labelsMap
    
    def fill(self):
        # maps label ID to index in list (labelsA/labelsB)
        self.labelsMapA = self.fillLabelsMap( self.labelsA )
        self.labelsMapB = self.fillLabelsMap( self.labelsB )
        
        self.rangeA, self.fullA   = self.getRange(self.labelsA, self.options.rangeA )
        self.requiredA,nok,nelems = self.getRequired(self.labelsA, self.options.required)
        if self.labelsB is None:
            self.rangeB, self.fullB   = self.getRange(self.labelsA, self.options.rangeB )
            self.requiredB            = self.requiredA
        else:
            self.rangeB, self.fullB   = self.getRange(self.labelsB, self.options.rangeB )
            self.requiredB,nok,nelems = self.getRequired(self.labelsB, self.options.required)

    def matrix_full(self):
        return self.options.matrix=='full'
    def matrix_half(self):
        return self.options.matrix=='half'
    def matrix_nodiag(self):
        return self.options.matrix=='nodiag'
    def matrix_one2one(self):
        return self.options.matrix.startswith('one2one')
    def matrix_one2many(self):
        return self.options.matrix.startswith('one2many')
    
    def selectPair(self, iA,iB):
        if self.options.keeporder:
            return True
        if self.matrix_full():
            return True
        if self.matrix_half():
            return iB>=iA
        return iB>iA

    def selectB0(self,iA):
        if self.matrix_full():
            return self.rangeB[0]
        if self.matrix_half():
            return iA
        return iA+1

    def isRequiredA( self, indA ):
        if (self.requiredA is None) or (len(self.requiredA)==0):
            return True
        return indA in self.requiredA

    def isRequiredB( self, indB ):
        if (self.requiredB is None) or (len(self.requiredB)==0):
            return True
        return indB in self.requiredB
    
    def indexPairOK(self, indA, indB ):
        if not self.fullA:
            if (indA < self.rangeA[0]) or (indA > self.rangeA[1]):
                return False

        if not self.fullB:
            if (indB < self.rangeB[0]) or (indB > self.rangeB[1]):
                return False

        if not (self.isRequiredA(indA) or self.isRequiredB(indB)):
            return False

        return True

    def labelPairOK(self, labelA, labelB ):
        indA = self.labelsMapA[labelA]
        indB = self.labelsMapA[labelB] if self.labelsMapB is None else self.labelsMapB[labelB]
        return self.indexPairOK(indA,indB)

    def makePair( self, A, B, swap = False ):
        P=[A,B]
        if not self.options.keeporder:
            if not self.matrix_full():
                P.sort()
        if swap:
            p = P[0]
            P[0] = P[1]
            P[1] = p
        
        return ' '.join(P)
    
    def getPairAA(self, iA,iB, swap = False):
        if (self.labelsA is None):
            return None
        if (len(self.labelsA) <= iA) or (len(self.labelsA) <= iB):
            return None
        A = self.labelsA[iA]
        B = self.labelsA[iB]
        return self.makePair(A,B,swap)

    def getPairAB(self, iA,iB, swap = False):
        if (self.labelsA is None):
            return None
        if (len(self.labelsA) <= iA) or (len(self.labelsB) <= iB):
            return None
        A = self.labelsA[iA]
        B = self.labelsB[iB]
        return self.makePair(A,B,swap)

    def getPairBB(self, iA,iB, swap = False):
        if (self.labelsB is None):
            return None
        if (len(self.labelsB) <= iA) or (len(self.labelsB) <= iB):
            return None
        A = self.labelsB[iA]
        B = self.labelsB[iB]
        return self.makePair(A,B,swap)
    
#---------------------------------------------------------------
def readTodoSet( ftodo ):
    '''
    read a single todo file and return a list of all labels
    '''
    if ftodo is None:
        return None

    labels=set()
    with open(ftodo,mode='r') as fh:
        for line in fh:
            line = line.strip()
            # skip empty of commented lines
            if (len(line) == 0) or (line[0] == '#'):
                continue

            labels.add(line)

        if len(labels)==0:
            labels = None

    return labels

def fillPairData( options ):
    '''
    read todo and possibly todo2 or children
    fill PairData
    '''
    pairdata = PairData( options )
    if options.pairs is not None:
        return pairdata
    
    setA = readTodoSet( options.todo )
    if setA is None:
        print("error: no data in file {}".format(options.todo))
        return None

    nA = len(setA)
    listB = None
    # options todo2 is never set if children are non-zero
    if options.todo2 is not None:
        setB = readTodoSet( options.todo2 )
        if setB is not None:
            listB = list(setB)
        
        # two todo files, one for A and one for B - set matrix to full
        #if ( pairdata.options.matrix == 'nodiag' ):
        #    pairdata.options.matrix = 'full'

    # possibly read children
    if len(options.children)>0:
        # A: remove all labels in children not in todo file
        # B: all labels from children
        setB = set()
        for c in pairdata.options.children:
            setC = readTodoSet( c )
            setA = setA.difference( setC )
            setB = setB.union( setC )

        listB = list(setB)
        
        # flag that pairs using only labelsA should be created
        pairdata.dopwcA = True

    nB = 0 if listB is None else len(listB)
    print('# Ninput(A) = {}'.format(nA))
    print('# Ninput(B) = {}'.format(nB))
    print('# Nfinal(A) = {}'.format(len(setA)))
    pairdata.labelsA = list(setA)
    pairdata.labelsA.sort()
    if listB is not None:
        pairdata.labelsB = listB
        pairdata.labelsB.sort()
    if pairdata.labelsA is None:
        print("error: no data in labelsA {}".format(options.todo))
        return None

    if pairdata.matrix_one2one():
        if len(pairdata.labelsA) != len(pairdata.labelsB):
            printerr("error: one2one option is set but label sets A and B are not of same size.")
            printerr("error: check options --todo and --todo2")
            return None
        
    pairdata.fill()
    return pairdata

def printPairsFromFile( pairData ):
    '''
    read a single pair file and return a list of all pairs
    '''
    if pairData.options.pairs is None:
        return None
    nok=None
    
    with open(pairData.options.pairs,mode='r') as fh:
        nok = 0
        for line in fh:
            line = line.strip()
            # skip empty of commented lines
            if (len(line) == 0) or (line[0] == '#'):
                continue
            elems = line.split()
            
            # make sure labels are valid
            if not pairData.labelPairOK(elems[0],elems[1]):
                continue
            
            # print pair
            print(line)
            nok += 1

    return nok


def printPairsFromTodo( pairData ):
    pairsAA = set()
    pairsAB = set()
    pairsBB = set()
    pairsDiag = set()
    nok = 0
    if pairData.labelsB is None:
        #--------------------------------------
        # create pairs from one todo
        #--------------------------------------
        pairsAA = set()
        for indA in range(pairData.rangeA[0],pairData.rangeA[1]+1):
            
            # select first B based on selected matrix
            B0 = pairData.selectB0(indA)

            # loop over B
            for indB in range(B0,pairData.rangeB[1]+1):
                if pairData.isRequiredA(indA) or pairData.isRequiredA(indB):
                    p = pairData.getPairAA(indA,indB)
                    if p is not None:
                        pairsAA.add(p)
    else:
        #--------------------------------------
        # separate label lists for A and B
        #--------------------------------------
        # TODO: save pairs, then depending on matrix, add diagonal, ie (A1,A1) etc
        if pairData.matrix_one2one():
            # A and B are the same size
            # pairs: (A1,B1), (A2,B2), ...
            for ind in range(pairData.rangeA[0],pairData.rangeA[1]+1):
                if pairData.isRequiredA(ind) or pairData.isRequiredB(ind):
                    # index are the same in A and B
                    # get A[ind] and B[ind]
                    p = pairData.getPairAB(ind,ind)
                    if p is not None:
                        pairsAB.add(p)
                    
                    if pairData.options.swapAB:
                        p = pairData.getPairAB(ind,ind,True)
                        if p is not None:
                            pairsAB.add(p)
                        
                    if pairData.options.diagA:
                        p = pairData.getPairAA(ind,ind)
                        if p is not None:
                            pairsDiag.add(p)
                            
                    if pairData.options.diagB:
                        p = pairData.getPairBB(ind,ind)
                        if p is not None:
                            pairsDiag.add(p)
                        
        elif pairData.matrix_one2many():
            # A and B are not necessarily the same size.
            # Generate all possible pairs where id's are sorted.
            # pairs: (A1,B1), (A2,B2), ...
            for indA in range(pairData.rangeA[0],pairData.rangeA[1]+1):
                if not pairData.isRequiredA(indA):
                    continue
                
                # add diagonal of A
                if pairData.options.diagA:
                    p = pairData.getPairAA(indA,indA)
                    if p is not None:
                        pairsAA.add(p)
                    
                for indB in range(pairData.rangeB[0],pairData.rangeB[1]+1):
                    if not pairData.isRequiredB(indB):
                        continue 
                    p = pairData.getPairAB(indA,indB)
                    if p is not None:
                        pairsAB.add(p)
                    
                    if pairData.options.swapAB:
                        p = pairData.getPairAB(indA,indB,True)
                        if p is not None:
                            pairsAB.add(p)

                    # add diagonal of B
                    if pairData.options.diagB:
                        p = pairData.getPairBB(indB,indB)
                        if p is not None:
                            pairsBB.add(p)

            # loop over A if dopwcA (has children)
            if pairData.dopwcA:
                for indA in range(pairData.rangeA[0],pairData.rangeA[1]+1):
            
                    # select first B based on selected matrix
                    B0 = pairData.selectB0(indA)

                    # loop over B
                    for indB in range(B0,pairData.rangeB[1]+1):
                        if pairData.isRequiredA(indA) or pairData.isRequiredA(indB):
                            p = pairData.getPairAA(indA,indB)
                            if p is not None:
                                pairsAA.add(p)


    # print all pairs
    if len(pairsAA)>0:
        print("#A vs A")
        for p in sorted(pairsAA):
            if (p not in pairsAB) and (p not in pairsDiag):
                print(p)
                nok += 1

    if len(pairsAB)>0:
        print("#A vs B")
        for p in sorted(pairsAB):
            if p not in pairsDiag:
                print(p)
                nok += 1

    if len(pairsBB)>0:
        print("#B vs B")
        for p in sorted(pairsBB):
            if p not in pairsDiag:
                print(p)
                nok += 1

    if len(pairsDiag)>0:
        print("# Diagonal")
        for p in sorted(pairsDiag):
            print(p)
            nok += 1

    return nok

        
    
def printPairs( pairData ):
    npairs = printPairsFromFile( pairData )
    if npairs is None:
        npairs = printPairsFromTodo( pairData )
    return npairs
    
    
def parse_matrix( options ):
    optmat = options.matrix.lower()
    mat = optmat[0]
    if mat not in ['h','f','n','o']:
        printerr("error: matrix must be one of f(ull),h(alf),n(odiag),o(ne2one), o(ne2many)")
        return None

    m = 'full'
    if mat == 'h':
        m = 'half'
    elif mat == 'n':
        m = 'nodiag'
    elif mat == 'o':
        if optmat == 'one2many':
            m = optmat
        else:
            m = 'one2one'
    
    # return matrix string
    return m
    
def process(options):
    if options.printHeader:
        print('# Todo       : {}'.format(options.todo))
        print('# Todo2      : {}'.format(options.todo2))
        print('# Children   : {}'.format(options.children))
        print('# Pairs      : {}'.format(options.pairs))
        print('# Required   : {}'.format(options.required))
        print('# Range A    : {}'.format(options.rangeA))
        print('# Range B    : {}'.format(options.rangeB))
        print('# Matrix     : {}'.format(options.matrix))
        print('# Keep order : {}'.format(options.keeporder))

    pairData = fillPairData( options )
    if pairData is None:
        return 255
    
    n = printPairs( pairData )
    print('# Nr of pairs = {}'.format(n))
    return 0

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


'''
This code generates a file with pairs of labels given the input.
Source:
 todo         - file containing labels
 todo2        - file containing labels, used as B
 pairs        - file containing pairs

Filter:
 required     - CSV list of required labels, if set at least one of a pair
               is reuired to be from this list
 rangeA       - range for first label in a pair, format: <first,last>
 rangeB       - same as rangeA but for second label

Matrix:
 * one todo file:
  full     : all-against-all pairs -> (A1,A1),(A1,A2),... (A2,A1)
  half     : like 'full' but not (A2,A1)
             order of labels in pair are sorted
  nodiag   : like 'half' but without diagonal pairs (A1,A1),...
  one2one  : N/A
  one2many : N/A
  diagA    : N/A
  diagB    : N/A

 * two todo files
  full     : N/A
  half     : N/A
  nodiag   : N/A
  one2one  : the two files must have the same nr of labels - (A1,B1), (A2,B2), ...
  one2many : all from A against all from B - (A1,B1), (A1,B2), ...
             order of labels in pair are sorted
  diagA    : add diagonal from A
  diagB    : add diagonal from B

 * one todo files + children nodes
  matrix -> one2many
  diagA  -> always set

 Note that two todo files and with matrix 'one2many', 'diagA' and 'diagB' is equivalent
 to using 'half' on a concatenation of the two files.

Possible input source combinations:
1. pairs         - use the pairs file
2. todo          - create pairs from todo file
3. todo/todo2    - A:todo ; B:todo2
4. todo+children - A:todo minus children ; B:children labels

'''
parser = argparse.ArgumentParser(description='Generate pairwise file')

parser.add_argument('-t', '--todo',     dest='todo',      help='todo file',        metavar='FILE', required=False, default=None)
parser.add_argument('-e', '--todo2',    dest='todo2',     help='second todo file', metavar='FILE',   required=False, default=None)
parser.add_argument('-p', '--pairs',    dest='pairs',     help='pairs file',       metavar='FILE', required=False, default=None)
parser.add_argument('-r', '--required', dest='required',  help='CSV list of required labels', metavar='CSV', required=False, default=None)
parser.add_argument('-A', '--rangeA',   dest='rangeA',    help='range for A, comma separated, if one is omitted the extreme value is assumed', metavar='RANGE', required=False, default=None)
parser.add_argument('-B', '--rangeB',   dest='rangeB',    help='range for B, comma separated, if one is omitted the extreme value is assumed', metavar='RANGE', required=False, default=None)

parser.add_argument('-m', '--matrix',   dest='matrix',    help='full,half,nodiag,one2one,one2many - ignored when --pairs is given', metavar='MATRIX', required=False, default='half')
parser.add_argument('-c', '--children', dest='children',  help='child node(s) ; labels from --todo present in children will be removed ; labels in children not in --todo are ignored', nargs='+')
parser.add_argument('-k', '--keeporder',dest='keeporder', help='keep order in pair; used when todo2 is defined',  action='store_true', default=False )
parser.add_argument('--diagA',dest='diagA', help='add diagonal pairs from A (--todo)  ; with --todo2',  action='store_true', default=False )
parser.add_argument('--diagB',dest='diagB', help='add diagonal pairs from B (--todo2) ; with --todo2',  action='store_true', default=False )
parser.add_argument('--swapAB',dest='swapAB', help='add the swapped pair (one2one or one2many) ; with --todo2',  action='store_true', default=False )
parser.add_argument('--header',dest='printHeader', help='print header',  action='store_true', default=False )

options = parser.parse_args()

# parse matrix type
options.matrix = parse_matrix( options )

# make sure matrix option is valid
if options.matrix is None:
    sys.exit(255)

# check if children is not set
if options.children is None:
    options.children = []

# make sure required options are set
if (options.pairs is None) and (options.todo is None):
    printerr("error: one of --todo or --pairs must be set")
    sys.exit(255)

# if pairs file is set
if options.pairs is not None:
    options.matrix = 'full'
    print('# info: ignoring any matrix setting since --pairs is given')
    if options.todo is not None:
        printerr('error: options --todo and --pairs are not allowed together')
        options.todo = None
        
    if options.todo2 is not None:
        printerr('error: options --todo2 and --pairs are not allowed together')
        options.todo2 = None

hasLabelB = False
# if todo2 is set
if options.todo2 is not None:
    hasLabelB = True
    if options.todo is None:
        printerr('error: option --todo must be set if --todo2 is set')
        sys.exit(255)

# if children is set ; option is only valid when only --todo is set
if len(options.children)>0:
    hasLabelB = True
    if options.todo2 is not None:
        printerr('error: option --children and --todo2 are mutually exclusif')
        sys.exit(255)
    if options.pairs is not None:
        printerr('error: option --children and --pairs are mutually exclusif')
        sys.exit(255)

# check matrix type if one or two todo files
if hasLabelB:
    # B labels are explicitely defined
    if not ( options.matrix == 'one2one' or options.matrix == 'one2many' ):
        printerr('error: illegal matrix type ({}) for given setup (two or more label sources)'.format(options.matrix))
        printerr('error: must be either one2one or one2many')
        sys.exit(255)
else:
    if options.matrix == 'one2one' or options.matrix == 'one2many':
        printerr('error: illegal matrix type ({}) for a single todo file'.format(options.matrix))
        printerr('error: must one of full, half or nodiag')
        sys.exit(255)
        
# open fasta file and read data
sys.exit( process( options ) )

