def parse_clustered_loci_file(filepath, gb_filepath, output_root,
                              circular, padding=1000, logger=None):
    """parses the clusters and returns a list of LociCluster objects

    Given a file from riboSelect or manually created (see specs in README)
    this parses the clusters and returns a list where [0] is sequence name
    and [1] is a list of loci in that cluster

    Args:
        filepath (str): path to loci file form riboSelect
        gb_filepath (str): path to genbank file from riboScan
        output_root (str): output root directory
        circular (bool): whether or not the gneome is circular; if so, the ends of
                         the sequences will be the padded
        padding (int):  how many base pairs to pad width
    Returns:
        (list): list of LociCluster objects
    Raises:
        Exception: one of the various things that can happen with bad files
        ValueError: No clusters found

    """
    assert logger is not None, "logging must be used!"
    clusters = []
    try:
        with open(filepath, "r") as f:
            file_contents = list(f)
    except Exception as e:
        logger.error("Cluster file could not be parsed!")
        raise e
    feature = None
    for line in file_contents:
        try:
            if line.startswith("#$ FEATURE"):
                try:
                    feature = line.split("FEATURE")[1].strip()
                except:
                    raise ValueError("Cannot extract FEATURE from '%s'" % line)
                continue
            elif line.startswith("#") or line.strip() == '':
                continue
            seqname = line.strip("\n").split(" ")[0]
            lt_list = [x for x in
                       line.strip("\n").split(" ")[1].split(":")]
        except Exception as e:
            logger.error("error parsing line: %s" % line)
            raise e
        # make and append the locus objects
        loci_list = []
        for i, loc in enumerate(lt_list):
            loci_list.append(Locus(index=i,
                                   locus_tag=loc,
                                   sequence_id=seqname))
        # make and append LociCluster objects
        clusters.append(LociCluster(mappings=[],
                                    output_root=output_root,
                                    sequence_id=seqname,
                                    loci_list=loci_list,
                                    padding=padding,
                                    circular=circular))
        # cluster_index = cluster_index + 1
    ### check feature i;f still none or starts with #$ (ie, no split)
    if feature is None:
        raise ValueError("no feature extracted from coords file! This " +
                         "has been made mandatory 20161108")
    ###
    if len(clusters) == 0:
        raise ValueError("No Clusters Found!!")
    for clu in clusters:
        clu.feat_of_interest = feature
    return clusters def pad_genbank_sequence(cluster, logger=None, verbose=False):
    """ circular genomes need "padding" to get flanking coordinates near the
    origin

    coords in coords list should be the 1st list item, with the index
    being 0th. Given a genbank record and a coord_list. this returns a seq
    padded on both ends by --padding bp, and returns a coord_list with coords
    adjusted accordingly.  Used to capture regions across origin.
    # as of 20161028, cluster object, not coord_list, is used
    """
    # take care of the coordinates
    if logger:
        logger.debug(str("adjusting coordinates by {0} to account for " +
                         "padding").format(cluster.padding))
    for loc in cluster.loci_list:
        if logger:
            logger.debug("pre-padded")
            logger.debug(str(loc.__dict__))
        start, end = loc.start_coord, loc.end_coord
        loc.start_coord, loc.end_coord = [start + cluster.padding,
                                          end + cluster.padding]
        if logger:
            logger.debug("post-padded")
            logger.debug(str(loc.__dict__))
    #TODO: check for interference with other clusters

    # take care of the sequence
    old_seq = cluster.seq_record.seq
    if cluster.padding > len(old_seq):
        logger.warning("padding cannot be greater than length of " +
                       "sequence! returning original sequence")
        return cluster
    new_seq = str(old_seq[-cluster.padding:]
                  + old_seq
                  + old_seq[0: cluster.padding])
    assert len(new_seq) == (len(old_seq) + (2 * cluster.padding)), \
        "Error within function! new seq should be len of " + \
        "seq plus 2x padding"
    cluster.seq_record = SeqRecord(Seq(new_seq))
    return cluster def get_genbank_record(input_genome_path, logger=None, first_only=True):
    """Returns a list of records in a GenBank file"""
    if logger:
        logger.info("Reading genbank file...")
    with open(input_genome_path) as input_genome_handle:
        records = list(SeqIO.parse(input_genome_handle, "genbank"))
    if first_only:
        return [records[0]]
    return records def combine_contigs(contigs_dir, pattern="*",
                    contigs_name="riboSeedContigs_aggregated",
                    ext=".fasta", verbose=False, logger=None):
    """changed over to biopython
    combine all *ext files in dir, return path to concatenated file
    requires Bio.SeqIO, glob, os
    """
    # make sure it ends in a path sep
    if contigs_dir[-1] != os.path.sep:
        contigs_dir = str(contigs_dir + os.path.sep)
    output = os.path.join(contigs_dir, str(contigs_name + ext))
    fastas = glob.glob(str(contigs_dir + pattern + ext))
    fastas.sort()
    if verbose:
        print(str("combining the following files matching pattern " +
                  "{0}:{1}".format(pattern, " ".join(fastas))))
    if logger:
        logger.debug(str("combining the following files matching pattern " +
                         "{0}:{1}".format(pattern, " ".join(fastas))))
    if len(fastas) == 0:
        if logger:
            logger.error("No matching files to combine found" +
                         " in {0}!".format(contigs_dir))
        sys.exit(1)
    with open(output, 'w') as w_file:
        for filen in fastas:
            with open(filen, 'r') as o_file:
                seq_records = SeqIO.parse(o_file, 'fasta')
                SeqIO.write(seq_records, w_file, 'fasta')

    return output def md5(fname, string=False):
    """straight from quantumsoup
    http://stackoverflow.com/questions/3431825/
        generating-an-md5-checksum-of-a-file
    updated 20160916 for strings; if true, fname can be a string
    """
    if string:
        return hashlib.md5(fname.encode('utf-8')).hexdigest()
    else:
        hash_md5 = hashlib.md5()
        with open(fname, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
            return hash_md5.hexdigest() def extract_coords_from_locus(cluster, feature="rRNA",
                              logger=None, verbose=False):
    """given a LociCluster object, ammend values
    20161028 returns a LociCluster
    """
    assert logger is not None, "logging must be used!"
    loc_number = 0  # index for hits
    locus_tags = [x.locus_tag for x in cluster.loci_list]
    if verbose:
        logger.debug("Locus tags for cluster %s: %s", cluster.index,
                     " ".join([x for x in locus_tags]))
    for feat in cluster.seq_record.features:
        if not feat.type in feature:
            continue
        if verbose:
            logger.debug("found {0} in the following feature : \n{1}".format(
                feature, feat))
        try:
            locus_tag = feat.qualifiers.get("locus_tag")[0]
        except:
            logger.error(str("found a feature ({0}), but there is no" +
                             "locus tag associated with it! Try formatting " +
                             "it by running scanScaffolds.sh").format(feat))
            raise ValueError

        # quick way of checking without using whole object
        if locus_tag in locus_tags:
            # make this_locus point to locus we are adding info to
            this_locus = next((x for x in cluster.loci_list if
                               x.locus_tag == locus_tag), None)
            #  SeqIO makes coords 0-based; the +1 below undoes that
            this_locus.start_coord = feat.location.start.position + 1
            this_locus.end_coord = feat.location.end.position
            this_locus.strand = feat.strand
            this_locus.product = feat.qualifiers.get("product")
            logger.debug("Added attributes for %s", this_locus.locus_tag)
            # logger.debug(str(this_locus.__dict__))
            logger.debug(
                "locus_tag: %s; coords: [%i-%i]; sequence_id: %s; " +
                "product: %s", this_locus.locus_tag,
                this_locus.start_coord, this_locus.end_coord,
                this_locus.sequence_id, this_locus.product[0])
            loc_number = loc_number + 1
        else:
            pass
            # logger.error("skipping %s", locus_tag)
    if not loc_number > 0:
        logger.error("no hits found in any record with feature %s! Double " +
                     "check your genbank file", feature)
        raise ValueError def add_gb_seqrecords_to_cluster_list(cluster_list, gb_filepath):
    """return cluster list with seq_records attached

    We need seqrecords in the lociCluster. So we look through the gb file
    to check for a sequence matching the LociClusters's sequence ID

    Args:
        cluster_list (list): List of LociCluster objects
        gb_filepath (str): oath to genbank file
    Returns:
        (list): the same cluster list, but with their seq records
    Raises:
        None

    """
    # match up seqrecords
    gb_records = SeqIO.index(gb_filepath, 'genbank')
    for clu in cluster_list:
        clu.seq_record = gb_records[clu.sequence_id]
    gb_records.close()
    return cluster_list def stitch_together_target_regions(cluster,
                                   flanking=500,
                                   logger=None, circular=False,
                                   revcomp=False):
    """
    given a single LociCluster object containing Locus objects (usually)
    of length 3 (16,5,and 23 rRNAs),
    return the object with ammended sequence info
    revamped 20161004
    """
    assert logger is not None, "Must have logger for this function"

    # TODO : make this safer. coord list is constructed sequentially but this
    # is a backup. Throws sort of a cryptic error. but as I said, its a backup
    if sorted([x.start_coord for x in cluster.loci_list]) != \
       [x.start_coord for x in cluster.loci_list]:
        logger.warning("Coords are not in increasing order; " +
                       "you've been warned")
    start_list = sorted([x.start_coord for x in cluster.loci_list])
    logger.debug("Start_list: {0}".format(start_list))

    logger.debug("stitching together the following coords:")
    for i in cluster.loci_list:
        logger.debug(str(i.__dict__))
    #  This works as long as coords are never in reverse order
    cluster.global_start_coord = min([x.start_coord for
                                      x in cluster.loci_list]) - flanking
    # if start is negative, just use 1, the beginning of the sequence
    if cluster.global_start_coord < 1:
        logger.warning("Caution! Cannot retrieve full flanking region, as " +
                       "the 5' flanking region extends past start of " +
                       "sequence. If this is a problem, try using a smaller " +
                       "--flanking region, and/or if  appropriate, run " +
                       "without with --linear.")
        cluster.global_start_coord = 1
    cluster.global_end_coord = max([x.end_coord for
                                    x in cluster.loci_list]) + flanking
    if cluster.global_end_coord > len(cluster.seq_record):
        logger.warning("Caution! Cannot retrieve full flanking region, as " +
                       "the 5' flanking region extends past start of " +
                       "sequence. If this is a problem, try using a smaller " +
                       "--flanking region, and/or if  appropriate, run " +
                       "without with --linear.")
        cluster.global_end_coord = len(cluster.seq_record)

    logger.debug("global start and end: %s %s", cluster.global_start_coord,
                 cluster.global_end_coord)
    #  the minus one makes things go from 1 based to zero based
    seq_with_ns = str(cluster.seq_record.seq[cluster.global_start_coord - 1:
                                             cluster.global_end_coord])
    seq_len = len(seq_with_ns[:])
    logger.debug(seq_len)
    # again, plus 1 corrects for 0 index.
    # len("AAAA") = 4 vs AAAA[-1] - AAAA[0] = 3
    logger.debug("\nexp length %i \nact length %i",
                cluster.global_end_coord - cluster.global_start_coord + 1,
                seq_len)

    ## Change coords in ID When using padding
    if not circular:
        seq_id = str(cluster.sequence_id + "_" + str(cluster.global_start_coord) +
                     ".." + str(cluster.global_end_coord))
    else:  # correct for padding
        seq_id = str(cluster.sequence_id + "_" + str(cluster.global_start_coord -
                                                     cluster.padding) +
                     ".." + str(cluster.global_end_coord - cluster.padding))

    # if most are on - strand, return sequence reverse complement
    strands = [x.strand for x in cluster.loci_list]
    if revcomp and \
       (sum([x == -1 for x in strands]) > sum([x == 1 for x in strands])):
        logger.info("returning the reverse compliment of the sequence")
        cluster.extractedSeqRecord = SeqRecord(
            Seq(seq_with_ns,
                IUPAC.IUPACAmbiguousDNA()).reverse_complement(),
            id=str(seq_id + "_RC"))
    else:
        cluster.extractedSeqRecord = SeqRecord(
            Seq(seq_with_ns,
                IUPAC.IUPACAmbiguousDNA()),
            id=seq_id)
    ### last minuete check
    cluster.extractedSeqRecord.description = "from riboSnag"
    for prop, value in vars(cluster).items():
        if value is None:
            logger.debug("%s has a value of None!", prop)
    return cluster def prepare_prank_cmd(outdir, combined_fastas, prank_exe,
                      add_args="", outfile_name="best_MSA",
                      logger=None):
    """returns command line for constructing MSA with
    PRANK and the path to results file
    """
    assert logger is not None, "Must use logger"
    if not os.path.exists(outdir):
        raise FileNotFoundError("output directory not found!")
    prank_cmd = "{0} {1} -d={2} -o={3}".format(
        prank_exe, add_args, combined_fastas,
        os.path.join(outdir, outfile_name))
    logger.debug("PRANK command: \n %s", prank_cmd)
    return (prank_cmd, os.path.join(outdir, str(outfile_name + ".best.fas"))) def prepare_mafft_cmd(outdir, combined_fastas, mafft_exe,
                      add_args="", outfile_name="best_MSA",
                      logger=None):
    """returns command line for constructing MSA with
    mafft and the path to results file.
    We call it best.fas to match the mafft result
    """
    assert logger is not None, "Must use logger"
    if not os.path.exists(outdir):
        raise FileNotFoundError("output directory not found!")
    mafft_cmd = "{0} {1} {2} > {3}".format(
        mafft_exe, add_args, combined_fastas,
        os.path.join(outdir, outfile_name + ".best.fas"))
    logger.debug("MAFFT command: \n %s", mafft_cmd)
    return (mafft_cmd, os.path.join(outdir, outfile_name + ".best.fas")) def calc_Shannon_entropy(matrix):
    """ $j$ has entropy $H(j)$ such that
    $H(j) = -sum_{i=(A,C,T,G)} p_i(j) log p_i(j)$
    """
    entropies = []
    for instance in matrix:
        unique = set(instance)
        proportions = {}
        for i in unique:
            proportions[i] = sum([x == i for x in instance]) / len(instance)
        entropy = -sum([prob * (math.log(prob, math.e)) for
                        prob in proportions.values()])
        entropies.append(entropy)
    return entropies def calc_entropy_msa(msa_path):
    """givn a path to an MSA in FASTA format, this gets the
    entropy of each position in batches, so as to not gobble memory
    return list
    """
    batch_size = 1000  # read seequences in chunks this long
    lengths = []
    with open(msa_path) as fh:
        msa_seqs = list(SeqIO.parse(fh, 'fasta'))
    seq_names = []
    for rec in msa_seqs:
        lengths.append(len(rec))
        seq_names.append(rec.id)
    if not all([i == lengths[0] for i in lengths]):
        raise ValueError("Sequences must all be the same length!")
    entropies = []
    tseq = []
    # save memory by reading in chunks
    for batch in range(0, (math.ceil(lengths[0] / batch_size))):
        # get all sequences into an array
        seq_array = []
        for nseq, record in enumerate(msa_seqs):
            seq_array.append(
                [x for x in record.seq[(batch * batch_size):
                                       ((batch + 1) * batch_size)]])
        # transpose
        tseq_array = list(map(list, zip(*seq_array)))
        tseq.extend(tseq_array)
        entropies.extend(calc_Shannon_entropy(tseq_array))
    # check length of sequence is the same as length of the entropies
    assert len(entropies) == lengths[0]
    return (entropies, seq_names, tseq) def annotate_msa_conensus(tseq_array, seq_file, barrnap_exe,
                          kingdom="bact",
                          pattern='product=(.+?)$',
                          countdashcov=True,   # include -'s in coverage
                          collapseNs=False,  # include -'s in consensus
                          excludedash=False,
                          logger=None):
    """ returns annotations (as gfflist),the consensus sequence as a
    list[base, cov], and named coords  as a list
    TODO: The 'next_best' thing fails is an N is most frequent. Defaults to a T
    """
    assert logger is not None, "Must use logger"
    if excludedash:
        logger.warning("CAUTION: excludedash selected. There is a known " +
                       "bug in the 'next_best' thing fails if an " +
                       "N is most frequent. Defaults to a T")
    consensus = []
    nseqs = len(tseq_array[0])
    logger.info("calc coverage for each of the %i positions", len(tseq_array))
    for position in tseq_array:
        if all([x == position[0] for x in position]):
            if position[0] == '-':
                if collapseNs:
                    continue
                elif not countdashcov:
                    consensus.append([position[0], 0])
            else:
                consensus.append([position[0], nseqs])
        else:
            max_count = 0  # starting count
            nextbest_count = 0
            best_nuc = None
            nextbest_nuc = None
            for nuc in set(position):
                count = sum([nuc == z for z in position])
                # if max count, swap with max to nextbest and update max
                if count > max_count:
                    nextbest_count = max_count
                    max_count = count  # update count if better
                    nextbest_nuc = best_nuc
                    best_nuc = nuc
                else:
                    pass
            # contol whether gaps are allowed in consensus
            if (
                    all([x != '-' for x in position]) and
                    best_nuc == '-' and
                    excludedash):
                # if we dont want n's, choose nextbest
                if nextbest_nuc is None:
                    nextbest_nuc = 't'  # I hate myself for this
                consensus.append([nextbest_nuc, nextbest_count])
            elif best_nuc == '-' and not countdashcov:
                consensus.append([best_nuc, 0])
            else:
                consensus.append([best_nuc, max_count])
    # if any are '-', replace with n's for barrnap
    seq = str(''.join([x[0] for x in consensus])).replace('-', 'n')

    # annotate seq
    with open(seq_file, 'w') as output:
        SeqIO.write(SeqRecord(
            Seq(seq, IUPAC.IUPACAmbiguousDNA())), output, "fasta"),
    barrnap_cmd = "{0} --kingdom {1} {2}".format(barrnap_exe,
                                                 kingdom, seq_file)
    barrnap_gff = subprocess.run(barrnap_cmd,
                                 shell=sys.platform != "win32",
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 check=True)
    results_list = [x.split('\t') for x in
                    barrnap_gff.stdout.decode("utf-8").split("\n")]

    ###  make [name, [start_coord, end_coord]] list
    named_coords = []
    for i in results_list:
        if i[0].startswith("#") or not len(i) == 9:
            logger.debug("skipping gff line: %s", i)
            continue
        m = re.search(pattern, i[8])
        if m:
            found = m.group(1)
        named_coords.append([found, [int(i[3]), int(i[4])]])
    if len(named_coords) == 0:
        raise ValueError(str("Error extracting coords from barrnap gff" +
                             "with pattern %s!" %  pattern))
    return (results_list, consensus, named_coords) def plot_scatter_with_anno(data,
                           consensus_cov,
                           anno_list,
                           names=["Position", "Entropy"],
                           title="$Strain Name$",
                           output_prefix="entropy_plot.png", pubplot=True):
    """Given annotation coords [feature, [start, end]],
    consensus cov list ['base', coverage_depth],
    entropy values (list) and consensus sequence
    (same length for consensus_cov and data, no funny business),
    plot out the entropies for each position,
    plot the annotations, and return 0
    """
    if len(consensus_cov) != len(data):
        raise ValueError("data and consensus different lengths!")
    df = pd.DataFrame({names[0]: range(1, len(data) + 1),
                       names[1]: data})  # columns=names)
    df_con = pd.DataFrame(consensus_cov, columns=["base",
                                                  "depth"])
    cov_max_depth = max(df_con['depth'])
    fig = Figure()
    FigureCanvas(fig)
    #fig.SubplotParams(hspace=0.0)
    # fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True,
    #                                gridspec_kw={'height_ratios': [4, 1]})
    gs = gridspec.GridSpec(2, 1, height_ratios=[4, 1])
    ax1 = fig.add_subplot(gs[0])
    ax2 = fig.add_subplot(gs[1])
    ax1.get_shared_x_axes().join(ax1, ax2)
    colors = ['#ff4c05', '#FFFB07', '#04FF08', '#06B9FF', '#6505FF', '#FF012F',
              '#ff4c05', '#FFFB07', '#04FF08', '#06B9FF', '#6505FF', '#FF012F']
    # ax1.set_title("Shannon Entropy by Position\n" +
    #               title, y=1.08, fontsize=20)
    if pubplot:
        fontscale = 1.5
    else:
        fontscale = 1
    ax1.set_title(title, y=1.08, fontsize=18 * fontscale)
    xmin, xmax = 0, len(data)
    ymin, ymax = -0.1, (max(data) * 1.2)
    ax1.set_xlim([xmin, xmax])
    ax1.set_ylim([ymin, ymax])
    if pubplot:
        # make minimal y labels for coeverage
        ax2.set_yticks([0, max(df_con["depth"])])
    yjust = -.1
    # add the highlighted bits sowing where te coding regions are
    for index, anno in enumerate(anno_list):
        rect1 = patches.Rectangle(
            (anno[1][0],  # starting x
             ymin),  # starting y
            anno[1][1] - anno[1][0],  # rel x end
            ymax - ymin,  # rel y end
            facecolor=matplotlib.colors.ColorConverter().to_rgba(
                colors[index], alpha=0.2),
            edgecolor=matplotlib.colors.ColorConverter().to_rgba(
                colors[index], alpha=0.2))
        rect2 = patches.Rectangle(
            (anno[1][0],  # starting x
             1),  # starting y
            anno[1][1] - anno[1][0],  # rel x end
            cov_max_depth,  # dont -1 beacuse start at 1
            facecolor=matplotlib.colors.ColorConverter().to_rgba(
                colors[index], alpha=0.2),
            edgecolor=matplotlib.colors.ColorConverter().to_rgba(
                colors[index], alpha=0.2))
        ax1.add_patch(rect1)
        ax2.add_patch(rect2)
        if not pubplot:
            ax1.text((anno[1][0] + anno[1][1]) / 2,    # x location
                     ymax - 0.48 - yjust,              # y location
                     anno[0][0:20],                     # text first 20 char
                     ha='center', color='red', weight='bold', fontsize=11)
            yjust = yjust * - 1
    ax1.scatter(x=df["Position"], y=df["Entropy"],
                marker='o', color='black', s=2)
    # add smoothing for kicks
    df["fit"] = savitzky_golay(df["Entropy"].values, 351, 3)  # window size 51, polynomial order 3
    ax1.scatter(x=df["Position"], y=df["fit"], color='red', s=1)
    #
    ax1.set_ylabel('Shannon Entropy', fontsize= 13 * fontscale)
    ax1.get_yaxis().set_label_coords(-.05, 0.7)
    ax2.set_xlim([xmin, xmax])
    ax2.invert_yaxis()  # we want the bars pointing down
    ax2.set_ylabel('Consensus Coverage', fontsize= 13 * fontscale)
    ax2.set_xlabel('Position (bp)', fontsize= 13 * fontscale)
    ax2.get_yaxis().set_label_coords(-.05, 0.5)
    # ax2.set_ylim([1, cov_max_depth + 1]) #, 1])
    ax2.bar(df_con.index, df_con["depth"],
            width=1, color='darkgrey', linewidth=0, edgecolor='darkgrey')
    # ax2.step(df_con.index, df_con["depth"],
    #          where='mid', color='darkgrey')
    for ax in [ax1, ax2]:
        ax.spines['right'].set_visible(False)
    # Only show ticks on the left and bottom spines
    ax1.spines['top'].set_visible(False)
    ax2.spines['bottom'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax2.xaxis.set_ticks_position('bottom')
    ax1.xaxis.set_ticks_position('top')
    ax1.tick_params(axis='y', colors='dimgrey', labelsize = 10 * fontscale)
    ax2.tick_params(axis='y', colors='dimgrey', labelsize = 10 * fontscale)
    ax1.tick_params(axis='x', colors='dimgrey', labelsize = 10 * fontscale)
    ax2.tick_params(axis='x', colors='dimgrey', labelsize = 10 * fontscale)
    ax1.yaxis.label.set_color('black')
    ax2.yaxis.label.set_color('black')
    ax1.xaxis.label.set_color('black')
    ax2.xaxis.label.set_color('black')
    fig.tight_layout()
    fig.subplots_adjust(hspace=0)
    fig.set_size_inches(12, 7.5)
    fig.savefig(str(output_prefix + '.png'), dpi=(300))
    fig.savefig(str(output_prefix + '.pdf'), dpi=(300))
    return 0 def get_all_kmers(alph="", length=3):
    """Given an alphabet of charactars, return a list of all permuations
    not actually used, I dont think I'll need this
    """
    mers = [''.join(p) for p in product(alph, repeat=length)]
    return mers def profile_kmer_occurances(rec_list, k, logger=None):
    """ given a list of seq records, an alphabet, and a value k,
    retrun counts dict of kmer occurances and list of seq names
    """
    counts = defaultdict(list)
    names_list = []
    # part 1: get all kmers from seqs so that we can have equal length lisst
    all_mers = []
    for rec in rec_list:
        all_mers.extend([str(rec.seq).lower()[x: x + k] for
                         x in range(0, len(rec.seq) - k + 1)])
    # unique_mers = list(set().union(all_mers))
    unique_mers = set(all_mers)
    for i in unique_mers:
        counts[i] = []  # initialixe counts dictionary with ker keys
    # part two: count 'em
    for n, rec in enumerate(rec_list):
        logger.info("counting kmer occurances for %s", rec.id)
        names_list.append(rec.id)
        logger.debug("converting to lower")
        string = str(rec.seq).lower()
        logger.debug("getting %imers from seq", k)
        string_mers = [string[x: x + k] for x in range(0, len(string) - k + 1)]
        logger.debug("counting mer occurances")
        # Add counts for those present
        for value in set(string_mers):
            counts[value].append(sum([value == mer for mer in string_mers]))
        # filling in where not found
        for missing in unique_mers - set(string_mers):
            counts[missing].append(0)
    return counts, names_list def plot_pairwise_least_squares(counts, names_list, output_prefix):
    """given a  list of sequence names, a prefix for plot file names,
    and a list of counts from plot_kmer_occurances,
    retruns a pandas df of least squares after plotting heatmaps
    """
    res_list = []
    counts_list = []
    for v in counts.values():
        counts_list.append(v)
    # this gives each an index and gets all the pairs
    all_pairs = [[index, value] for index, value in
                 enumerate(product(range(0, len(names_list)), repeat=2))]
    for i in all_pairs:
        this_pairs_diffs = []
        for row in counts_list:
            this_pairs_diffs.append((row[i[1][0]] - row[i[1][1]]) ** 2)
        res_list.append([names_list[i[1][0]],
                         names_list[i[1][1]],
                         sum(this_pairs_diffs)])
    lsdf_wNA = pd.DataFrame(res_list, columns=["locus_1", "locus_2", "sls"])
    wlsdf = lsdf_wNA.pivot(index='locus_1', columns='locus_2', values='sls')
    fig = Figure()
    FigureCanvas(fig)
    ax = fig.add_subplot(111)
    #fi, ax = plt.subplots(1, 1)
    # lsdf = lsdf_wNA.fillna(value=0)
    heatmap = ax.pcolormesh(wlsdf, cmap='Greens')
    # put the major ticks at the middle of each cell
    ax.set_yticks(np.arange(wlsdf.shape[0]) + 0.5, minor=False)
    ax.set_xticks(np.arange(wlsdf.shape[1]) + 0.5, minor=False)
    # # want a more natural, table-like display
    ax.invert_yaxis()
    ax.xaxis.tick_top()
    # # Set the labels
    ax.set_xticklabels([x.replace('_', '-') for x in wlsdf.columns.values],
                       minor=False, rotation=90)
    ax.set_yticklabels([x.replace('_', '-') for x in wlsdf.index],
                       minor=False)
    fig.colorbar(heatmap)  # add colorbar key
    fig.tight_layout()  # pad=0, w_pad=5, h_pad=.0)
    # fig.subplots_adjust(top=0.4, bottom=0.0, hspace=0, wspace=0.3 )
    fig.set_size_inches(8, 8)
    fig.savefig(str(output_prefix + "heatmap.png"), dpi=(300))
    fig.savefig(str(output_prefix + "heatmap.pdf"), dpi=(300))
    ####  plot clustered heatmap
#    plt.close('all')
#    plt.figure(1, figsize=(6, 6))
    # h = heatmapcluster(wlsdf.as_matrix(), row_labels=wlsdf.index,
    #                    col_labels=wlsdf.columns.values,
    #                    num_row_clusters=2, num_col_clusters=2,
    #                    label_fontsize=6,
    #                    xlabel_rotation=-75,
    #                    cmap=plt.cm.coolwarm,
    #                    show_colorbar=True,
    #                    top_dendrogram=True)
    # # plt.tight_layout()  # pad=0, w_pad=5, h_pad=.0)
    # # fig2.set_size_inches(16, 16)
    # plt.savefig(str(output_prefix + "clustered_heatmap.png"), dpi=(200))
    # plt.savefig(str(output_prefix + "clustered_heatmap.pdf"), dpi=(200))
    return lsdf_wNA def make_msa(msa_tool, unaligned_seqs, prank_exe, mafft_exe,
             args, outdir, logger=None):
    """returns msa cmd and results path
    """
    assert logger is not None, "Must use logger"
    if msa_tool == "prank":
        if shutil.which(prank_exe):
            msa_cmd, results_path = prepare_prank_cmd(
                outdir=outdir,
                outfile_name="best_MSA",
                combined_fastas=unaligned_seqs,
                prank_exe=prank_exe,
                add_args=args,
                logger=logger)
        else:
            raise ValueError("Construction of MSA skipped because " +
                             "%s is not a valid executable!", prank_exe)
    elif msa_tool == "mafft":
        if shutil.which(mafft_exe):
            msa_cmd, results_path = prepare_mafft_cmd(
                outdir=outdir,
                outfile_name="best_MSA",
                combined_fastas=unaligned_seqs,
                mafft_exe=mafft_exe,
                add_args=args,
                logger=logger)
        else:
            raise ValueError("Construction of MSA skipped because " +
                             "%s is not a valid executable!", mafft_exe)
    return(msa_cmd, results_path) def check_loci_file_not_genbank(filepath):
    """ raise error if "cluster file" ends in .gb, etc

    this covers common case where user submits genbank and cluster file
    in the wrong order.

    Args:
        filepath (str): path to what should be a loci file
    Returns:
        None
    Raises:
       FileNotFoundError: if the file smells like a genbank file

    """
    if not (os.path.isfile(filepath) and os.path.getsize(filepath) > 0):
        raise ValueError("Cluster File not found!")
    if filepath.endswith(("gb", "genbank", "gbk")):
        raise FileNotFoundError("Hmm, this cluster file looks like genbank;" +
                                " it ends in {0}".format(os.path.splitext(
                                    filepath)[1])) class LociCluster(object):
    """ organizes the clustering process instead of dealing with nested lists
    This holds the whole cluster of one to several individual loci
    """
    newid = itertools.count()

    def __init__(self, sequence_id, loci_list, padding=None,
                 global_start_coord=None, global_end_coord=None,
                 seq_record=None, feat_of_interest=None, mappings=None,
                 extractedSeqRecord=None, cluster_dir_name=None,
                 coverage_exclusion=None,
                 circular=False, output_root=None, final_contigs_path=None,
                 continue_iterating=True, keep_contigs=True):
        # int: unique identifier for cluster
        self.index = next(LociCluster.newid)
        # str: sequence name, usually looks like 'NC_17777373.1' or similar
        self.sequence_id = sequence_id
        # list: hold locus objects for each item in cluster
        self.loci_list = loci_list  # this holds the Locus objects
        # int: bounds ____[___.....rRNA....rRNA..rRNA....__]_________
        self.global_start_coord = global_start_coord
        self.global_end_coord = global_end_coord
        # int: how much to pad sequences y if treating as circular
        self.padding = padding
        # str: feature for filtering: rRNA, cDNA, exon, etc
        self.feat_of_interest = feat_of_interest
        # Bool: treat seqs as circular by padding the ends
        self.circular = circular
        # path: where your cluster-specific output goes
        self.cluster_dir_name = cluster_dir_name  # named dynamically
        # path: where the overall output goes
        self.output_root = output_root
        # list: lociMapping objects that hold mappinging paths
        self.mappings = mappings
        # SeqRecord: holds SeqIO Seqrecord for sequence_id
        self.seq_record = seq_record
        # SeqRecord: holds SeqIO Seqrecord for seq extracted from global coords
        self.extractedSeqRecord = extractedSeqRecord
        # path: for best contig after riboseed2 iterations
        self.keep_contigs = keep_contigs  # by default, include all
        self.continue_iterating = continue_iterating  # by default, keep going
        self.coverage_exclusion = coverage_exclusion
        self.final_contig_path = final_contigs_path
        self.name_mapping_dir()

    def name_mapping_dir(self):
        self.cluster_dir_name = str("{0}_cluster_{1}").format(
            self.sequence_id, self.index) class Locus(object):
    """ this holds the info for each individual Locus"
    """
    def __init__(self, index, sequence_id, locus_tag, strand=None,
                 start_coord=None, end_coord=None,
                 product=None,
                 feature_type=None):
        # int: unique identifier for cluster
        self.index = index
        # str: sequence name, usually looks like 'NC_17777373.1' or similar
        self.sequence_id = sequence_id  # I dont think this is this needed?
        # str: unique identifier from \locus_tag= of gb file
        self.locus_tag = locus_tag
        # int: 1 is +strand, -1 is -strand
        self.strand = strand
        # int:
        self.start_coord = start_coord
        self.end_coord = end_coord
        # str: from \product= of gb file
        self.product = product
        # str "feature" genbank identifier
        self.feature_type = feature_type
import copy
import sys
import logging
import os
import unittest
import shutil
import time
from Bio import SeqIO
from riboSeed.shared_methods import parse_clustered_loci_file, \
    pad_genbank_sequence, get_genbank_record, combine_contigs, md5, \
    extract_coords_from_locus, add_gb_seqrecords_to_cluster_list
from riboSeed.riboSnag import stitch_together_target_regions, \
    prepare_prank_cmd, prepare_mafft_cmd,\
    calc_Shannon_entropy, calc_entropy_msa,\
    annotate_msa_conensus, plot_scatter_with_anno, get_all_kmers,\
    profile_kmer_occurances, plot_pairwise_least_squares, make_msa, \
    check_loci_file_not_genbank
from riboSeed.classes import LociCluster, Locus
        if not os.path.exists(self.testdirname):
            os.makedirs(self.testdirname, exist_ok=True)

    def test_check_loci_not_gb(self):
        """ check that the loci file is not a genbank file
        """
        with self.assertRaises(FileNotFoundError):
            check_loci_file_not_genbank(self.test_gb_file)

    def test_parse_loci(self):
        """this checks the parsing of riboSelect ouput
        """
        clusters = parse_clustered_loci_file(
            filepath=self.test_loci_file,
            gb_filepath=self.test_gb_file,
            output_root=self.testdirname,
            padding=100,
            circular=True,
            logger=logger)
        ref_seq_name = 'NC_011751.1'
        ref_loci_list = ['ECUMN_16S_6', 'ECUMN_23S_6', 'ECUMN_5S_7']
        test_locus_tags = [x.locus_tag for x in clusters[0].loci_list]
        self.assertEqual(clusters[0].sequence_id, ref_seq_name)
        self.assertEqual(test_locus_tags, ref_loci_list)

    # def test_genbank_match_id(self):
    #     """ tests get_genbank_rec_from_multigb
    #     """
    #     records = get_genbank_record(self.test_gb_file)
    #     record = get_genbank_rec_from_multigb(recordID='NC_011751.1',
    #                                           genbank_records=records)
    #     self.assertEqual(records[0].seq, record.seq)
    #     with self.assertRaises(ValueError):
    #         get_genbank_rec_from_multigb(
    #             recordID='NC_011751.X', genbank_records=records)

    def test_extract_coords_from_locus(self):
        """todo: replace essentially unused get_genbank_record function
        """
        test_cluster = LociCluster(
            sequence_id='NC_011751.1',
            loci_list=[Locus(index=0,
                             sequence_id='NC_011751.1',
                             locus_tag="ECUMN_0004",
                             strand=None,
                             start_coord=None,
                             end_coord=None,
                             product=None)],
            padding=None,
            global_start_coord=None,
            global_end_coord=None)
        test_cluster.seq_record = get_genbank_record(self.test_gb_file)[0]
        extract_coords_from_locus(
            cluster=test_cluster,
            feature="CDS",
            logger=logger)
        cluster_post_extract = test_cluster
        self.assertEqual(cluster_post_extract.loci_list[0].index, 0)
        self.assertEqual(
            cluster_post_extract.loci_list[0].locus_tag, "ECUMN_0004")
        self.assertEqual(
            cluster_post_extract.loci_list[0].strand, 1)
        self.assertEqual(
            cluster_post_extract.loci_list[0].start_coord, 3692)
        self.assertEqual(
            cluster_post_extract.loci_list[0].end_coord, 4978)
        self.assertEqual(
            cluster_post_extract.loci_list[0].sequence_id, 'NC_011751.1')

    def test_pad_genbank_sequence(self):
        """Just tests the function, later we wll test that the same
        sequence is extracted padded and otherwise
        """
        padding_val = 50
        test_cluster = LociCluster(
            sequence_id='NC_011751.1',
            loci_list=[Locus(index=0,
                             sequence_id='NC_011751.1',
                             locus_tag="ECUMN_0004",
                             strand=None,
                             start_coord=None,
                             end_coord=None,
                             product=None)],
            padding=padding_val,
            global_start_coord=None,
            global_end_coord=None)
        test_cluster.seq_record = get_genbank_record(self.test_gb_file)[0]
        extract_coords_from_locus(cluster=test_cluster,
                                  feature="CDS",
                                  logger=logger)
        old_seq = test_cluster.seq_record
        # test excessive padding
        cluster_too_much_padding = copy.deepcopy(test_cluster)
        cluster_too_much_padding.padding = 10000000
        # with self.assertRaises(ValueError):
        #     pad_genbank_sequence(cluster_too_much_padding,
        #                          logger=logger)

        # test sucessful execution
        padded_cluster = pad_genbank_sequence(test_cluster,
                                              logger=logger)
        self.assertEqual(str(old_seq.seq),
                         str(padded_cluster.seq_record.seq[padding_val:
                                                          - padding_val]))

    def test_stitching(self):
        """  This is actually the thing needing the most testing, most likely
        """
        padding_val = 1000
        clusters = parse_clustered_loci_file(
            filepath=self.test_loci_file,
            gb_filepath=self.test_gb_file,
            output_root=self.testdirname,
            padding=padding_val,
            circular=True,
            logger=logger)
        clusters = add_gb_seqrecords_to_cluster_list(
            cluster_list=clusters,
            gb_filepath=self.test_gb_file)
        cluster = clusters[0]  # this should be reverse complimented

        extract_coords_from_locus(