#!/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):
        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=='one2one'
    
    def selectPair(self, iA,iB):
        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 ):
        P=[A,B]
        if not self.matrix_full():
            P.sort()
        return ' '.join(P)
    def getPairAA(self, iA,iB):
        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)

    def getPairAB(self, iA,iB):
        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)

    def getDiagPair(self, ind, useA):
        L=None
        if useA:
            if (self.labelsA is None):
                return None
            if (len(self.labelsA) <= ind):
                return None
            L = self.labelsA[ind]
        else:
            if (self.labelsB is None):
                return None
            if (len(self.labelsB) <= ind):
                return None
            L = self.labelsB[ind]
            
        return "{} {}".format(L,L)
            
    
#---------------------------------------------------------------
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 OLDfillPairData( options ):
    '''
    read todo and possibly todo2
    fill PairData
    '''
    pairdata = PairData( options )
    pairdata.labelsA, pairdata.labelsMapA = readTodo( options.todo )
    if options.todo2 is not None:
        pairdata.labelsB, pairdata.labelsMapB = readTodo( options.todo2 )
        # two todo files, one for A and one for B - set matrix to full
        if ( pairdata.matrix_nodiag() ):
            pairdata.options.matrix = 'full'

    if pairdata.options.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 fillPairData( options ):
    '''
    read todo and possibly todo2
    fill PairData
    '''
    pairdata = PairData( options )
    setA = readTodoSet( options.todo )
    if setA is None:
        print("error: no data in file {}".format(options.todo))
        return None

    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'

    if len(options.children)>0:
        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

    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()
    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:
                        pairsAA.add(p)

        else:
            # pairs: (A1,B1), (A1,B2), ... , (A2,B1), (A2,B2), ...
            for indA in range(pairData.rangeA[0],pairData.rangeA[1]+1):
                for indB in range(pairData.rangeB[0],pairData.rangeB[1]+1):
                    if pairData.isRequiredA(indA) or pairData.isRequiredB(indB):
                        p = pairData.getPairAB(indA,indB)
                        if p is not None:
                            pairsAB.add(p)

            # loop over A if dopwcA
            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)

            if not pairData.matrix_nodiag():
                print('# A diagonal')
                # get the Ax,Ax pairs
                for ind in range(pairData.rangeA[0],pairData.rangeA[1]+1):
                    if pairData.isRequiredA(ind):
                        p = pairData.getDiagPair(ind,True)
                        if p is not None:
                            pairsDiag.add(p)

                # get the Bx,Bx pairs only if NO children (dopwcA is false)
                if not pairData.dopwcA:
                    print('# B diagonal')
                    for ind in range(pairData.rangeB[0],pairData.rangeB[1]+1):
                        if pairData.isRequiredB(ind):
                            p = pairData.getDiagPair(ind,False)
                            if p is not None:
                                pairsDiag.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(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 ):
    options.matrix = options.matrix.lower()[0]
    if options.matrix not in ['h','f','n','o']:
        printerr("error: matrix must be one of f(ull),h(alf),n(odiag),o(ne2one)")
        return None

    m = 'full'
    if options.matrix == 'h':
        m = 'half'
    elif options.matrix == 'n':
        m = 'nodiag'
    elif options.matrix == 'o':
        m = 'one2one'
    
    # return matrix string
    return m
    
def process(options):
    print('# Todo     : {}'.format(options.todo))
    print('# Todo2    : {}'.format(options.todo2))
    print('# Children : {}'.format(options.children))
    ###print('# One2one  : {}'.format(options.one2one))
    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))
    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:
 matrix       - pairwise matrix type (only with --todo)
 one2one      - if A and B sets are of same size, create one-to-one pairs

Possible input source combinations:
1. pairs      - use the pairs file           allowed: required,rangeA,rangeB,one2one
2. todo       - create pairs from todo file  allowed: all options
3. todo/todo2 - A:todo and B:todo2           allowed: required,rangeA,rangeB,one2one

'''
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('-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('-e', '--todo2',    dest='todo2',    help='second todo file', metavar='FILE',   required=False, default=None)
###parser.add_argument('-x', '--one2one',  dest='one2one',  help='if set, create pairs from todo/todo2 using a 1-to-1 correspondance', action='store_true', required=False, default=False)
parser.add_argument('-m', '--matrix',   dest='matrix',   help='full,half,nodiag,one2one - ignored when --pairs is given', metavar='MATRIX', required=False, default='half')
parser.add_argument('-c', '--child',    dest='children', help='child node(s) ; labels from --todo present in children will be removed ; labels in children not in --todo are ignored', nargs='+')

options = parser.parse_args()

options.matrix = parse_matrix( options )
if options.matrix is None:
    sys.exit(255)

if (options.pairs is None) and (options.todo is None):
    printerr("error: one of --todo or --pairs must be set")
    sys.exit(255)
    
if options.pairs is not None:
    options.matrix = 'full'
    if options.todo is not None:
        print('warning: ignoring option --todo since --pairs is given')
        options.todo = None
        
    if options.todo2 is not None:
        printerr('error: only one of options --pairs and --todo2 may be given')
        sys.exit(255)

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

if options.children is None:
    options.children = []
    
if len(options.children)>0:
    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)

# open fasta file and read data
sys.exit( process( options ) )

