#/usr/bin/env python

import os
import argparse
import string
import re
from collections import defaultdict
from pbcore.io import DataSet, ContigSet, openDataSet, openDataFile
from pbcore.io.dataset import InvalidDataSetIOError
from pbcore.io.dataset.utils import _swapPath
from pbcore.io.dataset.DataSetMembers import Filters, OPMAP
from pbcore.io.dataset.DataSetValidator import validateFile
import logging

log = logging.getLogger(__name__)

def summarizeXml(args):
    dset = openDataSet(args.infile, strict=args.strict)

    # check to see if there was an error updating the dataset length:
    numFlag = ""
    if dset.numRecords == 0:
        dset.updateCounts()
        if not dset._countsUpdated:
            numFlag = " Unable to update counts!"
    print "DataSet Type          : {f}".format(f=dset.datasetType)
    print "Name                  : {f}".format(f=dset.name)
    print "Id                    : {f}".format(f=dset.uuid)
    print "Number of records     : {r}{f}".format(r=dset.numRecords,
                                                  f=numFlag)
    print "Total number of bases : {r}{f}".format(r=dset.totalLength,
                                                  f=numFlag)
    print "# of Resources        : {r}".format(r=len(dset.toExternalFiles()))
    print "Filters               : {r}".format(r=str(dset.filters) if
                                               dset.filters else "None")
    for fname in dset.toExternalFiles():
        print fname
    return 0

def summarize_options(parser):
    parser.description = "Summarize a DataSet XML file"
    parser.add_argument("infile", type=str,
                        help="The xml file to summarize")
    parser.set_defaults(func=summarizeXml)

def createXml(args):
    if args.dsType is None:
        dset = openDataFile(*args.infile, strict=args.strict,
                            skipCounts=args.skipCounts,
                            generateIndices=args.generateIndices)
    else:
        dsTypes = DataSet.castableTypes()
        dset = dsTypes[args.dsType](*args.infile, strict=args.strict,
                                    skipCounts=args.skipCounts,
                                    generateIndices=args.generateIndices)
    if args.generateIndices:
        # we generated the indices with the last open, lets capture them with
        # this one:
        dset = dsTypes[args.dsType](*args.infile, strict=args.strict,
                                    skipCounts=args.skipCounts)
    if args.dsName != '':
        dset.name = args.dsName
    if args.metadata:
        dset.loadMetadata(args.metadata)
    log.debug("Dataset created")
    dset.write(args.outfile, validate=args.novalidate, modPaths=True,
               relPaths=args.relative)
    log.debug("Dataset written")
    return 0


def create_options(parser):
    parser.description = ('Create an XML file from a fofn or bam. Possible '
                          'types: SubreadSet, AlignmentSet, ReferenceSet, '
                          'HdfSubreadSet, BarcodeSet, ConsensusAlignmentSet, '
                          'ConsensusReadSet, ContigSet')
    parser.add_argument("outfile", type=str, help="The XML to create")
    #parser.add_argument("infile", type=validate_file, nargs='+',
    parser.add_argument("infile", type=str, nargs='+',
                        help="The fofn or BAM file(s) to make into an XML")
    parser.add_argument("--type", type=str, default=None,
                        dest='dsType', help="The type of XML to create")
    parser.add_argument("--name", type=str, default='',
                        dest='dsName', help="The name of the new DataSet")
    parser.add_argument("--generateIndices", action='store_true',
                        default=False, help="Generate index files (.pbi and .bai for BAM, .fai for FASTA).  Requires samtools/pysam and pbindex.")
    parser.add_argument("--metadata", type=str, default=None,
                        help=("A metadata.xml file (or DataSet) to supply "
                              "metadata"))
    parser.add_argument("--novalidate", action='store_false', default=True,
                        help=("Don't validate the resulting XML, don't touch "
                              "paths"))
    parser.add_argument("--relative", action='store_true', default=False,
                        help=("Make the included paths relative instead of "
                              "absolute (not compatible with --novalidate)"))
    parser.set_defaults(func=createXml)

def pad_separators(base_set):
    """e.g. 'gt' will hit 'length', so we pad it to ' gt '"""
    padded = []
    for sep in base_set:
        new_sep = []
        if sep[0] in string.ascii_lowercase:
            new_sep.append(" ")
        new_sep.append(sep)
        if sep[-1] in string.ascii_lowercase:
            new_sep.append(" ")
        padded.append(''.join(new_sep))
    return padded

def parse_filter_list(filtStrs):
    filters = defaultdict(list)
    # pull them from the filter parser:
    separators = OPMAP.keys()
    # pad the ones that start and end with letters
    separators = pad_separators(separators)
    for filt in filtStrs:
        for sep in separators:
            if sep in filt:
                param, condition = filt.split(sep)
                condition = (sep.strip(), condition.strip())
                filters[param.strip()].append(condition)
                break
    return filters

def filterXml(args):
    if args.infile.endswith('xml'):
        dataSet = openDataSet(args.infile, strict=args.strict)
        filters = parse_filter_list(args.filters)
        dataSet.filters.addRequirement(**filters)
        dataSet.updateCounts()
        log.info("{i} filters added".format(i=len(filters)))
        dataSet.write(args.outfile)
    else:
        raise IOError("No files found/found to be compatible")
    return 0

def filter_options(parser):
    pbiFilterOptions = set(Filters()._pbiMappedVecAccMap().keys())
    bamFilterOptions = set(Filters()._bamAccMap.keys())
    parser.description = ('Add filters to an XML file. Suggested fields: '
                          '{f}. More expensive fields: {b}'.format(
        f=sorted(list(pbiFilterOptions)),
        b=sorted(list(bamFilterOptions - pbiFilterOptions))))
    #parser.add_argument("infile", type=validate_file,
    parser.add_argument("infile", type=str,
                        help="The xml file to filter")
    parser.add_argument("outfile", type=str, help="The resulting xml file")
    parser.add_argument("filters", type=str, nargs='+',
                        help="The values and thresholds to filter (e.g. "
                        "'rq>0.85')")
    parser.set_defaults(func=filterXml)

def splitXml(args):
    log.debug("Starting split")
    dataSet = openDataSet(args.infile, strict=args.strict)
    chunks = len(args.outfiles)
    if args.chunks:
        chunks = args.chunks
    if isinstance(dataSet, ContigSet):
        dss = dataSet.split(chunks)
    else:
        dss = dataSet.split(chunks=chunks,
                            ignoreSubDatasets=(not args.subdatasets),
                            contigs=args.contigs,
                            maxChunks=args.maxChunks,
                            breakContigs=args.breakContigs,
                            targetSize=args.targetSize,
                            zmws=args.zmws,
                            barcodes=args.barcodes,
                            byRecords=(not args.byRefLength),
                            updateCounts=(not args.noCounts))
    log.debug("Splitting into {i} chunks".format(i=len(dss)))
    infix = 'chunk{i}'
    nSuf = -2 if re.search(r".+\.\w+set\.xml", args.infile) else -1
    if not args.outfiles:
        if not args.outdir:
            args.outfiles = ['.'.join(args.infile.split('.')[:nSuf] +
                                      [infix.format(i=chNum)] +
                                      args.infile.split('.')[nSuf:])
                             for chNum in range(len(dss))]
        else:
            args.outfiles = ['.'.join(args.infile.split('.')[:nSuf] +
                                      [infix.format(i=chNum)] +
                                      args.infile.split('.')[nSuf:])
                             for chNum in range(len(dss))]
            args.outfiles = [os.path.join(args.outdir,
                                          os.path.basename(outfn))
                             for outfn in args.outfiles]
            num = len(dss)
            end = ''
            if num > 5:
                num = 5
                end = '...'
            log.debug("Emitting {f} {e}".format(
                f=', '.join(args.outfiles[:num]),
                e=end))
    log.debug("Finished splitting, now writing")
    for out_fn, dset in zip(args.outfiles, dss):
        dset.write(out_fn)
    log.debug("Done writing files")
    return 0

def split_options(parser):
    parser.description = "Split the dataset"
    parser.add_argument("infile", type=str,
                        help="The xml file to split")
    parser.add_argument("--contigs", default=False, action='store_true',
                        help="Split on contigs")
    parser.add_argument("--barcodes", default=False, action='store_true',
                        help="Split on barcodes")
    parser.add_argument("--zmws", default=False, action='store_true',
                        help="Split on zmws")
    parser.add_argument("--byRefLength", default=True, action='store_true',
                        help="Split contigs by contig length")
    parser.add_argument("--noCounts", default=False, action='store_true',
                        help="Update dataset counts after split")
    parser.add_argument("--chunks", default=0, type=int,
                        help="Split contigs into <chunks> total windows")
    parser.add_argument("--maxChunks", default=0, type=int,
                        help="Split contig list into at most <chunks> groups")
    parser.add_argument("--targetSize", default=5000, type=int,
                        help="Target number of records per chunk")
    parser.add_argument("--breakContigs", default=False, action='store_true',
                        help="Break contigs to get closer to maxCounts")
    parser.add_argument("--subdatasets", default=False, action='store_true',
                        help="Split on subdatasets")
    parser.add_argument("--outdir", default=False, type=str,
                        help="Specify an output directory")
    parser.add_argument("outfiles", nargs=argparse.REMAINDER,
                        type=str, help="The resulting xml files")
    parser.set_defaults(func=splitXml)

def mergeXml(args):
    dss = [openDataSet(infn, strict=args.strict) for infn in args.infiles]
    allds = reduce(lambda ds1, ds2: ds1 + ds2, dss)
    if not allds is None:
        allds.updateCounts()
        allds.write(args.outfile)
    else:
        raise InvalidDataSetIOError("Merge failed, likely due to "
                                    "conflicting Filters")
    return 0

def merge_options(parser):
    parser.description = 'Combine XML files'
    parser.add_argument("outfile", type=str,
                        help="The resulting XML file")
    #parser.add_argument("infiles", type=validate_file, nargs='+',
    parser.add_argument("infiles", type=str, nargs='+',
                        help="The XML files to merge")
    parser.set_defaults(func=mergeXml)


def relativizeXml(args):
    dss = openDataSet(args.infile, strict=args.strict)
    dss.write(args.infile, relPaths=True)
    return 0

def relativize_options(parser):
    # doesn't make sense to have an outdir, that would screw with the relative
    # paths...
    parser.description = 'Make the paths in an XML file relative'
    parser.add_argument("infile", type=str,
                        help="The XML file to relativize")
    parser.set_defaults(func=relativizeXml)

def absolutizeXml(args):
    dss = openDataSet(args.infile, strict=args.strict)
    outfn = args.infile
    if args.outdir:
        if os.path.isdir(args.outdir):
            outfn = _swapPath(args.outdir, args.infile)
        else:
            outfn = args.outdir
    dss.write(outfn, relPaths=False)
    return 0

def absolutize_options(parser):
    parser.description = 'Make the paths in an XML file absolute'
    parser.add_argument("infile", type=str,
                        help="The XML file to absolutize")
    parser.add_argument("--outdir", default=None, type=str,
                        help="Specify an optional output directory")
    parser.set_defaults(func=absolutizeXml)

def copyToXml(args):
    dss = openDataSet(args.infile, strict=args.strict)
    outfn = args.outdir
    if os.path.isdir(args.outdir):
        outfn = _swapPath(args.outdir, args.infile)
    dss.copyTo(os.path.split(outfn)[0])
    dss.write(outfn, relPaths=args.relative)
    return 0

def copyTo_options(parser):
    parser.description = 'Copy a dataset and resources to a new location'
    parser.add_argument("infile", type=str,
                        help="The XML file to copy")
    parser.add_argument("outdir", type=str,
                        help="The copy to directory")
    parser.add_argument("--relative", action='store_true', default=False,
                        help=("Make the included paths relative instead of "
                              "absolute"))
    parser.set_defaults(func=copyToXml)

def newUuidXml(args):
    dss = openDataSet(args.infile, strict=args.strict)
    dss.newUuid(random=args.random)
    dss.write(args.infile)
    return 0

def newUniqueId_options(parser):
    parser.description = "Refresh a DataSet's UniqueId"
    parser.add_argument("infile", type=str,
                        help="The XML file to refresh")
    parser.add_argument("--random", action='store_true', default=False,
                        help=("Generate a random UUID, instead of a hash"))
    parser.set_defaults(func=newUuidXml)

def loadStatsXml(args):
    dset = openDataSet(args.infile, strict=args.strict)
    if len(dset.externalResources) > 1:
        log.info("More than one ExternalResource found, adding the "
                 "sts.xml nested external resource to the first one")
    dset.externalResources[0].sts = args.statsfile
    if args.outfile:
        dset.write(args.outfile, validate=False)
    else:
        dset.write(args.infile, validate=False)
    return 0

def loadStatsXml_options(parser):
    parser.description = 'Load an sts.xml file into a DataSet XML file'
    parser.add_argument("infile", type=str,
                        help="The XML file to modify")
    parser.add_argument("statsfile", type=str,
                        help="The .sts.xml file to load")
    parser.add_argument("--outfile", type=str, default=None,
                        help="The XML file to output")
    parser.set_defaults(func=loadStatsXml)

def loadMetadataXml(args):
    dset = openDataSet(args.infile, strict=args.strict)
    dset.loadMetadata(args.metadata)
    if args.outfile:
        dset.write(args.outfile, validate=False)
    else:
        dset.write(args.infile, validate=False)
    return 0

def loadMetadataXml_options(parser):
    parser.description = 'Load an .metadata.xml file into a DataSet XML file'
    parser.add_argument("infile", type=str,
                        help="The XML file to modify")
    parser.add_argument("metadata", type=str,
                        help=("The .metadata.xml file to load (or DataSet "
                              "to borrow from)"))
    parser.add_argument("--outfile", type=str, default=None,
                        help="The XML file to output")
    parser.set_defaults(func=loadMetadataXml)

def validateXml(args):
    validateFile(args.infile, args.skipFiles)
    print("{f} is valid DataSet XML with valid ResourceId "
          "references".format(f=args.infile))
    return 0

def validate_options(parser):
    parser.description = 'Validate XML and ResourceId files'
    parser.add_argument("infile", type=str,
                        help="The XML file to validate")
    parser.add_argument("--skipFiles",
                        default=False, action='store_true',
                        help="Skip validating external resources")
    parser.set_defaults(func=validateXml)

def consolidateXml(args):
    """Combine BAMs and apply the filters described in the XML file, producing
    one consolidated XML"""
    dset = openDataSet(args.infile)
    dset.consolidate(args.datafile, numFiles=args.numFiles, useTmp=(not
                     args.noTmp))
    dset.write(args.xmlfile)
    return 0

def consolidate_options(parser):
    parser.description = 'Consolidate the XML files'
    #parser.add_argument("infile", type=validate_file,
    parser.add_argument("--numFiles", type=int, default=1,
                        help="The number of data files to produce (1)")
    parser.add_argument("--noTmp", default=False, action='store_true',
                        help="Don't copy to a tmp location to ensure local"
                             " disk use")
    parser.add_argument("infile", type=str,
                        help="The XML file to consolidate")
    parser.add_argument("datafile", type=str,
                        help="The resulting data file")
    parser.add_argument("xmlfile", type=str,
                        help="The resulting XML file")
    parser.set_defaults(func=consolidateXml)
