#!/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
        
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 mapping requires from pipeline data the following:
    # > Rawdata
    #   - *.fs.maptxt            mapping between of fasta labels
    # > Sequences
    #   - *.fs_selected          final processed fasta
    #   - *.fs_selected.mm*|dmnd DB files (mmseqs or diamond)
    #   - *.fs_selected.seqlen   sequence lengths
    #   - *.fs_selected.stats    fasta stats
    #   - *.fs_selected.clstr    cdhit clusters
    # > PWC
    #   - PWC/A/A_A.align        diagonal of alignment matrix
    # > Cluster
    #   - Cluster/*.og           final cluster results, internal id's
    #   - Cluster/*.spcov        cluster coverage per tax id
    # > todo
    #   - todo/<root>.todo
    #
    files = []
    missing = []
    # todo file
    files,missing = add_file( options, 'todo/{}.todo'.format(rootNode), 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 )

    for L in leaves:
        # 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(L), files, missing )
        for S in [ 'seqlen', 'stats', 'clstr' ]:
            files,missing = add_file( options, 'Sequences/{}.fs_selected.{}'.format(L,S), files, missing )

        # DB files
        # 1. mmseqs
        fDB = 'Sequences/{}.fs_selected.mm'.format(L)
        if os.path.isfile(fDB):
            files,missing = add_file( options, fDB, files, missing )
            for S in [ '.dbtype', '_h', '_h.dbtype', '_h.index', '.index', '.lookup', '.source' ]:
                files,missing = add_file( options, '{}{}'.format(fDB,L,S), files, missing )

        # 2. diamond
        fDB = 'Sequences/{}.fs_selected.dmnd'.format(L)
        if os.path.isfile(fDB):
            files,missing = add_file( options, fDB, files, missing )

        # PWC files - the diagonal
        files,missing = add_file( options, 'PWC/{0}/{0}_{0}.align'.format(L), files, missing )
        
    return files, missing

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate runs 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)

    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)

    # for each node, 
    for node in nodes:
        # get all nodes below current node
        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
        options.maxLevel = 1
        options.setroot = node
        leaves = GR.process(options)
        files,missing = make_file_list( options, node, leaves )
        tarFile = output_file_base( args, node )+'.tgz'
        childrenFile = output_file_base( args, node )+'.children'
        print('creating {} with {} files'.format(tarFile,len(files)))
        with open(childrenFile,'w') as fh:
            if len(subnodes)==0:
                fh.write('#EMPTY\n')
            else:
                for c in subnodes:
                    fh.write('{}\n'.format(c))
        
        with tarfile.open(tarFile,'w:gz') as fh:
            for f in files:
                fh.add(f[1],arcname=f[0])
            
        
