#!/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 generate_runs as GR
import tarfile

class Options(GR.Options):
    def __init__(self, args):
        GR.Options.__init__(self)
        self.nwtree  = args.nwtree
        self.source  = args.source
        self.output  = args.output
        self.prefix  = args.prefix
        self.postfix = args.postfix
        self.arcname = args.arcname
        self.ssprefix = args.ssprefix

def output_file_base( options, rootNode ):
    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,rootNode,postfix)
    return os.path.join(options.output,base)

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 make_file_list( options, rootNode, leaves ):
    # 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 = []
    # root todo file
    files,missing = add_file( options, 'todo/{}.todo'.format(rootNode), files, missing )
    
    # subsampled todo file
    if options.ssprefix is not None:
        files,missing = add_file( options, 'todo/g{}_{}.todo'.format(rootNode,options.ssprefix), files, missing )
    
    # cluster file
    files,missing = add_file( options, 'Cluster/{}.og'.format(rootNode), files, missing )

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

    ssid_base = ''
    if options.ssprefix is not None:
        ssid_base = 'g{}_{}_'.format(rootNode, options.ssprefix)
    
    for L in leaves:
        ssid = '{}{}'.format(ssid_base,L)
        # gene map files
        files,missing = add_file( options, 'Rawdata/{}.fs.maptxt'.format(L), 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 )
        
    return files, missing

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Package OrthoDB mapping data from a newick tree')

    parser.add_argument('-t', '--nwtree',  dest='nwtree',  help='tree in newick format',   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('-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 = None)

    args = parser.parse_args()

    options = Options( args )

    # extract root node
    options.root = True
    output = GR.process(options)
    rootNode = output[0]

    # get all nodes
    options.root = False
    options.nodes = True
    options.nokids = True
    options.setroot = rootNode
    options.maxLevel = 0
    nodes = GR.process(options)

    create_children = True if options.ssprefix is None else False
    
    # for each node, 
    for node in nodes:
        # get all nodes below current node
        if create_children:
            options.nodes = True
            options.leaves = False
            options.maxLevel = 0
            options.setroot = node
            subnodes = GR.process(options)
            if node in subnodes:
                subnodes.remove(node)

        # get all UNIQUE leaves per node (maxLevel=1)
        options.nodes = False
        options.leaves = True
        if create_children:
            options.maxLevel = 1
        else:
            options.maxLevel = 0
        options.setroot = node
        leaves = GR.process(options)
        files,missing = make_file_list( options, node, leaves )

        if len(missing)>0:
            print('error: missing files for node {}!'.format(node))

        if create_children:
            childrenFile = output_file_base( args, node )+'.children'
            print('creating children file {} with {} nodes'.format(childrenFile,len(subnodes)))
            with open(childrenFile,'w') as fh:
                if len(subnodes)==0:
                    fh.write('#EMPTY\n')
                else:
                    for c in subnodes:
                        fh.write('{}\n'.format(c))

        tarFile = output_file_base( args, node )+'.tgz'
        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])
            
        
