__author__ = 'Alexendar Perez & Yuri Pritykin'

#####################
#                   #
#   Introduction    #
#                   #
#####################

"""Analyze k-mers and extract guideRNAs with all relevant info about them.

analyze_guides(): main function, create, modify and save trie, print log
build_kmers_trie(): create initial trie of k-mers and their coordinates
label_multimapping(): label multimapping k-mers in trie with correct counts
filter_trie_mismatch_similarity(): find keys in trie similar to other keys
filter_keys_trie(): using labels in trie, choose good k-mers from a list

trie_get_approximate(): wrapper of trie.trie.get_approximate()
trie_get_approximate_hamming(): deprecated; use trie.get_approximate_hamming()

save_trie(): save trie to file
load_trie(): load trie from file
restore_trie_arrays(): restore numpy.array from str in loaded kmer trie
load_restore_trie(): load and restore trie of all genomic kmers for a project
"""

#################
#               #
#   Libraries   #
#               #
#################

import os
import gzip
from collections import defaultdict

import numpy as np
# from Bio import trie

import util
import trie

#########################
#                       #
#   Auxillary Function  #
#                       #
#########################

def save_trie(t, filename):
    """Save trie to file.

    Note: saving and loading tries (especially with np.array values)
    may be system-dependent (e.g., depend on 64bit or 32bit arithmetic),
    need to test.

    Args:
    t: trie.trie object
    filename: where to save
    """
    f = open(filename, 'w')
    trie.save(f, t)
    f.close()

def load_trie(filename):
    """Load trie from file.

    Note: saving and loading tries (especially with np.array values)
    may be system-dependent (e.g., depend on 64bit or 32bit arithmetic),
    need to test.
    """
    f = open(filename)
    t = trie.load(f)
    f.close()
    return t

def restore_trie_arrays(kmers_trie, keysfile):
    """Restore numpy arrays from str in values of kmer trie loaded from disk.

    When loaded from disk, numpy arrays in values of a previously stored trie
    are loaded as string characters. This function restores the initial arrays
    using numpy.fromstring(). To save memory, avoid building a list of keys
    of the trie and take them from file.

    Note: saving and loading tries (especially with np.array values)
    may be system-dependent (e.g., depend on 64bit or 32bit arithmetic
    and endianness), use caution when transferring files between systems.

    Args:
    kmers_trie: trie.trie object loaded from dist (e.g., using load_trie())
    keysfile: name of file where first field of each line is a k-mer, assume
              file is gzipped; loop only over the k-mers in this file
    """
    util.check_file_exists(keysfile)
    f = gzip.open(keysfile)
    for line in f:
        kmer = line.split()[0]
        if kmers_trie.has_key(kmer):
            value = kmers_trie[kmer]
            if isinstance(value, basestring):
                kmers_trie[kmer] = np.fromstring(value, dtype=int)
    f.close()
    return kmers_trie

def load_restore_trie(name):
    """Fully load previously stored trie of all genomic kmers for a project.

    name: project name, used to get project args and in all output

    Return:
    loaded kmers_trie
    """
    util.print_log('load trie...')
    trie_filename = '%s/%s_trie.dat' % (name, name)
    kmers_trie = load_trie(trie_filename)
    util.print_log('done')
    keysfile = '%s/%s_kmers_shuffled.txt.gz' % (name, name)
    util.print_log('restore numpy arrays in trie values...')
    restore_trie_arrays(kmers_trie, keysfile)
    util.print_log('done')
    return kmers_trie

def trie_get_approximate(tr, seq, dist):
    """Wrapper of trie.trie.get_approximate()

    Find all keys in trie tr at edit (Levenshtein) distance at most dist
    from seq.

    Args:
    tr: trie.trie object
    seq: key
    dist: distance

    Return: dict {key : distance}
    """
    d = defaultdict(lambda : np.inf)
    for p in tr.get_approximate(seq, dist):
        d[p[0]] = min(d[p[0]], p[2])
    return d.items()

def trie_get_approximate_hamming(tr, seq, dist):
    """Find all keys in trie tr at Hamming distance at most dist from seq.

    Deprecated: now trie has method get_approximate_hamming implemented.

    Args:
    tr: trie.trie object
    seq: key
    dist: distance

    Return: dict {key : distance}
    """
    pairs = trie_get_approximate(tr, seq, dist)
    d = ((k, util.hamming(k, seq)) for k,edit_dist in pairs)
    d = dict((k, h) for k,h in d if h <= dist)
    return d

def build_kmers_trie(filename, genome, altpam=[], pampos='end', maxcount=10,
                     goodkeysfile='',badkeysfile=''):
    """Read k-mers and their coordinates from file and store them in a trie.

    The resulting trie is of the form {<k-mer> : <np.array of int values>}.
    In each array, store int-transformed coordinates of occurrences of
    the k-mer, starting from position 1 of the array. Store at most
    'maxcount' coordinates.
    Position 0 in the array is reserved for labeling:
    0: good guideRNA,
    positive: show how many occurrences this k-mer has in the genome
    [other labels: decide later]
    In this building stage, label all k-mers with alternative PAM
    and all k-mers with more than one occurrence in the genome
    as bad guideRNAs.
    Optionally, also store in a separate file all k-mers that are still
    considered good candidate guideRNAs. This means only filtering out
    k-mers with alternative PAM, because detecting multi-mapping k-mers
    labeling them as bad guideRNAs may happen after they were first read.

    Note: make sure lines with k-mers in input file are randomly shuffled.
    This is to ensure that for k-mers with more than 'maxcount' occurrences,
    we store artibrary 'maxcount' of them without any bias.

    Args:
    filename: name of file with k-mers, assume file is gzipped and
              lines are randomly shuffled
    genome: list of pairs [(<chromosome name>, <chromosome length>)]
    altpam: list of alternative PAM sequences, all k-mers starting
            or ending (depending on argument 'pampos') with these
            sequences are labeled as bad guideRNAs
    pampos: position of alternative PAM in k-mer ('start' or 'end')
    maxcount: store at most this many coordinates for each k-mer
    goodkeysfile: where to store potentially good candidate guideRNAs;
                  use only if altpam is not empty, otherwise all input
                  keys from filename will be stored which is redundant

    Output:
    return trie.trie object {<k-mer> : <np.array of int values>}
    optionally produce file goodkeysfile with candidate guideRNAs
    """
    util.check_file_exists(filename)
    if goodkeysfile:
        util.warn_file_exists(goodkeysfile)
        goodkeys = gzip.open(goodkeysfile, 'w')
    if badkeysfile:
        util.warn_file_exists(badkeysfile)
        badkeys = gzip.open(badkeysfile,'w')
    kmers_trie = trie.trie()
    f = gzip.open(filename)
    count_added = 0
    for line in f:
        kmer, coord = line.split()
        if kmers_trie.has_key(kmer):
            arr = kmers_trie[kmer]
            if len(arr) < maxcount + 1:
                coord_int = util.map_coord_to_int(coord, genome)
                arr = np.append(arr, coord_int)
                arr[0] = len(arr) - 1
                kmers_trie[kmer] = arr
        else:
            coord_int = util.map_coord_to_int(coord, genome)
            label = 0
            if pampos == 'start' and any(kmer.startswith(p) for p in altpam):
                label = 1
            if pampos == 'end' and any(kmer.endswith(p) for p in altpam):
                label = 1
            kmers_trie[kmer] = np.array([label, coord_int])
            count_added += 1
            if goodkeysfile and label == 0:
                goodkeys.write('%s\n' % kmer)
            if badkeysfile and label != 0:
                badkeys.write('%s\n' %kmer)
    f.close()
    if goodkeysfile:
        goodkeys.close()
    if badkeysfile:
        badkeys.close()
    print '%s keys added to trie' % count_added
    return kmers_trie

def label_multimapping(kmers_trie, filename):
    """Read multimapping k-mers and counts from file and add counts to trie.

    For each kmer and its count, find it in input trie and update the count
    of the kmer stored in position 0 of the array in the value of the trie.

    Args:
    kmers_trie: trie.trie object of the form described in and returned
                by build_kmers_trie()
    filename: name of file with k-mers and counts for k-mers occurring too
              many times in the genome,
              e.g., file <name>/<name>_kmers_counts.txt.gz
              produced by kmers.sort_count_kmers();
              assume file is gzipped
    
    Return:
    modified input trie.trie object
    """
    util.check_file_exists(filename)
    f = gzip.open(filename)
    count_labeled = 0
    for line in f:
        kmer, count = line.split()
        if kmers_trie.has_key(kmer):
            kmers_trie[kmer][0] = count
            count_labeled += 1
    f.close()
    print '%s k-mers assigned counts' % count_labeled
    return kmers_trie

def filter_keys_trie(kmers_trie, keysinputfile, keysoutputfile,nonCandidatekeysoutputfile):
    """Select k-mers from file that have label 0 in trie and write to file.

    Args:
    kmers_trie: trie.trie object of the form described in and returned
                by build_kmers_trie()
    keysinputfile: name of file where first field of each line is a k-mer,
                   assume file is gzipped
    keysoutputfile: name of file where to write selected keys one per line,
                    gzipped
    """
    util.check_file_exists(keysinputfile)
    keysinput = gzip.open(keysinputfile)
    util.warn_file_exists(keysoutputfile)
    keysoutput = gzip.open(keysoutputfile, 'w')
    non_candidate_keysoutput = gzip.open(nonCandidatekeysoutputfile, 'w')
    read_count = 0
    write_count = 0
    for line in keysinput:
        kmer = line.split()[0]
        read_count += 1
        if kmers_trie.has_key(kmer) and kmers_trie[kmer][0] == 0:
            keysoutput.write('%s\n' % kmer)
            write_count += 1
        elif kmers_trie.has_key(kmer) and kmers_trie[kmer][0] != 0:
            non_candidate_keysoutput.write('%s\n' % kmer)
    keysoutput.close()
    non_candidate_keysoutput.close()
    keysinput.close()
    print '%s keys read' % read_count
    print '%s keys written' % write_count

def filter_trie_mismatch_similarity(kmers_trie, sim, keysfile):
    """Find keys in trie few mismatches away from other keys.

    Args:
    kmers_trie: trie.trie object of the form described in and returned
                by build_kmers_trie()
    sim: if a key has another key at Hamming distance at most this,
         label it as bad guideRNA
    keysfile: name of file where first field of each line is a k-mer, assume
              file is gzipped; loop only over the k-mers in this file
    """
    util.check_file_exists(keysfile)
    f = gzip.open(keysfile)
    count_labeled = 0
    for line in f:
        kmer = line.split()[0]
        if not kmers_trie.has_key(kmer):
            continue
        if kmers_trie[kmer][0] != 0:
            continue
#        dist = trie_get_approximate_hamming(kmers_trie, kmer, sim)
        all_dist = kmers_trie.get_approximate_hamming(kmer, sim)
#        if any(v > 0 for v in dist.itervalues()):
        if any(dist > 0 for seq,arr,dist in all_dist):
            kmers_trie[kmer][0] = 1
            count_labeled += 1
#        for k,v in dist.iteritems():
        for seq,arr,dist in all_dist:
            # if seq is similar to kmer, can label seq as bad too
            # this avoids excessive trie.get_approximate_hamming() calls
            if dist > 0 and kmers_trie[seq][0] == 0:
                kmers_trie[seq][0] = 1
                count_labeled += 1
    f.close()
    print '%s k-mers labeled as bad guideRNAs' % count_labeled
    return kmers_trie

def analyze_guides(name):
    """Analyze k-mers and find all candidate guideRNAs and their off-targets.

    Load project arguments, build and analyze a trie, find guideRNAs.
    Run after all files were generated by kmers.extract_process_kmers()

    Produce files:
    trie with all k-mers, values store label for good or bad candidate
        guideRNA and coordinates in the genome: <name>/<name>_trie.dat
    intermediate files with candidate guideRNA k-mers used as keys
        to the trie: <name>/<name>_triekeys_v?.txt.gz
    final list of guideRNAs: <name>/<name>_guides.txt.gz

    Args:
    name: project name, used to get project args and in all output

    Return:
    trie.trie object with all k-mers, their coordinates in the genome,
    and labels of good and bad candidate guideRNAs
    """
    util.print_log('start analyze_guides()')
    util.print_log('load arguments...')
    args = util.load_args(name)
    util.print_args(args)
    util.print_log('done')

    if os.path.exists('%s%s' % (name,'/blacklist')):
        util.print_log('blacklist directory already exists \n')
        pass
    else:
        os.mkdir(('%s%s' % (name,'/blacklist')))
        util.print_log('blacklist directory made \n')

    util.print_log('construct trie...')
    kmers_filename = '%s/%s_kmers_shuffled.txt.gz' % (name, name)
    util.print_log('load k-mers from %s' % kmers_filename)
    genome = args['genome']
    goodkeysfile = '%s/%s_triekeys_v1.txt.gz' % (name, name) \
                   if args['altpam'] else ''
    badkeysfile = '%s/%s/%s_nonCandidate_triekeys_with_altpams.txt.gz' % (name,'blacklist',name) if args['altpam'] else ''
    if goodkeysfile:
        util.print_log('print candidate guideRNAs to %s' % goodkeysfile)
    kmers_trie = build_kmers_trie(kmers_filename, genome,
                                  altpam=args['altpam'], pampos=args['pampos'],
                                  maxcount=args['maxoffpos'],
                                  goodkeysfile=goodkeysfile,badkeysfile=badkeysfile)
    util.print_log('done')

    util.print_log('label as bad guideRNAs multimapping k-mers in trie...')
    keysinputfile = goodkeysfile if goodkeysfile else kmers_filename
    keysoutputfile = '%s/%s_triekeys_v2.txt.gz' % (name, name)
    nonCandidatekeysoutputfile = '%s/%s/%s_nonCandidate_triekeys_targetSites_with_multiple_perfect_hits.txt.gz' %\
                                 (name,'blacklist',name)
    util.print_log('read keys from %s and write to %s'
                   % (keysinputfile, keysoutputfile))
    filter_keys_trie(kmers_trie, keysinputfile, keysoutputfile,nonCandidatekeysoutputfile)
    util.print_log('done')

    util.print_log('assign correct counts to multimapping k-mers in trie...')
    count_filename = '%s/%s_kmers_counts.txt.gz' % (name, name)
    util.print_log('read counts from %s' % count_filename)
    kmers_trie = label_multimapping(kmers_trie, count_filename)
    util.print_log('done')

    util.print_log('label as bad guideRNAs k-mers in trie few mismatches away'
                   ' from other k-mers...')
    sim = args['sim'] - 1
    util.print_log('label as bad k-mers with other k-mer at distance <=%s'
                   % sim)
    keysfile = '%s/%s_triekeys_v2.txt.gz' % (name, name)
    util.print_log('read keys from %s' % keysfile)
    filter_trie_mismatch_similarity(kmers_trie, args['sim'] - 1, keysfile)
    util.print_log('done')

    util.print_log('produce list of good guideRNAs...')
    keysinputfile = keysfile
    keysoutputfile = '%s/%s_guides.txt.gz' % (name, name)
    nonCandidatekeysoutputfile = '%s/%s/%s_nonCandidate_guides_with_mismatch_neighbors.txt.gz' % (name,'blacklist',name)
    util.print_log('read keys from %s and write to %s'
                   % (keysinputfile, keysoutputfile))
    filter_keys_trie(kmers_trie, keysinputfile, keysoutputfile,nonCandidatekeysoutputfile)
    util.print_log('done')

    util.print_log('save trie...')
    trie_filename = '%s/%s_trie.dat' % (name, name)
    util.print_log('save in file %s' % trie_filename)
    save_trie(kmers_trie, trie_filename)
    util.print_log('done')

    return kmers_trie
