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

def load_selected(fsel):
    selected={}
    ispair=None
    with open(fsel) as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            # skip empty line
            if len(line)==0:
                continue

            # skip commented or empty lines
            if len(line)==0 or line[0] in '#':
                continue

            # split line in elements
            elems = line.split()

            # check if pairs or single
            if ispair is None:
                ispair = len(elems)==2
            if ispair:
                A=elems[0]
                B=elems[1]
                if A not in selected:
                    selected[A]=[]
                selected[A].append(B)
                if B not in selected:
                    selected[B]=[]
                selected[B].append(A)
            else:
                print(elems)
                selected[elems[0]]=True
    return ispair,selected

def check_if_selected(A,B,ispair,selected):
    if ispair:
        if A in selected:
            if B in selected[A]:
                return True
        return False
    else:
        if A in selected and B in selected:
            return True
    return False
            
def process(opts):
    ispair,selected = load_selected(opts.fselect)
    if ispair is None:
        return
    
    with open(opts.filename) as fh:
        for line in fh:
            line = line.strip().rstrip('\n')

            # print commented or empty lines
            if len(line)==0 or line[0] in '#':
                print(line)
                continue

            # split line in elements
            elems = line.split()

            # extract tx ids for A and B
            eA = elems[0].split(':')
            eB = elems[1].split(':')
            if check_if_selected(eA[0],eB[0],ispair,selected):
                print(line)

parser = argparse.ArgumentParser(description='split global alignment file according to taxids')
parser.add_argument('-f', '--file', dest='filename',
                    help='input pairwise data', metavar='FILE', required=True)
parser.add_argument('-s', '--select', dest='fselect',
                    help='file with selected taxids (single or double column)', metavar='FILE', required=True)



options = parser.parse_args()

process(options)
