#!/usr/bin/env python
'''Concatenate (fastphase) outputs for single-file input to 
"Haplotype-based profiling of subtle allelic imbalance with SNP arrays"
   This file by Jerry Fowler
'''



EPILOG = '''Tell me what I should put here.
'''

CAT_VERSION = [1, 0]

import os
import sys

import argparse
from collections import OrderedDict
import logging
import random
import re
from io import StringIO
import subprocess

from labtools import misc
from labtools import const
from labtools.ArgumentParserValidator import ArgumentParserValidator
from labtools import timeutil
from labtools.version_tool import version_tool

from haploh import haploh

cat_version = version_tool(CAT_VERSION).version()

ACCEPTABLE_FASTPHASE = 'Output from fastPHASE'

VERSION = haploh.VERSION
HELP = haploh.HELP
LOGFILE = haploh.LOGFILE
LOGLEVEL = haploh.LOGLEVEL
PREFIX = haploh.PREFIX
DESTDIR = haploh.PERLCONST.DESTDIR
HAPID = haploh.PERLCONST.HAPID
CHROMOSOMES = 'chromosomes'
ALL_CHROMOSOMES = 'all'
CHROMOSOME_LIST = [str(chr+1) for chr in range(22)]

def parse(input=None, stdout=sys.stdout, stderr=sys.stderr):
    '''establish the parsing rules for the command line, and then
    parse and perform simple semantic tests on the parsed inputs.
    '''
    parser = ArgumentParserValidator(description=__doc__,
                                     formatter_class=argparse.RawTextHelpFormatter,
                                     add_help=False,
                                     epilog=EPILOG,
                                     prefix_chars=const.DASH+const.PLUS)

    # basic args
    parser.add_argument(const.DASHES+PREFIX, dest=PREFIX,
                        required=False, default=None,
                        help='common prefix for phasing, switch probabilities, and BAF files')
    parser.add_argument(const.DASHES+CHROMOSOMES, dest=CHROMOSOMES,
                        required=False, default=ALL_CHROMOSOMES,
                        help="list of chromosomes to look for ('%s' for all)" % (ALL_CHROMOSOMES))
    parser.add_argument(const.DASH+haploh.TYPE_PHASES, dest=haploh.TYPE_PHASES, default=True, action='store_false',
                       help='''don't process hapguess files''')
    parser.add_argument(const.DASH+haploh.TYPE_BAFS, dest=haploh.TYPE_BAFS, default=True, action='store_false',
                       help='''don't process baf files''')
    parser.add_argument(const.PLUS+haploh.TYPE_SWITCHPROBS, dest=haploh.TYPE_SWITCHPROBS, default=False, action='store_true',
                       help='''process switchprobs files''')
    parser.add_argument(const.DASHES+DESTDIR, dest=DESTDIR, default='CWD',
                        help='directory into which output files should be written, default .')
    parser.add_argument(const.DASHES+VERSION, dest=VERSION, default=False, action='store_true',
                       help='''Display version and quit''')
    parser.add_argument(const.DASHES+HELP, dest=HELP, default=0, action='count',
                       help='''Display usage and quit (--help --help for more comprehensive help)''')
    parser.add_argument(const.DASHES+LOGFILE, dest=LOGFILE, default=None,
                       help='''file in which log output will be written''')
#FIXME INFO
    parser.add_argument(const.DASHES+LOGLEVEL, dest=LOGLEVEL, default='DEBUG',
                       help='''log output level (INFO, DEBUG, default INFO)''')

    if len(input) > 1:
        args = parser.parse_args(input if input else None)
    else:
        args = argparse.Namespace()
        args.help = 1
    
    args.versionstring = 'cat-%s' % (cat_version(aslist=False))
    if args.version or args.help:
        print(args.versionstring, file=stdout)
        if args.help:
            print(const.EMPTY, file=stdout)
            parser.print_help(file=stdout)
        return None
    else:
        print('%s\n' % (args.versionstring), file=stderr)

    args.invocation = const.SPACE.join(['Invocation:'] + input)
    if args.destdir == 'CWD':
        print('Warning, placing output in current working directory', file=stderr)
        args.destdir = misc.CWD
    if args.is_valid_file(DESTDIR, accept_directory=True, test_writable=True):
        ok_to_log = True
        if not os.path.isdir(args.destdir):
            try:
                os.makedirs(args.destdir)
            except OSError as e:
                args.add_error('mkdir %s: %s' % (args.destdir, e.strerror))
                ok_to_log = False
    else:
        ok_to_log = False

    if ok_to_log:
        if not args.logfile:
            args.logfile = misc.suffixed_name(args.versionstring, 'log')
        args.lfile = args.logfile if os.path.dirname(args.logfile) else os.path.join(args.destdir, args.logfile)
        if args.loglevel in haploh.loglevel:
            llevel = haploh.loglevel[args.loglevel]
        else:
            llevel = logging.INFO
            args.loglevel = 'INFO'
        root = logging.getLogger()
        root.setLevel(llevel)
        handler = logging.FileHandler(filename=args.lfile, mode='w', delay=True)
        handler.setLevel(llevel)
        root.addHandler(handler)

    logging.info(args.versionstring)

    #logging.warn(args)
    args.types = []
    if args.phases: args.types.append(haploh.TYPE_PHASES)
    if args.bafs: args.types.append(haploh.TYPE_BAFS)
    if args.switchprobs: args.types.append(haploh.TYPE_SWITCHPROBS)
    no_output_list = [const.DASH+haploh.TYPE_PHASES, const.DASH+haploh.TYPE_BAFS]
    if not len(args.types):
        args.add_error('Combining %s specifies no output!' % (no_output_list))

    if args.chromosomes == ALL_CHROMOSOMES:
        args.chromosomes = CHROMOSOME_LIST
        allgood = True
    else:
        chromosomes = args.chromosomes.split(const.COMMA)
        allgood = min([(chr in CHROMOSOME_LIST) for chr in chromosomes])
        if not allgood:
            args.add_error("Invalid chromosome spec '%s' (must be comma-separated in the range 1 to 22)" %
                            (args.chromosomes))
        else:
            args.chromosomes = chromosomes

    if not args.prefix:
        args.add_error('%s is a required argument' % (const.DASHES+PREFIX))
    elif allgood:
        thusfar = len(args.validation_errors())
        for chr in args.chromosomes:
            for type in args.types:
                name = haploh.standard_input_name(args.prefix, type, chr)
                t = args.is_valid_filepath(name)
                #logging.debug('%s: %s found' % (name, 'not' if not t else ''))
        if len(args.validation_errors()) - thusfar > 0:
            report = (len(args.chromosomes)*len(args.types) +
                      thusfar - len(args.validation_errors()))
            del args.validation_errors()[thusfar:]
            args.add_error("Found %s files with prefix '%s' (Expected %d files including %s)" % 
                           (str(report) if report else 'no', args.prefix,
                            len(args.chromosomes)*len(args.types),
                            ', '.join([haploh.standard_input_name(os.path.basename(args.prefix),
                                                                  type, args.chromosomes[0])
                                              for type in args.types])))

    return args


def main(input=None, stdout=sys.stdout, stderr=sys.stderr):
    '''Parse input and run the driver, choosing an exit code based on how far we
    get successfully
    '''
    args = parse(input, stdout=stdout, stderr=stderr)
    if not args:
        # request for help or something
        return 0

    if args.validation_errors():
        args.print_errors(stream=stderr)
        return 1

    print('Logging at level %s will appear in' % (args.loglevel), args.lfile, file=stderr)
    logging.info(args.invocation)
    if args.loglevel not in loglevel:
        lmsg = "Loglevel '%s' not recognized, using INFO" % (args.loglevel)
        logging.WARN(lmsg)
        print(lmsg, file=stderr)

    rc = driver(args, stdout, stderr)
    if args.validation_errors():
        args.print_errors(stream=stderr)

    return rc


def extend_stream(thelist, thestream, thetype):
    pass

def driver(args, stdout=sys.stdout, stderr=sys.stderr):
    '''For each listed chromosome, for each suffix, build a single output for each
sample found in the phased input
    '''
    if max([(type in [haploh.TYPE_PHASES, haploh.TYPE_SWITCHPROBS]) for type in args.types]):
        process_phased(args)
    if (haploh.TYPE_BAFS in args.types):
        bafs = process_bafs(args, stdout, stderr)
        if not args.validation_errors():
            output_bafs(args, bafs)


ILLUMINA_BAF_KEY = 'B Allele Freq'
def baf_key(sample):
    return misc.suffixed_name(sample, ILLUMINA_BAF_KEY)
ILLUMINA_GTYPE_KEY = 'GType'
def gtype_key(sample):
    return misc.suffixed_name(sample, ILLUMINA_GTYPE_KEY)


class sample_baf():

    def __init__(self, name, bafindex, gtindex):
        self.name = name
        self.bafindex = bafindex
        self.gtindex = gtindex
        self.bafs = []
        self.genotypes = []
        logging.debug('Sample %s: %d, %d' % (name, bafindex, gtindex))

    def add_bafs(self, data):
        if len(data) < self.bafindex:
            logging.warn('FATAL: Sample %s (%s - %d): %d, %d' % (self.name, self.urlname(), len(self.bafs), self.bafindex, len(data)))
        if not len(self.bafs)%500:
            logging.debug('Sample %s (%s): %d, %d' % (self.name, self.urlname(), self.bafindex, len(data)))
        self.bafs.append(data[self.bafindex])

    def add_genotypes(self, data):
        self.genotypes.append(data[self.gtindex])
    
    def urlname(self):
        return misc_url_quote(self.name)


def extract_bafsamplenames(args, input, bafs):
    try:
        # the first column is an id field
        headerline = input.readline().strip().split(const.TAB)
        keys = OrderedDict()
        for key in [t.split(const.DOT)[0] for t in headerline if t.endswith(ILLUMINA_BAF_KEY)]:
            keys[key] = sample_baf(key, headerline.index(baf_key(key)), headerline.index(gtype_key(key)))
        if bafs:
            allfound = min([key in bafs for key in keys])
            nonenew = min([key in keys for key in bafs])
            if not (allfound and nonenew):
                args.add_error('bafname extract error: %s' % ('!@#$%%%%'))
                return None
            else:
                return bafs
        else:
            logging.info('Found samples %s' % (list(keys.keys())))
            if not len(keys):
                args.add_error('No bafnames found')
            return keys
    except Exception as e:
        args.add_error('bafname extract error: %s' % e.message)
        return None


def process_bafinput(args, bafs, input):
    line = input.readline()
    if line == const.EMPTY: return True
    newdata = line.strip().split()
    if not newdata: return False
    for sample in bafs:
        bafs[sample].add_bafs(newdata)
        bafs[sample].add_genotypes(newdata)

    return False


def process_bafs(args, stdout=sys.stdout, stderr=sys.stderr):
    bafs = None
    for chr in args.chromosomes:
        logging.info('Processing chromosome %s' % (chr))
        
        chrfile = haploh.standard_input_name(args.prefix, haploh.TYPE_BAFS, chr)
        with open(chrfile) as input:
            bafs = extract_bafsamplenames(args, input, bafs)
            if not bafs:
                return None
            if args.phases:
                phased_not_baf = [k for k in args.samples if k not in [bafs[b].urlname() for b in bafs]]
                baf_not_phased = [bafs[baf].urlname() for baf in bafs if bafs[baf].urlname() not in args.samples]
                if phased_not_baf:
                    print('Samples %r were found in the phased input but not the baf' % (phased_not_baf))
                if baf_not_phased:
                    print('Samples %r were found in the baf input but not the phased' % (baf_not_phased))

            done = process_bafinput(args, bafs, input)
            while not done:
                done = process_bafinput(args, bafs, input)
    return bafs


def output_bafs(args, bafs):
    for sample in list(bafs.keys()):
        outname = haploh.standard_intermediate_name(args.destdir, haploh.TYPE_BAFS, bafs[sample].urlname())
        with open(outname, 'w') as output:
            tab = const.EMPTY
            for baf in bafs[sample].bafs:
                output.write(tab+baf)
                tab = const.TAB
            output.write(const.NEWLINE)
            logging.info('Wrote %s' % (outname))


def output_genotypes(args, bafs):
    for sample in list(bafs.keys()):
        outname = os.path.join(args.destdir, 'sample-%s.genotypes' % (sample))
        with open(outname, 'w') as output:
            tab = const.EMPTY
            for gtype in bafs[sample].genotypes:
                output.write(tab+gtype)
                tab = const.TAB
            output.write(const.NEWLINE)
            logging.info('Wrote %s' % (outname))


def process_phased(args):
    sampledict = OrderedDict()
    for chr in args.chromosomes:
        logging.info('Processing chromosome %s' % (chr))
        
        chrfile = haploh.standard_input_name(args.prefix, haploh.TYPE_PHASES, chr)
        with open(chrfile) as input:
            isFastphase = ACCEPTABLE_FASTPHASE in input.read(200)
            if not isFastphase:
                args.add_error("Can currently concatenate only fastphase output, sorry.")
                return args

            process_phased_header(input)

            probname = haploh.standard_input_name(args.prefix, haploh.TYPE_SWITCHPROBS, chr)
            probfile = open(probname) if haploh.TYPE_SWITCHPROBS in args.types else None

            sample = extract_samplename(input)
            while sample is not None:
                alreadyStarted = sample in sampledict
                if not alreadyStarted:
                    logging.info('Adding sample %s' % (sample))
                    sampledict[sample] = OrderedDict()
                    sampledict[sample][haploh.TYPE_PHASES] = [const.EMPTY]*2

                spc = const.SPACE if alreadyStarted else const.EMPTY
                sampledict[sample][haploh.TYPE_PHASES][0] += (spc) + input.readline().strip()
                sampledict[sample][haploh.TYPE_PHASES][1] += (spc) + input.readline().strip()

                if probfile:
                    probline = probfile.readline().strip()
                    if alreadyStarted:
                        problist = probline.split(const.TAB)
                        problist[0] = '0.5'
                        sampledict[sample][haploh.TYPE_SWITCHPROBS][0] += const.TAB + const.TAB.join(problist)
                    else:
                        sampledict[sample][haploh.TYPE_SWITCHPROBS] = [probline]

                sample = extract_samplename(input)
                pass
            pass
        pass

    args.samples = list(sampledict.keys())
    for sample in list(sampledict.keys()):
        odict = sampledict[sample]
        for type in list(odict.keys()):
            data = odict[type]
            outfile = haploh.standard_intermediate_name(args.destdir, type, sample)
            with open(outfile, 'w') as output:
                for line in data:
                    output.write(line + const.NEWLINE)

                
FASTPHASE_START_TAG = 'BEGIN GENOTYPES'
FASTPHASE_END_TAG = 'END GENOTYPES'

def extract_samplename(phasestream):
    '''assume the next line of the phased data is a sample name
    and extract it
    '''
    sampleline = phasestream.readline().strip()
    if const.EMPTY == sampleline or sampleline.startswith(FASTPHASE_END_TAG):
        return None
    name = re.match(r'(#\s*)([^#]*).*$', sampleline).group(2)
    if not name:
        raise UserWarning("Found no sample name in line '%s'" % (sampleline))
    return misc_url_quote(name)


def process_phased_header(phasestream):
    '''strip the header off the phased file.
    At this instant, we assume fastphase output.
    '''
    line = 'start'
    while line:
        line = phasestream.readline()
        if line.startswith(FASTPHASE_START_TAG):
            break


import urllib.request, urllib.parse, urllib.error
def misc_url_quote(text):
    return urllib.parse.quote_plus(text, safe='+')

def misc_url_unquote(url):
    return urllib.parse.unquote_plus(url, safe='+')

if __name__ == '__main__':
    sys.exit(main(sys.argv))
