#!/usr/bin/env python

# ----------------------------------------------------------------------------
# Copyright (c) 2015--, The WGS-HGT Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

# Implement Distance Method for HGT detection based on algorithm described
#    in:
#        Wei. X et al., "A Distance-Based Method for Detecting HGT in Whole
#        Genomes", International Symposium on Bioinformatics Research and
#        Applications (ISBRA), 2008, pages 26-37
#
#    The workflow follows the algorithm:
#    1. For each gene in target genome,
#        i.    BLAST sequence against all other genes in the reference
#              genomes;
#        ii.   Go to step 3 if gene has more than threshold number of homologs
#              (min-num-homologs), otherwise go to next gene in target genome;
#        iii.  Compute multiple sequence alignment on homolog genes using
#              CLUSTAL;
#        iv.   Compute pairwise distance matrix using PHYLIP's protdist
#              function and Z-score normalize the set of pairwise distances
#              for each gene family and species;
#        v.    Add distance matrix for all pairwise distances into global
#              distance matrix storing results for all genes
#
#    2. Cluster gene families by species,
#        vi.   Compute all species sets (sets of genes whose orthologs are
#              detectable in exactly the same subset of the considered
#              species);
#        vii.  Cluster genes to each core species set cluster using the
#              Hamming distance clustering algorithm;
#        viii. Run outlier detection algorithm on each cluster (paragraph 2
#              of section 'Detecting Outlier Genes' in original paper)
#
#    Requires protdist version 3.696
#

import sys
import click
import numpy
import operator
import threading
import subprocess
import traceback
import shlex
from os.path import join, basename, isdir, exists, getsize
from os import mkdir

from glob import glob

import skbio.io

class Command(object):
    """Run subprocess commands in a different thread with TIMEOUT option.

    Based on jcollado's solution:
    http://stackoverflow.com/questions/1191374/subprocess-with-timeout/4825933#4825933
    https://gist.github.com/kirpit/1306188
    """
    process = None
    status = None
    output, error = '', ''

    def __init__(self, command):
        if isinstance(command, str):
            command = shlex.split(command)
        self.command = command

    def run(self, timeout=None, **kwargs):
        """ Run a command then return: (status, output, error). """
        def target(**kwargs):
            try:
                self.process = subprocess.Popen(self.command, **kwargs)
                self.output, self.error = self.process.communicate()
                self.status = self.process.returncode
            except subprocess.CalledProcessError:
                self.error = traceback.format_exc()
                self.status = -1
        # default stdout and stderr
        if 'stdout' not in kwargs:
            kwargs['stdout'] = subprocess.PIPE
        if 'stderr' not in kwargs:
            kwargs['stderr'] = subprocess.PIPE
        # thread
        thread = threading.Thread(target=target, kwargs=kwargs)
        thread.start()
        thread.join(timeout)
        if thread.is_alive():
            self.process.terminate()
            thread.join()
        return self.status, self.output, self.error

def hamming(str1, str2):
    """Compute the Hamming distance between two strings.

    Parameters
    ----------
    str1: string
        string
    str2: string
        string
    """
    assert len(str1) == len(str2)
    return sum(map(operator.ne, str1, str2))

def preprocess_data(working_dir,
                    target_proteomes_dir,
                    extensions,
                    verbose=False):
    """ Map each gene to sudo name (ex. 1_1 for species 1, gene 1).

    Parameters
    ----------
    working_dir:  string
        path to working directory
    target_proteomes_dir: string
        path to directory holding proteomes for all target organisms
    extensions: list
        list of extensions for reference proteomes
    verbose: boolean, optional
        output details about the running processes of this function

    Returns
    -------
    gene_map: dictionary
        "two-way" dictionary storing gene names as keys and their pseudo
        names as values, and vica versa
    ref_db: dictionary
        dictionary storing FASTA label as key and sequence as value for the
        reference databases
    species: integer
        the number of species in the reference databases

    Notes
    -----
        This will facilitate easier output comparison and the 10 character
        name limitation in PHYLIP output. This format is limited up to 9999
        species and 99999 genes per species.
    """
    gene_map = {}
    ref_db = {}
    if verbose:
        sys.stdout.write("Target organism\tNumber of genes\n")
    # each file contains genes for species
    files = [f
             for ext in extensions
             for f in glob("%s/*%s" % (target_proteomes_dir, ext))]
    for species, _file in enumerate(files):
        if verbose:
            sys.stdout.write("%s. %s\t" % (
                species+1, basename(_file)))
        for gene, seq in enumerate(skbio.io.read(_file, format='fasta')):
            label = seq.metadata['id']
            ref_db[label] = seq
            sudo_label = "%s_%s" % (species, gene)
            if label in gene_map:
                raise ValueError("Duplicate sequence labels are "
                                 "not allowed: %s" % label)
            gene_map[label] = sudo_label
            gene_map[sudo_label] = label
        if verbose:
            sys.stdout.write("%s\n" % gene)
    return gene_map, ref_db, species+1

def launch_diamond(query_proteome_fp,
                   ref_fp,
                   working_dir,
                   tmp_dir,
                   e_value=10e-20,
                   threads=1,
                   debug=False):
    """ Launch DIAMOND for a query and a reference database of proteomes.

    Parameters
    ----------
    query_proteome_fp: string
      filepath to query proteome
    ref_fp: string
      filepath to reference proteome
    working_dir: string
      working directory path
    tmp_dir:
      temporary working directory for DIAMOND
    e_value: float, optional
      the cutoff E-value for BLASTP results
    threads: integer
      number of threads to use for running DIAMOND BLASTP
    debug: boolean
      if True, run function in debug mode

    Returns
    -------
    out_file_fp: string
      filepath to tabular alignment file output by DIAMOND
    """
    db_file_fp = join(working_dir, "%s" % basename(ref_fp))
    # build DIAMOND database
    makediamonddb_command = ["diamond",
                             "makedb",
                             "--in", ref_fp,
                             "-d", db_file_fp,
                             "--threads", str(threads)]
    proc = subprocess.Popen(makediamonddb_command,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            close_fds=True)
    proc.wait()
    stdout, stderr = proc.communicate()
    if (stderr and debug):
        print("[DEBUG] %s\n" % stderr)

    # launch DIAMOND
    out_file_fp = join(
        working_dir, "%s.daa" % basename(query_proteome_fp))
    diamond_command = ["diamond",
                       "blastp",
                       "-t", tmp_dir,
                       "--db", "%s.dmnd" % db_file_fp,
                       "--query", query_proteome_fp,
                       "--evalue", str(e_value),
                       "--threads", str(threads),
                       "--daa", out_file_fp,
                       "--sensitive"]
    proc = subprocess.Popen(diamond_command,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            close_fds=True)
    proc.wait()
    stdout, stderr = proc.communicate()
    if (stderr and debug):
        print("[DEBUG] %s\n" % stderr)

    # convert output to tab delimited file
    out_file_conv_fp = join(
        working_dir, "%s.m8" % basename(query_proteome_fp))
    diamond_convert_command = ["diamond",
                               "view",
                               "--daa", out_file_fp,
                               "-f", "tab",
                               "-o", out_file_conv_fp]
    proc = subprocess.Popen(diamond_convert_command,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            close_fds=True)
    proc.wait()
    stdout, stderr = proc.communicate()
    if (stderr and debug):
        print("[DEBUG] %s\n" % stderr)

    return out_file_conv_fp

def launch_blast(query_proteome_fp,
                 ref_fp,
                 working_dir,
                 e_value=10e-20,
                 threads=1,
                 debug=False):
    """ Launch BLASTp for a query and a reference database of proteomes.

    Parameters
    ----------
    query_proteome_fp: string
      filepath to query proteome
    ref_fp: string
      filepath to reference proteome
    working_dir: string
      working directory path
    e_value: float, optional
      the cutoff E-value for BLASTP results
    threads: integer
      number of threads to use for running BLASTP
    debug: boolean
      if True, run function in debug mode

    Returns
    -------
    out_file_fp: string
      filepath to tabular alignment file output by
      BLASTP
    """
    db_file_fp = join(working_dir, "%s" % basename(ref_fp))
    # build blast database
    makeblastdb_command = ["makeblastdb",
                           "-in", ref_fp,
                           "-out", db_file_fp,
                           "-dbtype", "prot"]
    proc = subprocess.Popen(makeblastdb_command,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            close_fds=True)
    proc.wait()
    stdout, stderr = proc.communicate()
    if (stderr and debug):
        print("[DEBUG] %s\n" % stderr)

    # launch blast
    out_file_fp = join(
        working_dir, "%s.blast" % basename(query_proteome_fp))
    blastp_command = ["blastp",
                      "-db", db_file_fp,
                      "-query", query_proteome_fp,
                      "-evalue", str(e_value),
                      "-num_threads", str(threads),
                      "-outfmt", "6 std qcovs",
                      "-task", "blastp",
                      "-out", out_file_fp]
    proc = subprocess.Popen(blastp_command,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            close_fds=True)
    proc.wait()
    stdout, stderr = proc.communicate()
    if (stderr and debug):
        print("[DEBUG] %s\n" % stderr)

    return out_file_fp

def parse_blast(alignments_fp,
                hits,
                gene_map,
                debug=False):
    """ Parse BLASTp alignment file into a dictionary.

    Parameters
    ----------
    alignments_fp: string
      filepath to tabular alignment file output by BLASTP
    hits: dictionary
      dictionary storing query (gene) names as keys and the best aligning
      reference sequences as values (one alignment per reference sequence)
    gene_map: dictionary
      "two-way" dictionary storing gene names as keys and their pseudo
      names as values, and vica versa
    debug: boolean
      if True, run function in debug mode

    Notes
    -----
        The keys are the queries and the values are all the reference
        sequences to which the query mapped with E-value cutoff score.
    """
    # read blastp results
    with open(alignments_fp, 'r') as alignments_f:
        for line in alignments_f:
            if debug:
                sys.stdout.write("[DEBUG] %s" % line)
            query, ref = line.split()[:2]
            if query not in hits:
                hits[query] = [ref]
            else:
                # check that the query mapped to a different species
                # since we only want the best homolog per species
                if gene_map[ref].split('_')[0] not in [
                        gene_map[gene].split('_')[0] for gene in hits[query]]:
                    hits[query].append(ref)

def launch_msa(fasta_in_fp,
               clustal_command_fp,
               gene_map,
               ref_db,
               hits,
               query,
               timeout):
    """ Create MSA for all gene othologs using Clustalw.

    Parameters
    ----------
    fasta_in_fp: string
      filepath to FASTA file of protein sequences to use as input to
      Clustalw
    clustal_command_fp: string
      filepath to Clustalw command (interactive)
    gene_map: dictionary
      "two-way" dictionary storing gene names as keys and their pseudo
      names as values, and vica versa
    ref_db: dictionary
      dictionary storing FASTA label as key and sequence as value for the
      reference databases
    hits: dictionary
      dictionary storing query (gene) names as keys and the best aligning
      reference sequences as values (one alignment per reference sequence)
    query: string
      query gene name
    timeout: integer
      number of seconds to allow Clustalw to run before terminating the
      process
    """
    with open(fasta_in_fp, 'w') as in_f:
        for ref in hits[query]:
            in_f.write(">%s\n%s\n" % (gene_map[ref], ref_db[ref]))

    with open(clustal_command_fp, 'r') as clustal_command_f:
        clustalw_command = Command("clustalw")
        status, output, error = clustalw_command.run(
            timeout=timeout,
            stdin=clustal_command_f,
            close_fds=True)
        if status < 0:
            sys.stdout.write(
                "status: %s\noutput: %s\terror: %s\t" % (
                    status, output, error))

def compute_distances(phylip_command_fp,
                      warnings=False):
    """ Compute distances between each pair of sequences in the MSA.

    Parameters
    ----------
    phylip_command_fp: string
      filepath to the PHYLIP command (interactive)
    warnings: boolean, optional
      print warnings output by PHYLIP

    Notes
    -----
        Use PHYLIP's protdist function.
    """
    with open(phylip_command_fp, 'r') as phylip_command_f:
        proc = subprocess.Popen("protdist",
                                stdin=phylip_command_f,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                close_fds=True)
        proc.wait()
        stdout, stderr = proc.communicate()
        if stderr and warnings:
            print(stderr)

def normalize_distances(phylip_fp,
                        full_distance_matrix,
                        num_species,
                        full_distance_matrix_offset,
                        species_set_dict,
                        gene_bitvector_map,
                        debug=False):
    """ Parse and normalize the output file of PHYLIP's protdist function.

    Parameters
    ----------
    phylip_fp: string
        filepath to distance matrix output by PHYLIP's protdist function
    full_distance_matrix: dictionary
        complete distance matrix for pairwise alignments between all species
        for every gene
    num_species: integer
        number of species in the reference database
    full_distance_matrix_offset: integer
        the index offset for elements in full_distance_matrix where to write
        the next array
    species_set_dict: dictionary
        dictionary containing the binary indicator vectors as keys and the
        number of genes with identical species set represented by the binary
        vectors as values
    gene_bitvector_map: list
        list containing the binary indicator vector for each query gene
    debug: boolean
        if True, run function in debug mode

    Notes
    -----
        Parse PHYLIP's protdist output containing the distance matrix, Z-score
        normalize the set of pairwise distances between the gene in a species
        and all other species and stores the results in a separate array.

        Each normalized distance matrix is then sorted by species name and
        added to the complete array storing distance matrices for all genes.
        In addition, a list of missing species (species which did not include
        a certain gene) is also maintained and used for setting nan's in array
        cells which represent those species.

        Below is an example of a parsed distance matrix
        for 3 genes and 3 species:
            0         1         2        (genes)
        0_0 nan       nan       nan
        0_1 0.53099   0.878855  0.83673
        0_2 0.642856  1.083039  1.083039
        1_0 0.300297  0.300297  0.702003
        1_1 nan       nan       nan
        1_2 0.399722  0.379156  0.356543
        2_0 0.53099   0.53099   0.83673
        2_1 0.399722  0.399722  0.356543
        2_2 nan       nan       nan

        (species pairs)

        Example of Z-score normalized distance matrix from
        above:
            0            1           2          (genes)
        0_0 nan          nan         nan
        0_1 -1.40548346  0.83861735  0.56686611
        0_2 -1.41421356  0.70710678  0.70710678
        1_0 -0.70710678 -0.70710678  1.41421356
        1_1 nan          nan         nan
        1_2 1.20493966   0.03869341 -1.24363308
        2_0 -0.70710678 -0.70710678  1.41421356
        2_1 0.70710678   0.70710678 -1.41421356
        2_2 nan          nan         nan

        (species pairs)
    """
    # assume a pairwise alignment exists for all species
    missing_species = [str(x) for x in range(0, num_species)]
    # scan through file and remove species that exist
    # from missing_species list
    if exists(phylip_fp) and getsize(phylip_fp) > 0:
        with open(phylip_fp, 'r') as phylip_f:
            next(phylip_f)
            for line in phylip_f:
                if not line.startswith(' '):
                    species = line.split()[0].split('_')[0]
                    missing_species.remove(species)
    else:
        raise ValueError('%s does not exist or is empty' % phylip_fp)

    # scan through file again, collecting alignment
    # distances
    orig_order_labels = []
    p = numpy.empty(shape=(num_species, num_species))
    p.fill(numpy.nan)
    idx = 0
    with open(phylip_fp, 'r') as phylip_f:
        alignment_list = []
        # skip first line containing number of lines in
        # the file
        next(phylip_f)
        for line in phylip_f:
            if debug:
                sys.stdout.write("[DEBUG] %s" % line)
            alignment_dist = line.strip().split()
            if line.startswith(' '):
                alignment_list.extend(alignment_dist)
            else:
                # new species alignment pairs
                if alignment_list:
                    for i in range(0, len(missing_species)):
                        alignment_list.append(None)
                    a = numpy.asarray(alignment_list[1:], dtype=float)
                    a[idx] = numpy.nan
                    p[idx] = (a - numpy.nanmean(a)) / numpy.nanstd(a)
                    idx += 1
                    orig_order_labels.append(alignment_list[0])
                alignment_list = alignment_dist

    # add distance on final line
    for i in range(0, len(missing_species)):
        alignment_list.append(None)
    a = numpy.asarray(alignment_list[1:], dtype=float)
    a[idx] = numpy.nan
    p[idx] = (a - numpy.nanmean(a)) / numpy.nanstd(a)
    orig_order_labels.append(alignment_list[0])

    # add the missing species names to the labels array
    bitvector_gene = 'I' * num_species
    for species in missing_species:
        orig_order_labels.append("%s_X" % species)
        # indicate missing gene for current species
        x = list(bitvector_gene)
        x[int(species)] = 'O'
        bitvector_gene = ''.join(x)

    # update species set counts
    if bitvector_gene not in species_set_dict:
        species_set_dict[bitvector_gene] = 1
    else:
        species_set_dict[bitvector_gene] += 1

    gene_bitvector_map[full_distance_matrix_offset] = bitvector_gene

    # sort the distance matrix based on species names (S1, S2, S3 ..)
    # in order to be consistent across all gene families
    x = numpy.argsort(numpy.array(orig_order_labels))

    # re-order rows and columns by ordered species name (0,1,2 ..)
    p2 = numpy.zeros(shape=(num_species, num_species))
    for idx_a, arr in enumerate(p):
        t = numpy.zeros(shape=num_species)
        for idx_b, el in enumerate(arr):
            t[x[idx_b]] = el
        p2[x[idx_a]] = t
    del p

    # add normalized distance matrix for current gene
    # to full distance matrix
    full_distance_matrix[full_distance_matrix_offset] = p2

def cluster_distances(species_set_dict,
                      species_set_size,
                      hamming_distance):
    """ Hamming distance clustering algorithm

    Parameters
    ----------
    species_set_dict: dictionary
        dictionary containing the binary indicator vectors as
        keys and the number of genes with identical species
        set represented by the binary vectors as values
    species_set_size: integer
        threshold number of genes in a species set to
        allow it to form a core cluster
    hamming_distance: integer
        maximum number of mismatches between two binary
        indicator vectors (ex. IIII and I0II) for the
        genes in a candidate vector to be merged into the
        core cluster

    Returns
    -------
    gene_clusters_list: list of tuples
        list of tuples containing core species sets and all belonging species
        sets (determined by the Hamming distance clustering algorithm)

    Notes
    -----
        Cluster gene families by species with detectable orthologs in exactly
        the same subset of the considered species.

        Ex. Assume we have 4 genes and 5 species with the following distance
        matrix:

            0             1             2             3
        0_0 nan           nan           nan           nan
        0_1 -1.59564844   -1.388031632  -0.9634704748 -1.342272936
        0_2 -0.4259542606 nan           0.7035215923  1.223837777
        0_3 -1.55041393   -1.51499567   -0.9634704748 -1.330178178
        0_4 -0.3659762821 0.8346037464  0.6565725705  1.15274682
        ..
        ..

        There are two binary indicator vectors to represent the species
        present in the four genes: IIIII (gene 0, 2 and 3), II0II (gene 1). If
        the core set threshold was 3, then there would be 1 core species set
        represented by IIIII.
    """
    sorted_species_set = sorted(list(species_set_dict.items()),
                                key=operator.itemgetter(1), reverse=True)
    # determine core clusters (initial species sets with more than
    # species_set_size genes)
    gene_clusters_list = []
    # if the largest species set contains less than threshold
    # (species_set_size) elements, set the only core cluster to the largest
    # species set
    if sorted_species_set[0][1] < species_set_size:
        cluster_core = (sorted_species_set[0][0], [])
        gene_clusters_list = []
    for bitvector in sorted_species_set:
        if bitvector[1] >= species_set_size:
            cluster_core = (bitvector[0], [])
            gene_clusters_list.append(cluster_core)
    # assign species sets with fewer than species_set_size species to core
    # clusters if the Hamming distance between the two bitvectors is less than
    # hamming_distance
    species_set_assigned = []
    for idx, cluster_core in enumerate(gene_clusters_list):
        for bitvector in sorted_species_set:
            bv = bitvector[0]
            if (bv not in species_set_assigned and
                    hamming(cluster_core[0], bv) <= hamming_distance):
                gene_clusters_list[idx][1].append(bv)
                species_set_assigned.append(bv)
    # assign the remaining species sets to the cluster with the closest core
    # Hamming distance
    for bitvector in sorted_species_set:
        bv = bitvector[0]
        if bv not in species_set_assigned:
            min_hamming_cluster = -1
            min_hamming_distance = sys.maxsize
            # find cluster core with smallest Hamming distance to species set
            for idx, cluster_core in enumerate(gene_clusters_list):
                dist = hamming(cluster_core[0], bv)
                if dist < min_hamming_distance:
                    min_hamming_distance = dist
                    min_hamming_cluster = idx
            if min_hamming_cluster >= 0:
                gene_clusters_list[min_hamming_cluster][1].append(bv)

    return gene_clusters_list

def detect_outlier_genes(species_set,
                         gene_bitvector_map,
                         full_distance_matrix,
                         stdev_offset,
                         outlier_hgt,
                         num_species,
                         total_genes,
                         debug=False):
    """ Detect outlier genes.

    Parameters
    ----------
    species_set: list
        list of bitvectors representing species clusters to use in detecting
        outlier genes
    gene_bitvector_map: list
        list containing the binary indicator vector for each query gene
    full_distance_matrix: dictionary
        complete distance matrix for pairwise alignments between all species
        for every gene
    stdev_offset: integer
        the number of standard deviations a gene's normalized distance is from
        the mean to identify it as an outlier for a species pair
    outlier_hgt: float
        the fraction (value between (0,1]) of normalized pairwise distances
        over all species-pair vectors belonging to the same gene that are
        z-score standard deviations from the mean
    num_species: integer
        number of species in the reference database
    total_genes: integer
        total number of genes in the query genome with at least
        min_num_homologs (determined by BLAST search)
    debug: boolean
        if True, run function in debug mode

    Returns
    -------
    outlier_genes: set
        set of atypical genes

    Notes
    -----
        Algorithm described in section "Detecting `Outlier' Genes" of the Wei.
        X et al. paper. The full distance matrix is represented in the format:

        full_distance_matrix[#genes][#species][#species] =
        [[[0_0, 0_1, 0_2, .., 0_n]
          [1_0, 1_1, 1_2, .., 1_n]
          ..
          [n_0, n_1, n_2, .., n_n]]

         [[0_0, 0_1, 0_2, .., 0_n]
          [1_0, 1_1, 1_2, .., 1_n]
          ..
          [n_0, n_1, n_2, .., n_n]]

          ..
         [[0_0, 0_1, 0_2, .., 0_n]
          [1_0, 1_1, 1_2, .., 1_n]
          ..
          [n_0, n_1, n_2, .., n_n]]]

        The mean and standard deviation are computed for each species pair
        including all genes.
    """
    numpy.around(full_distance_matrix, decimals=5, out=full_distance_matrix)
    outlier_flag_matrix = numpy.zeros(
        shape=(total_genes, num_species, num_species), dtype=bool)
    distance_vector = numpy.zeros(total_genes)
    if debug:
        sys.stdout.write("[DEBUG] species_species\t")
        for k in range(total_genes):
            sys.stdout.write("gene # %s".ljust(12) % k)
        sys.stdout.write("[low_bound, up_bound]\n")
    for i in range(num_species):
        for j in range(num_species):
            if i != j:
                for k in range(total_genes):
                    distance_vector[k] = full_distance_matrix[k][i][j]
                mean = numpy.nanmean(distance_vector)
                stdev = numpy.nanstd(distance_vector)
                low_bound = round(mean - stdev_offset*stdev, 5)
                up_bound = round(mean + stdev_offset*stdev, 5)
                if debug:
                    sys.stdout.write("[DEBUG] %s_%s\t".ljust(20) % (i, j))
                for k, distance in enumerate(distance_vector):
                    spaces = "".ljust(2)
                    if distance < 0:
                        spaces = "".ljust(1)
                    if (distance != numpy.nan and
                       ((distance < low_bound) or (distance > up_bound))):
                        outlier_flag_matrix[k][i][j] = 1
                        if debug:
                            sys.stdout.write(
                                "%s\033[92m%s\033[0m" % (spaces, distance))
                    elif debug:
                        sys.stdout.write("%s%s" % (spaces, distance))
                if debug:
                    sys.stdout.write("\t[%s, %s]\n" % (low_bound, up_bound))

    # traverse outlier_matrix by gene and count the number of outlier
    # distances by species
    outlier_count_matrix = numpy.zeros(
        shape=(total_genes, num_species), dtype=int)
    for i in range(total_genes):
        for j in range(num_species):
            for k in range(num_species):
                if outlier_flag_matrix[i][j][k]:
                    outlier_count_matrix[i][k] += 1

    # if number of outlier distances exceeds threshold, label gene as outlier
    outlier_genes = set()
    for i in range(total_genes):
        for j in range(num_species):
            if outlier_count_matrix[i][j] > num_species*outlier_hgt:
                outlier_genes.add(i)

    return outlier_genes

def output_full_matrix(matrix, num_species):
    """ Output distance matrix to stdout
    """
    for i in range(num_species):
        for j in range(num_species):
            # for gene number
            for k in range(len(matrix)):
                sys.stdout.write("%s\t" % matrix[k][i][j])
            sys.stdout.write("\n")

def distance_method(query_proteome_fp,
                    target_proteomes_dir,
                    working_dir,
                    output_hgt_fp,
                    align_software,
                    tabular_alignments_fp=None,
                    ext=['fa', 'fasta', 'faa'],
                    min_num_homologs=3,
                    e_value=10e-20,
                    threads=1,
                    stdev_offset=2.326,
                    outlier_hgt=0.5,
                    species_set_size=30,
                    hamming_distance=2,
                    verbose=False,
                    debug=False,
                    warnings=False,
                    timeout=120):
    """ Run Distance Method algorithm

    Parameters
    ----------
    query_proteome_fp: string
        filepath to query proteome
    target_proteomes_dir: string
        dirpath to target proteomes
    working_dir: string
        dirpath to working directory
    output_hgt_fp: string
        filepath to output file for storing detected HGTs
    align_software: string
        software to use for sequence alignment (BLAST or DIAMOND)
    tabular_alignments_fp: string, optional
        filepath to tabular sequence alignments
    ext: list, optional
        list of file extensions to open in the target proteomes directory
    min_num_homologs: integer, optional
        the mininum number of homologs (determined by BLAST search)
        for each gene to test
    e_value: float, optional
        the E-value cutoff to identify orthologous genes using BLASTP
    threads: integer, optional
        number of threads to use for sequence alignment
    stdev_offset: float, optional
        the number of standard deviations a gene's normalized distance
        is from the mean to identify it as an outlier for a species pair
    outlier_hgt: float, optional
        the fraction (value between (0,1]) of normalized pairwise distances
        over all species-pair vectors belonging to the same gene that are
        z-score standard deviations from the mean
    species_set_size: integer, optional
        threshold number of genes to consider a species set large (a species
        set is a set of genes whose orthologs are detectable in exactly the
        same subset of the considered species)
    hamming_distance: integer, optional
        distance between two binary vectors indicating the species in which
        the corresponding ortholog gene appears
    verbose: boolean, optional
        if True, run in verbose mode
    debug: boolean, optional
        if True, run in debug mode
    warnings: boolean, optional
        if True, output warnings
    timeout: integer, optional
        number of seconds to allow Clustalw to run per call
    """
    if verbose:
        sys.stdout.write(
            "Begin whole-genome HGT detection using the Distance method.\n\n")
        sys.stdout.write("Query genome: %s\n" % query_proteome_fp)

    extensions = set(['fa', 'fasta', 'faa'])
    extensions.update(ext)

    # create working directory if doesn't exist
    if not isdir(working_dir):
        mkdir(working_dir)

    gene_map, ref_db, num_species = preprocess_data(
        working_dir=working_dir,
        target_proteomes_dir=target_proteomes_dir,
        extensions=extensions,
        verbose=verbose)

    if debug:
        sys.stdout.write("\n[DEBUG] gene map:\n")
        for gene in gene_map:
            sys.stdout.write("[DEBUG] %s: %s\n" % (gene, gene_map[gene]))

    if verbose:
        sys.stdout.write("\nRunning BLASTp ..\n")
    hits = {}

    # tabular alignments provided
    if tabular_alignments_fp is not None:
        # generate a dictionary of orthologous genes
        parse_blast(alignments_fp=tabular_alignments_fp,
                    hits=hits,
                    gene_map=gene_map,
                    debug=debug)
    # tabular alignments to be created
    else:
        files = [f
                 for e in extensions
                 for f in glob("%s/*%s" % (target_proteomes_dir, e))]
        for _file in files:
            # launch BLASTp
            if align_software == "blast":
                alignments_fp = launch_blast(
                    query_proteome_fp=query_proteome_fp,
                    ref_fp=_file,
                    working_dir=working_dir,
                    e_value=e_value,
                    threads=threads,
                    debug=debug)
            elif align_software == "diamond":
                alignments_fp = launch_diamond(
                    query_proteome_fp=query_proteome_fp,
                    ref_fp=_file,
                    working_dir=working_dir,
                    tmp_dir=working_dir,
                    e_value=e_value,
                    threads=threads,
                    debug=debug)
            else:
                raise ValueError(
                    "Software not supported: %s" % align_software)

            # generate a dictionary of orthologous genes
            parse_blast(alignments_fp=alignments_fp,
                        hits=hits,
                        gene_map=gene_map,
                        debug=debug)

    # keep only genes with >= min_num_homologs
    hits_min_num_homologs = {}
    max_homologs = 0
    for query in hits:
        len_hits = len(hits[query])
        if query in hits[query]:
            len_hits -= 1
        if len_hits >= min_num_homologs:
            if query in hits_min_num_homologs:
                raise ValueError("Duplicate gene names found: %s" % query)
            hits_min_num_homologs[query] = hits[query]
            if len_hits > max_homologs:
                max_homologs = len_hits
    hits.clear()

    if verbose:
        sys.stdout.write(
            "Total number of orthologous gene families with at "
            "least %s genes: %s\n" % (
                min_num_homologs, len(hits_min_num_homologs)))
    if debug:
        sys.stdout.write("[DEBUG] Blast matches:\n")
        for query in hits_min_num_homologs:
            sys.stdout.write(
                "[DEBUG] %s: %s\n" % (query, hits_min_num_homologs[query]))
    # generate command for CLUSTALW
    phy_msa_fp = join(working_dir, "msa.phy")
    open(phy_msa_fp, 'a').close()
    dnd_msa_fp = join(working_dir, "msa.dnd")
    open(dnd_msa_fp, 'a').close()
    phylip_fp = join(working_dir, "msa.dis")
    open(phylip_fp, 'a').close()
    # create fasta file for each gene family and run CLUSTALW
    fasta_in_fp = join(working_dir, "input.faa")
    clustal_command_fp = join(working_dir, "clustal_command.txt")
    with open(clustal_command_fp, 'w') as clustal_command_f:
        clustal_command_f.write(
            '1\n%s\n2\n9\n1\n4\n\n1\n%s\n%s\nX\n\nX\n' % (
                fasta_in_fp, phy_msa_fp, dnd_msa_fp))
    phylip_command_fp = join(working_dir, "phylip_command.txt")
    with open(phylip_command_fp, 'w') as phylip_command_f:
        phylip_command_f.write('%s\nF\n%s\nR\nY\n' % (phy_msa_fp, phylip_fp))

    total_genes = len(hits_min_num_homologs)
    if verbose:
        sys.stdout.write("\nRunning CLUSTALW and PROTDIST ..\n")
    if max_homologs > num_species:
        raise ValueError(
            "max_homologs > num_species: %s > %s " % (
                max_homologs, num_species))
    # distance matrix containing distances between all ortholog genes
    full_distance_matrix = numpy.zeros(
        shape=(total_genes, num_species, num_species), dtype=float)
    # dictionary to store all subsets of orthologs (keys) and
    # their number of occurrences (values) (maximum occurrences
    # is equal to the number of genes)
    species_set_dict = {}
    gene_bitvector_map = {}
    gene_id = {}
    for i, query in enumerate(hits_min_num_homologs):
        if verbose:
            print("Computing MSA and distances for gene %s .. (%s/%s)" % (
                query, i+1, total_genes))
        gene_id[i] = query
        # generate a multiple sequence alignment
        # for each orthologous gene family
        launch_msa(fasta_in_fp=fasta_in_fp,
                   clustal_command_fp=clustal_command_fp,
                   ref_db=ref_db,
                   gene_map=gene_map,
                   hits=hits_min_num_homologs,
                   query=query,
                   timeout=timeout)

        # compute distances between each pair of sequences in MSA
        compute_distances(phylip_command_fp=phylip_command_fp,
                          warnings=warnings)

        # Z-score normalize distance matrix and add results
        # to full distance matrix (for all genes)
        normalize_distances(phylip_fp=phylip_fp,
                            full_distance_matrix=full_distance_matrix,
                            num_species=num_species,
                            full_distance_matrix_offset=i,
                            species_set_dict=species_set_dict,
                            gene_bitvector_map=gene_bitvector_map,
                            debug=debug)

    # output_full_matrix(full_distance_matrix, num_species)

    # cluster gene families by species
    gene_clusters_dict = cluster_distances(
        species_set_dict=species_set_dict,
        species_set_size=species_set_size,
        hamming_distance=hamming_distance)

    # detect outlier genes per core cluster of genes
    with open(output_hgt_fp, 'w') as output_hgt_f:
        output_hgt_f.write("\n# Candidate HGT genes: \n")
        for core_cluster in gene_clusters_dict:
            outlier_genes = detect_outlier_genes(
                species_set=gene_clusters_dict[core_cluster],
                gene_bitvector_map=gene_bitvector_map,
                full_distance_matrix=full_distance_matrix,
                stdev_offset=stdev_offset,
                outlier_hgt=outlier_hgt,
                num_species=num_species,
                total_genes=total_genes,
                debug=debug)

            if outlier_genes:
                for gene in outlier_genes:
                    output_hgt_f("%s\n" % gene_id[gene])

    # output_full_matrix(outlier_genes, num_species)

@click.command()
@click.argument('query-proteome-fp', required=True,
                type=click.Path(resolve_path=True, readable=True, exists=True,
                                file_okay=True))
@click.argument('target-proteomes-dir', required=True,
                type=click.Path(resolve_path=True, readable=True, exists=True,
                                file_okay=True))
@click.argument('working-dir', required=True,
                type=click.Path(resolve_path=True, readable=True, exists=False,
                                file_okay=True))
@click.argument('output-hgt-fp', required=True,
                type=click.Path(resolve_path=True, readable=True, exists=False,
                                file_okay=True))
@click.option('--align-software', type=click.Choice(['diamond', 'blast']),
              required=False, default=['diamond'], show_default=True,
              help="Software to use for blasting sequences")
@click.option('--tabular-alignments-fp', required=False,
              type=click.Path(resolve_path=True, readable=True, exists=False,
                              file_okay=True),
              help="Tabular alignments in m6 format (output from BLAST or "
                   "DIAMOND)")
@click.option('--ext', multiple=True, type=str, required=False,
              default=['fa', 'fasta', 'faa'], show_default=True,
              help="File extensions of target proteomes (multiple extensions "
                   "can be given by calling --ext ext1 --ext ext2)")
@click.option('--min-num-homologs', type=int, required=False, default=3,
              show_default=True, help="The mininum number of homologs "
                                      "(determined by BLAST search) for each "
                                      "gene to test")
@click.option('--e-value', type=float, required=False, default=10e-20,
              show_default=True, help="The E-value cutoff to identify "
                                      "orthologous genes using BLASTP")
@click.option('--threads', type=int, required=False, default=1,
              show_default=True, help="Number of threads to use")
@click.option('--stdev-offset', type=float, required=False, default=2.326,
              show_default=True, help="The number of standard deviations a "
                                      "gene's normalized distance is from "
                                      "the mean to identify it as an outlier "
                                      "for a species pair")
@click.option('--outlier-hgt', type=float, default=0.5, show_default=True,
              required=False, help="The fraction (value between (0,1]) of "
                                   "normalized pairwise distances over all "
                                   "species-pair vectors belonging to the "
                                   "same gene that are z-score standard "
                                   "deviations from the mean")
@click.option('--species-set-size', type=int, required=False, default=30,
              show_default=True, help="Threshold number of genes to consider "
                                      "a species set large (a species set is "
                                      "a set of genes whose orthologs are "
                                      "detectable in exactly the same subset "
                                      "of the considered species)")
@click.option('--hamming-distance', type=int, required=False, default=2,
              show_default=True, help="Distance between two binary vectors "
                                      "indicating the species in which the "
                                      "corresponding ortholog gene appears")
@click.option('--verbose', type=bool, required=False, default=False,
              show_default=True, help="Run in verbose mode")
@click.option('--debug', type=bool, required=False, default=False,
              show_default=True, help="Run in debug mode")
@click.option('--warnings', type=bool, required=False, default=False,
              show_default=True, help="Print program warnings")
@click.option('--timeout', type=int, required=False, default=120,
              show_default=True, help="Number of seconds to allow Clustalw "
                                      "to run per call")
def distance_method_main(query_proteome_fp,
                         target_proteomes_dir,
                         working_dir,
                         output_hgt_fp,
                         align_software,
                         tabular_alignments_fp,
                         ext,
                         min_num_homologs,
                         e_value,
                         threads,
                         stdev_offset,
                         outlier_hgt,
                         species_set_size,
                         hamming_distance,
                         verbose,
                         debug,
                         warnings,
                         timeout):
    """ Run the Distance-Method HGT detection algorithm.
    """
    distance_method(query_proteome_fp=query_proteome_fp,
                    target_proteomes_dir=target_proteomes_dir,
                    working_dir=working_dir,
                    output_hgt_fp=output_hgt_fp,
                    align_software=align_software,
                    tabular_alignments_fp=tabular_alignments_fp,
                    ext=ext,
                    min_num_homologs=min_num_homologs,
                    e_value=e_value,
                    threads=threads,
                    stdev_offset=stdev_offset,
                    outlier_hgt=outlier_hgt,
                    species_set_size=species_set_size,
                    hamming_distance=hamming_distance,
                    verbose=verbose,
                    debug=debug,
                    warnings=warnings,
                    timeout=timeout)

if __name__ == "__main__":
    distance_method_main()


"""

This module is used for constructing a `PyMC3' model of the power spectrum using
the outputs from asy_peakbag as priors.

"""

import numpy as np
import pymc3 as pm
import warnings
from .plotting import plotting

class peakbag(plotting):
    """ Class for the final peakbagging.

    This class is used after getting the frequency intervals from asy_peakbag,
    that include the $l=0,2$ mode pairs.

    The

    Examples
    --------
    Using peakbag from the star class instance (recommended)

    >>> st = pbjam.star(ID='KIC4448777', pg=pg, numax=[220.0, 3.0],
                           dnu=[16.97, 0.01], teff=[4750, 100],
                           bp_rp = [1.34, 0.01])
    >>> st.run_kde()
    >>> st.run_asy_peakbag(norders=7)
    >>> st.run_peakbag()

    Using peakbag on it's own. Requires output from `asy_peakbag'.

    >>> pb = pbjam.peakbag(st.asy_fit)
    >>> pb()

    Parameters
    ----------
    f : float, array
        Array of frequency bins of the spectrum (muHz). Truncated to the range
        around numax.
    snr : float, array
        Array of SNR values for the frequency bins in f (dimensionless).
    asy_fit : asy_fit
        The result from the asy_peakbag method.
    init : bool
        If true runs make_start and trim_ladder to prepare starting
        guesses for sampling and transforms the data onto a ladder.

    Attributes
    ----------
    f : float, ndarray
        Array of frequency bins of the spectrum (muHz). Truncated to the range
        around numax.
    snr : float, ndarray
        Array of SNR values for the frequency bins in f (dimensionless).
    asy_fit : asy_fit
        The result from the asy_peakbag method.
        This is a dictionary of 'modeID' and 'summary'.
        'modeID' is a DataFrame with a list of modes and basic properties.
        'summary' are summary statistics from the asymptotic_fit.
        See asy_peakbag asymptotic_fit for more details.

    """

    def __init__(self, starinst, init=True, path=None,  verbose=False):

        self.pg = starinst.pg
        self.f = starinst.f
        self.s = starinst.s
        self.asy_fit = starinst.asy_fit
        self.norders = self.asy_fit.norders
        if init:
            self.make_start()
            self.trim_ladder(verbose=verbose)
        self.gp0 = [] # Used for gp linewidth info.

        starinst.references._addRef('pymc3')

        starinst.peakbag = self

    def make_start(self):
        """ Set the starting model for peakbag

        Function uses the result of the asymptotic peakbagging and builds a
        dictionary of starting values for the peakbagging methods.

        """

        idxl0 = self.asy_fit.modeID.ell == 0
        idxl2 = self.asy_fit.modeID.ell == 2

        l0 = self.asy_fit.modeID.loc[idxl0, 'nu_med'].values.flatten()
        l2 = self.asy_fit.modeID.loc[idxl2, 'nu_med'].values.flatten()

        l0, l2 = self.remove_outsiders(l0, l2)

        width = 10**(np.ones(len(l0)) * self.asy_fit.summary.loc['mode_width', 'mean']).flatten()
        height =  (10**self.asy_fit.summary.loc['env_height', 'mean'] * \
                 np.exp(-0.5 * (l0 - 10**self.asy_fit.summary.loc['numax', 'mean'])**2 /
                 (10**self.asy_fit.summary.loc['env_width', 'mean'])**2)).flatten()
        back = np.ones(len(l0))

        self.parnames = ['l0', 'l2', 'width0', 'width2', 'height0', 'height2',
                         'back']

        pars = [l0, l2, width, width, height, 0.7*height, back]

        self.start ={x:y for x,y in zip(self.parnames, pars)}

        self.n = np.linspace(0.0, 1.0, len(self.start['l0']))[:, None]

    def remove_outsiders(self, l0, l2):
        """ Drop outliers

        Drops modes where the guess frequency is outside of the supplied
        frequency range.

        Parameters
        ----------

        l0 : ndarray
            Array of l0 mode frequencies
        l2 : ndarray
            Array of l2 mode frequencies

        """

        sel = np.where(np.logical_and(l0 < self.f.max(), l0 > self.f.min()))
        return l0[sel], l2[sel]

    def trim_ladder(self, lw_fac=10, extra=0.01, verbose=False):
        """ Turns mode frequencies into list of pairs

        This function turns the list of mode frequencies into pairs and then
        selects only the pairs in the ladder that have modes that are to be fit.

        Each pair is constructed so that the central frequency is
        the mid point between the l=0 and l=2 modes as determined by the
        information in the asy_fit dictionary.

        Parameters
        ----------
        lw_fac: float
            The factor by which the mode line width is multiplied in order
            to contribute to the pair width.
        extra: float
            The factor by which dnu is multiplied in order to contribute to
            the pair width.

        """

        d02 = 10**self.asy_fit.summary.loc['d02', 'mean']
        d02_lw = d02 + lw_fac * 10**self.asy_fit.summary.loc['mode_width', 'mean']
        w = d02_lw + (extra * 10**self.asy_fit.summary.loc['dnu', 'mean'])
        bw = self.f[1] - self.f[0]
        w /= bw
        if verbose:
            print(f'w = {int(w)}')
            print(f'bw = {bw}')
        ladder_trim_f = np.zeros([len(self.start['l0']), int(w)])
        ladder_trim_s = np.zeros([len(self.start['l0']), int(w)])
        for idx, freq in enumerate(self.start['l0']):
            loc_mid_02 = np.argmin(np.abs(self.f - (freq - d02/2.0)))
            if loc_mid_02 == 0:
                warnings.warn('Did not find optimal pair location')
            if verbose:
                print(f'loc_mid_02 = {loc_mid_02}')
                print(f'w/2 = {int(w/2)}')
            ladder_trim_f[idx, :] = \
                self.f[loc_mid_02 - int(w/2): loc_mid_02 - int(w/2) + int(w)]
            ladder_trim_s[idx, :] = \
                self.s[loc_mid_02 - int(w/2): loc_mid_02 - int(w/2) + int(w) ]
        self.ladder_f = ladder_trim_f
        self.ladder_s = ladder_trim_s

    def lor(self, freq, w, h):
        """ Simple Lorentzian profile

        Calculates N Lorentzian profiles, where N is the number of pairs in the
        frequency list.

        Parameters
        ----------
        freq : float, ndarray
            Central frequencies the N Lorentzians
        w : float, ndarray
            Widths of the N Lorentzians
        h : float, ndarray
            Heights of the N Lorentzians

        Returns
        -------
        lors : ndarray
           A list containing one Lorentzian per pair.

        """

        norm = 1.0 + 4.0 / w**2 * (self.ladder_f.T - freq)**2

        return h / norm

    def model(self, l0, l2, width0, width2, height0, height2, back):
        """
        Calcuates a simple model of a flat backgroud plus two lorentzians
        for each of the N pairs in the list of frequencies under consideration.

        Parameters
        ----------
        l0 : ndarray
            Array of length N, of the l=0 mode frequencies.
        l2 : ndarray
            Array of length N, of the l=2 mode frequencies.
        width0 : ndarray
            Array of length N, of the l=0 mode widths.
        width2 : ndarray
            Array of length N, of the l=2 mode widths.
        height0 : ndarray
            Array of length N, of the l=0 mode heights.
        height2 : ndarray
            Array of length N, of the l=2 mode heights.
        back : ndarray
            Array of length N, of the background levels.

        Returns
        -------
        mod : ndarray
            A 2D array (or 'ladder') containing the calculated models for each
            of the N pairs.

        """

        mod = np.ones(self.ladder_f.shape).T * back
        mod += self.lor(l0, width0, height0)
        mod += self.lor(l2, width2, height2)
        return mod.T

    def init_model(self, model_type):
        """ Initialize the pymc3 model for peakbag

        Sets up the pymc3 model to sample, to perform the final peakbagging.

        Two treatements of the mode widths are available, the default
        independent mode widths for each pair, or modeling the mode widths as
        a function of freqeuency as a Gaussian Process.

        Parameters
        ----------
        model_type : str
            Model choice for the mode widths. The default is to treat the all
            mode widths independently. Alternatively they can be modeled as
            a GP.

        """

        self.pm_model = pm.Model()

        dnu = 10**self.asy_fit.summary.loc['dnu', 'mean']
        dnu_fac = 0.03 # Prior on mode frequency has width 3% of Dnu.
        height_fac = 0.4 # Lognorrmal prior on height has std=0.4.
        width_fac = 1.0 # Lognorrmal prior on width has std=1.0.
        back_fac = 0.5 # Lognorrmal prior on back has std=0.5.
        N = len(self.start['l2'])

        with self.pm_model:

            if model_type != 'model_gp':
                if model_type != 'simple': # defaults to simple if bad input
                    warnings.warn('Model not defined - using simple model')
                width0 = pm.Lognormal('width0', mu=np.log(self.start['width0']),
                                  sigma=width_fac, shape=N)
                width2 = pm.Lognormal('width2', mu=np.log(self.start['width2']),
                                  sigma=width_fac, shape=N)

                self.init_sampler = 'adapt_diag'
                self.target_accept = 0.9

            elif model_type == 'model_gp':
                warnings.warn('This model is developmental - use carefully')
                # Place a GP over the l=0 mode widths ...
                m0 = pm.Normal('gradient0', 0, 10)
                c0 = pm.Normal('intercept0', 0, 10)
                sigma0 = pm.Lognormal('sigma0', np.log(1.0), 1.0)
                ls = pm.Lognormal('ls', np.log(0.3), 1.0)
                mean_func0 = pm.gp.mean.Linear(coeffs=m0, intercept=c0)
                cov_func0 = sigma0 * pm.gp.cov.ExpQuad(1, ls=ls)
                self.gp0 = pm.gp.Latent(cov_func=cov_func0, mean_func=mean_func0)
                ln_width0 = self.gp0.prior('ln_width0', X=self.n)
                width0 = pm.Deterministic('width0', pm.math.exp(ln_width0))
                # and on the l=2 mode widths
                m2 = pm.Normal('gradient2', 0, 10)
                c2 = pm.Normal('intercept2', 0, 10)
                sigma2 = pm.Lognormal('sigma2', np.log(1.0), 1.0)
                mean_func2 = pm.gp.mean.Linear(coeffs=m2, intercept=c2)
                cov_func2 = sigma2 * pm.gp.cov.ExpQuad(1, ls=ls)
                self.gp2 = pm.gp.Latent(cov_func=cov_func2, mean_func=mean_func2)
                ln_width2 = self.gp2.prior('ln_width2', X=self.n)
                width2 = pm.Deterministic('width2', pm.math.exp(ln_width2))

                self.init_sampler = 'advi+adapt_diag'
                self.target_accept = 0.99

            l0 = pm.Normal('l0', self.start['l0'], dnu*dnu_fac, shape=N)

            l2 = pm.Normal('l2', self.start['l2'], dnu*dnu_fac, shape=N)

            height0 = pm.Lognormal('height0', mu=np.log(self.start['height0']),
                                    sigma=height_fac, shape=N)
            height2 = pm.Lognormal('height2', mu=np.log(self.start['height2']),
                                    sigma=height_fac, shape=N)
            back = pm.Lognormal('back', mu=np.log(1.0), sigma=back_fac, shape=N)

            limit = self.model(l0, l2, width0, width2, height0, height2, back)

            pm.Gamma('yobs', alpha=1, beta=1.0/limit, observed=self.ladder_s)

    def _addPPRatio(self):
        """ Add the prior/posterior width ratio to summary

        Computes the ratio of the prior width and the posterior width. This is a
        quantity which indicates which probability predominantly informs the
        resulting mode frequency. If the ratio is < 1 the prior dominates, and
        vice versa for ratios > 1.

        No cut-off is made based on this ratio, it is merely to inform the user.

        """

        dnu = 10**self.asy_fit.summary.loc['dnu', 'mean']

        log_ppr = np.log10((0.03*dnu))-np.log10(self.summary['sd'])

        idx = np.array(['l' in name for name in self.summary.index], dtype = 'bool')

        self.summary['log_ppr'] = np.nan

        self.summary.at[idx, 'log_ppr'] = log_ppr[idx]

    def __call__(self, model_type='simple', tune=1500, nthreads=1, maxiter=4,
                     advi=False):
        """ Perform all the steps in peakbag.

        Initializes and samples the `PyMC3' model that is set up using the
        outputs from asy_peakbag as priors.

        Parameters
        ----------
        model_type : str
            Defaults to 'simple'.
            Can be either 'simple' or 'model_gp' which sets the type of model
            to be fitted to the data.
        tune : int, optional
            Numer of tuning steps passed to pym3.sample. Default is 1500.
        nthreads : int, optional
            Number of cores to use - passed to pym3.sample. Default is 1.
        maxiter : int, optional
            Number of times to attempt to reach convergence. Default is 4.
        advi : bool, optional
            Whether or not to fit using the fullrank_advi option in pymc3.
            Default is False.

        """

        self.init_model(model_type=model_type)

        # REMOVE THIS WHEN pymc3 v3.8 is a bit older.
        try:
            rhatfunc = pm.diagnostics.gelman_rubin
            warnings.warn('pymc3.diagnostics.gelman_rubin is depcrecated; upgrade pymc3 to v3.8 or newer.', DeprecationWarning)
        except:
            rhatfunc = pm.stats.rhat

        if advi:
            with self.pm_model:
                cb = pm.callbacks.CheckParametersConvergence(every=1000,
                                                             diff='absolute',
                                                             tolerance=0.01)

                mean_field = pm.fit(n=200000, method='fullrank_advi',
                                    start=self.start,
                                    callbacks=[cb])
                self.traces = mean_field.sample(1000)
        else:
            Rhat_max = 10
            niter = 1
            while Rhat_max > 1.05:
                if niter > maxiter:
                    warnings.warn('Did not converge!')
                    break
                with self.pm_model:
                    self.traces = pm.sample(tune=tune * niter, cores=nthreads,
                                             start=self.start,
                                             init=self.init_sampler,
                                             target_accept=self.target_accept,
                                             progressbar=False)
                Rhat_max = np.max([v.max() for k, v in rhatfunc(self.traces).items()])
                niter += 1

        # REMOVE THIS WHEN pymc3 v3.8 is a bit older
        try:
            self.summary = pm.summary(self.traces)
        except:
            self.summary = pm.stats.summary(self.traces)

        self.par_names = self.summary.index

        samps = np.array([self.traces[x] for x in self.traces.varnames if not x.endswith('_log__')])
        self.samples = np.array([]).reshape((samps.shape[1], 0))
        for i in range(samps.shape[0]):
            self.samples = np.concatenate((self.samples, samps[i, :, :]), axis =1)


from __future__ import print_function

import string
import sys
from collections import deque

import numpy as np
from sklearn.base import BaseEstimator
from sklearn.utils import check_array, check_random_state
from sklearn.utils.validation import check_is_fitted

from . import _hmmc
from .utils import normalize, logsumexp, iter_from_X_lengths

DECODER_ALGORITHMS = frozenset(("viterbi", "map"))

class ConvergenceMonitor(object):
    """Monitors and reports convergence to :data:`sys.stderr`.

    Parameters
    ----------
    thresh : double
        Convergence threshold. The algorithm has convereged eitehr if
        the maximum number of iterations is reached or the log probability
        improvement between the two consecutive iterations is less than
        threshold.

    n_iter : int
        Maximum number of iterations to perform.

    verbose : bool
        If ``True`` then per-iteration convergence reports are printed,
        otherwise the monitor is mute.

    Attributes
    ----------
    history : deque
        The log probability of the data for the last two training
        iterations. If the values are not strictly increasing, the
        model did not converge.

    iter : int
        Number of iterations performed while training the model.
    """
    fmt = "{iter:>10d} {logprob:>16.4f} {delta:>+16.4f}"

    def __init__(self, thresh, n_iter, verbose):
        self.thresh = thresh
        self.n_iter = n_iter
        self.verbose = verbose
        self.history = deque(maxlen=2)
        self.iter = 1

    def report(self, logprob):
        if self.history and self.verbose:
            delta = logprob - self.history[-1]
            message = self.fmt.format(
                iter=self.iter, logprob=logprob, delta=delta)
            print(message, file=sys.stderr)

        self.history.append(logprob)
        self.iter += 1

    @property
    def converged(self):
        return (self.iter == self.n_iter or
                (len(self.history) == 2 and
                 self.history[1] - self.history[0] < self.thresh))

class _BaseHMM(BaseEstimator):
    """Hidden Markov Model base class.

    Representation of a hidden Markov model probability distribution.
    This class allows for easy evaluation of, sampling from, and
    maximum-likelihood estimation of the parameters of a HMM.

    See the instance documentation for details specific to a
    particular object.

    Parameters
    ----------
    n_components : int
        Number of states in the model.

    startprob_prior : array, shape (n_components, )
        Initial state occupation prior distribution.

    transmat_prior : array, shape (n_components, n_components)
        Matrix of prior transition probabilities between states.

    algorithm : string, one of the ``DECODER_ALGORITHMS```
        Decoder algorithm.

    random_state: RandomState or an int seed (0 by default)
        A random number generator instance.

    n_iter : int, optional
        Maximum number of iterations to perform.

    thresh : float, optional
        Convergence threshold.

    verbose : bool, optional
        When ``True`` per-iteration convergence reports are printed
        to :data:`sys.stderr`. You can diagnose convergence via the
        :attr:`monitor_` attribute.

    params : string, optional
        Controls which parameters are updated in the training
        process.  Can contain any combination of 's' for startprob,
        't' for transmat, and other characters for subclass-specific
        emmission parameters. Defaults to all parameters.

    init_params : string, optional
        Controls which parameters are initialized prior to
        training.  Can contain any combination of 's' for
        startprob, 't' for transmat, and other characters for
        subclass-specific emmission parameters. Defaults to all
        parameters.

    Attributes
    ----------
    startprob_ : array, shape (n_components, )
        Initial state occupation distribution.

    transmat_ : array, shape (n_components, n_components)
        Matrix of transition probabilities between states.
    """

    # This class implements the public interface to all HMMs that
    # derive from it, including all of the machinery for the
    # forward-backward and Viterbi algorithms.  Subclasses need only
    # implement _generate_sample_from_state(), _compute_log_likelihood(),
    # _init(), _initialize_sufficient_statistics(),
    # _accumulate_sufficient_statistics(), and _do_mstep(), all of
    # which depend on the specific emission distribution.
    #
    # Subclasses will probably also want to implement properties for
    # the emission distribution parameters to expose them publicly.

    def __init__(self, n_components=1,
                 startprob_prior=1.0, transmat_prior=1.0,
                 algorithm="viterbi", random_state=None,
                 n_iter=10, thresh=1e-2, verbose=False,
                 params=string.ascii_letters,
                 init_params=string.ascii_letters):
        self.n_components = n_components
        self.monitor_ = ConvergenceMonitor(thresh, n_iter, verbose)
        self.params = params
        self.init_params = init_params
        self.startprob_prior = startprob_prior
        self.transmat_prior = transmat_prior
        self.algorithm = algorithm
        self.random_state = random_state
        self.n_iter = n_iter
        self.thresh = thresh

    def score_samples(self, X, lengths=None):
        """Compute the log probability under the model and compute posteriors.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        logprob : float
            Log likelihood of ``X``.

        posteriors : array, shape (n_samples, n_components)
            State-membership probabilities for each sample in ``X``.

        See Also
        --------
        score : Compute the log probability under the model.
        decode : Find most likely state sequence corresponding to ``X``.
        """
        check_is_fitted(self, "startprob_")
        self._check()

        X = check_array(X)
        n_samples = X.shape[0]
        logprob = 0
        posteriors = np.zeros((n_samples, self.n_components))
        for i, j in iter_from_X_lengths(X, lengths):
            framelogprob = self._compute_log_likelihood(X[i:j])
            logprobij, fwdlattice = self._do_forward_pass(framelogprob)
            logprob += logprobij

            bwdlattice = self._do_backward_pass(framelogprob)
            posteriors[i:j] = self._compute_posteriors(fwdlattice, bwdlattice)
        return logprob, posteriors

    def score(self, X, lengths=None):
        """Compute the log probability under the model.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        logprob : float
            Log likelihood of ``X``.

        See Also
        --------
        score_samples : Compute the log probability under the model and
            posteriors.
        decode : Find most likely state sequence corresponding to ``X``.
        """
        check_is_fitted(self, "startprob_")
        self._check()

        # XXX we can unroll forward pass for speed and memory efficiency.
        logprob = 0
        for i, j in iter_from_X_lengths(X, lengths):
            framelogprob = self._compute_log_likelihood(X[i:j])
            logprobij, _fwdlattice = self._do_forward_pass(framelogprob)
            logprob += logprobij
        return logprob

    def _decode_viterbi(self, X):
        framelogprob = self._compute_log_likelihood(X)
        return self._do_viterbi_pass(framelogprob)

    def _decode_map(self, X):
        _, posteriors = self.score_samples(X)
        logprob = np.max(posteriors, axis=1).sum()
        state_sequence = np.argmax(posteriors, axis=1)
        return logprob, state_sequence

    def decode(self, X, lengths=None, algorithm=None):
        """Find most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.
        algorithm : string, one of the ``DECODER_ALGORITHMS``
            decoder algorithm to be used

        Returns
        -------
        logprob : float
            Log probability of the produced state sequence.

        state_sequence : array, shape (n_samples, )
            Labels for each sample from ``X`` obtained via a given
            decoder ``algorithm``.

        See Also
        --------
        score_samples : Compute the log probability under the model and
            posteriors.

        score : Compute the log probability under the model.
        """
        check_is_fitted(self, "startprob_")
        self._check()

        algorithm = algorithm or self.algorithm
        if algorithm not in DECODER_ALGORITHMS:
            raise ValueError("Unknown decoder {0!r}".format(algorithm))

        decoder = {
            "viterbi": self._decode_viterbi,
            "map": self._decode_map
        }[algorithm]

        X = check_array(X)
        n_samples = X.shape[0]
        logprob = 0
        state_sequence = np.empty(n_samples, dtype=int)
        for i, j in iter_from_X_lengths(X, lengths):
            # XXX decoder works on a single sample at a time!
            logprobij, state_sequenceij = decoder(X[i:j])
            logprob += logprobij
            state_sequence[i:j] = state_sequenceij

        return logprob, state_sequence

    def predict(self, X, lengths=None):
        """Find most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        state_sequence : array, shape (n_samples, )
            Labels for each sample from ``X``.
        """
        _, state_sequence = self.decode(X, lengths)
        return state_sequence

    def predict_proba(self, X, lengths=None):
        """Compute the posterior probability for each state in the model.

        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        posterios : array, shape (n_samples, n_components)
            State-membership probabilities for each sample from ``X``.
        """
        _, posteriors = self.score_samples(X, lengths)
        return posteriors

    def sample(self, n_samples=1, random_state=None):
        """Generate random samples from the model.

        Parameters
        ----------
        n_samples : int
            Number of samples to generate.

        random_state: RandomState or an int seed (0 by default)
            A random number generator instance. If ``None``, the object's
            random_state is used.

        Returns
        -------
        X : array, shape (n_samples, n_features)
            Feature matrix.
        state_sequence : array, shape (n_samples, )
            State sequence produced by the model.
        """
        check_is_fitted(self, "startprob_")

        if random_state is None:
            random_state = self.random_state
        random_state = check_random_state(random_state)

        startprob_cdf = np.cumsum(self.startprob_)
        transmat_cdf = np.cumsum(self.transmat_, axis=1)

        currstate = (startprob_cdf > random_state.rand()).argmax()
        state_sequence = [currstate]
        X = [self._generate_sample_from_state(
            currstate, random_state=random_state)]

        for t in range(n_samples - 1):
            currstate = (transmat_cdf[currstate] > random_state.rand()) \
                .argmax()
            state_sequence.append(currstate)
            X.append(self._generate_sample_from_state(
                currstate, random_state=random_state))

        return np.atleast_2d(X), np.array(state_sequence, dtype=int)

    def fit(self, X, lengths=None):
        """Estimate model parameters.

        An initialization step is performed before entering the
        EM-algorithm. If you want to avoid this step for a subset of
        the parameters, pass proper ``init_params`` keyword argument
        to estimator's constructor.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
        lengths : array-like of integers, shape (n_sequences, )
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        self : object
            Returns self.
        """
        X = check_array(X)
        self._init(X, lengths=lengths, params=self.init_params)
        self._check()

        for iter in range(self.n_iter):
            stats = self._initialize_sufficient_statistics()
            curr_logprob = 0
            for i, j in iter_from_X_lengths(X, lengths):
                framelogprob = self._compute_log_likelihood(X[i:j])
                logprob, fwdlattice = self._do_forward_pass(framelogprob)
                curr_logprob += logprob
                bwdlattice = self._do_backward_pass(framelogprob)
                posteriors = self._compute_posteriors(fwdlattice, bwdlattice)
                self._accumulate_sufficient_statistics(
                    stats, X[i:j], framelogprob, posteriors, fwdlattice,
                    bwdlattice, self.params)

            self.monitor_.report(curr_logprob)
            if self.monitor_.converged:
                break

            self._do_mstep(stats, self.params)

        return self

    def _do_viterbi_pass(self, framelogprob):
        n_observations, n_components = framelogprob.shape
        state_sequence, logprob = _hmmc._viterbi(
            n_observations, n_components, np.log(self.startprob_),
            np.log(self.transmat_), framelogprob)
        return logprob, state_sequence

    def _do_forward_pass(self, framelogprob):
        n_observations, n_components = framelogprob.shape
        fwdlattice = np.zeros((n_observations, n_components))
        _hmmc._forward(n_observations, n_components, np.log(self.startprob_),
                       np.log(self.transmat_), framelogprob, fwdlattice)
        return logsumexp(fwdlattice[-1]), fwdlattice

    def _do_backward_pass(self, framelogprob):
        n_observations, n_components = framelogprob.shape
        bwdlattice = np.zeros((n_observations, n_components))
        _hmmc._backward(n_observations, n_components, np.log(self.startprob_),
                        np.log(self.transmat_), framelogprob, bwdlattice)
        return bwdlattice

    def _compute_posteriors(self, fwdlattice, bwdlattice):
        log_gamma = fwdlattice + bwdlattice
        # gamma is guaranteed to be correctly normalized by logprob at
        # all frames, unless we do approximate inference using pruning.
        # So, we will normalize each frame explicitly in case we
        # pruned too aggressively.
        log_gamma += np.finfo(float).eps
        log_gamma -= logsumexp(log_gamma, axis=1)[:, np.newaxis]
        out = np.exp(log_gamma)
        normalize(out, axis=1)
        return out

    def _compute_log_likelihood(self, X):
        pass

    def _generate_sample_from_state(self, state, random_state=None):
        pass

    def _init(self, X, lengths, params):
        init = 1. / self.n_components
        if 's' in params or not hasattr(self, "startprob_"):
            self.startprob_ = np.full(self.n_components, init)
        if 't' in params or not hasattr(self, "transmat_"):
            self.transmat_ = np.full((self.n_components, self.n_components),
                                     init)

    def _check(self):
        self.startprob_ = np.asarray(self.startprob_)
        if len(self.startprob_) != self.n_components:
            raise ValueError("startprob_ must have length n_components")
        if not np.allclose(self.startprob_.sum(), 1.0):
            raise ValueError("startprob_ must sum to 1.0 (got {0:.4f})"
                             .format(self.startprob_.sum()))

        self.transmat_ = np.asarray(self.transmat_)
        if self.transmat_.shape != (self.n_components, self.n_components):
            raise ValueError(
                "transmat_ must have shape (n_components, n_components)")
        if not np.allclose(self.transmat_.sum(axis=1), 1.0):
            raise ValueError("rows of transmat_ must sum to 1.0 (got {0})"
                             .format(self.transmat_.sum(axis=1)))

    # Methods used by self.fit()

    def _initialize_sufficient_statistics(self):
        stats = {'nobs': 0,
                 'start': np.zeros(self.n_components),
                 'trans': np.zeros((self.n_components, self.n_components))}
        return stats

    def _accumulate_sufficient_statistics(self, stats, seq, framelogprob,
                                          posteriors, fwdlattice, bwdlattice,
                                          params):
        stats['nobs'] += 1
        if 's' in params:
            stats['start'] += posteriors[0]
        if 't' in params:
            n_observations, n_components = framelogprob.shape
            # when the sample is of length 1, it contains no transitions
            # so there is no reason to update our trans. matrix estimate
            if n_observations <= 1:
                return

            lneta = np.zeros((n_observations - 1, n_components, n_components))
            _hmmc._compute_lneta(n_observations, n_components, fwdlattice,
                                 np.log(self.transmat_),
                                 bwdlattice, framelogprob, lneta)
            stats['trans'] += np.exp(logsumexp(lneta, axis=0))

    def _do_mstep(self, stats, params):
        # Based on Huang, Acero, Hon, "Spoken Language Processing",
        # p. 443 - 445
        if 's' in params:
            self.startprob_ = self.startprob_prior - 1.0 + stats['start']
            normalize(self.startprob_)
        if 't' in params:
            self.transmat_ = self.transmat_prior - 1.0 + stats['trans']
            normalize(self.transmat_, axis=1)