#!/usr/bin/python3
"""
Package pipeline results for mapping using the tree of nodes.
For each node, only the files specific for the node are saved.
In addition, a file containing all children nodes is produced.
The idea is that reading one node will trigger a read of all
the children.
"""
import os
import sys
import argparse
import tarfile

def add_file( options, f, files, missing ):
    full = os.path.join( options.source, f )
    if not os.path.isfile(full):
        missing += [ full ]
        print('error: missing file {}'.format(full))
        return files,missing
    arcname = os.path.join('{}'.format(options.arcname),f) if options.arcname is not None else f
    files.append( [arcname,full] )
    return files,missing

def create_package( options ):
    # The fasta related files have two different labels:
    # 1. odb taxid    : <odbtxid> : plain taxid <taxid>
    # 2. subsample id : <sstxid>  : g<node>_<subsample label>_<taxid>
    #
    # Note the subsample id is only applied if --sspref is set.
    # Otherwise only id 1 is applied.
    #
    # The mapping requires from pipeline data the following:
    # > Rawdata
    #   - <odbtxid>.fs.maptxt            mapping between of fasta labels
    # > Sequences
    #   - <sstxid>.fs_selected          final processed fasta
    #   - <sstxid>.fs_selected.mm*      mmseqs DB files
    #   - <sstxid>.fs_selected.dmnd     diamond DB files
    #   - <sstxid>.fs_selected.seqlen   sequence lengths
    #   - <sstxid>.fs_selected.stats    fasta stats
    # > PWC
    #   - PWC/A/A_A.align        diagonal of alignment matrix, A = <sstxid>
    # > Cluster
    #   - Cluster/<root>.og           final cluster results, internal id's
    #   - Cluster/<root>.spcov        cluster coverage per tax id
    # > todo
    #   - todo/<root>.todo
    #   - todo/<ssroot>.todo
    #
    files = []
    missing = []
    ssnode='g{}_{}'.format(options.node,options.ssprefix)
    sstodo='todo/{}.todo'.format(ssnode)
    orgtodo='todo/{}.todo'.format(options.node)

    srctodo = orgtodo
    if options.srctodo is not None:
        srctodo = options.srctodo
        
    # read subsampled todo
    print('read labels at {}',srctodo)
    todoLabels=[]
    with open(srctodo,'r') as fh:
        for line in fh:
            L = line.rstrip()
            if (L[0] == '#') or (len(L) == 0):
                continue
            todoLabels.append(L)

    if len(todoLabels) == 0:
        print('error: found no labels in file {}'.format(orgtodo))
        return 255

    print('add todo files')
    # root todo file
    files,missing = add_file( options, orgtodo, files, missing )
    
    # subsampled todo file
    files,missing = add_file( options, sstodo, files, missing )

    print('add cluster files')
    # cluster file
    files,missing = add_file( options, 'Cluster/{}.og'.format(options.node), files, missing )

    # cluster coverage
    files,missing = add_file( options, 'Cluster/{}.spcov'.format(options.node), files, missing )

    print('add files per species')
    # loop over all species
    for label in todoLabels:

        ssid = '{}_{}'.format(ssnode,label)
        print('---> {}'.format(ssid))
        # gene map files
        files,missing = add_file( options, 'Rawdata/{}.fs.maptxt'.format(label), files, missing )

        # sequence files
        files,missing = add_file( options, 'Sequences/{}.fs_selected'.format(ssid), files, missing )
        for S in [ 'seqlen', 'stats' ]:
            files,missing = add_file( options, 'Sequences/{}.fs_selected.{}'.format(ssid,S), files, missing )

        ### inpar type 9 with ( really needed???)
        ### files,missing = add_file( options, 'Sequences/{}.inpar_selected'.format(L), files, missing )
        
        # diamond DB
        missdb=[]
        files,missdb = add_file( options, 'Sequences/{}.fs_selected.dmnd'.format(ssid), files, missdb )
        
        # mmseqs DB files
        if len(missdb) != 0:
            missdb=[]
            files,missdb = add_file( options, 'Sequences/{}.fs_selected.mm'.format(ssid), files, missdb )
            for S in [ '.dbtype', '_h', '_h.dbtype', '_h.index', '.index', '.lookup', '.source' ]:
                files,missdb = add_file( options, 'Sequences/{}.fs_selected.mm{}'.format(ssid,S), files, missdb )

        if len(missdb) != 0:
            missing.extend( missdb )
        
        files,missing = add_file( options, 'PWC/{0}/{0}_{0}.align'.format(ssid), files, missing )

    rval=0
    # create tar file
    if len(files)>0:
        prefix  = '{}_'.format(options.prefix)  if options.prefix  is not None else ''
        postfix = '_{}'.format(options.postfix) if options.postfix is not None else ''
        base = '{}{}_mapping{}'.format( prefix, options.node, postfix )
        tarFile = os.path.join(options.output,base)
        print('creating tar file {} with {} files'.format(tarFile,len(files)))
        with tarfile.open(tarFile,'w:gz') as fh:
            for f in files:
                fh.add(f[1],arcname=f[0])
    else:
        print('error: no files found, missing files {}'.format(len(missing)))
        rval=1

    return rval

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Package OrthoDB mapping data from a subsample todo file')

    parser.add_argument('-n', '--node',    dest='node',  help='taxid of node to be package',   required=True)
    parser.add_argument('-s', '--source',  dest='source',  help='orthologer pipeline dir', required=True)
    parser.add_argument('-o', '--output',  dest='output',  help='packaged output dir',     required=True)
    parser.add_argument('-t', '--srctodo', dest='srctodo', help='labels to select (true odb labels)', required=False, default = None)
    parser.add_argument('-a', '--arcname', dest='arcname', help='root directory for archive',      required=False, default = None)
    parser.add_argument('-p', '--prefix',  dest='prefix',  help='output file prefix',      required=False, default = None)
    parser.add_argument('-x', '--postfix', dest='postfix', help='output file postfix',     required=False, default = None)
    parser.add_argument('-S', '--sspref',  dest='ssprefix',help='subsample fasta label prefix', required=False, default = 'n8ssc')
    

    options = parser.parse_args()

    # extract root node
    rval = create_package( options )
    
    sys.exit(rval)
        
