def chi_square_fit(cdf, params, data, ndivs=20, minsamples=5, plot=False,
                   start=-util.INF, end=util.INF):

    from rasmus import gnuplot
    import scipy
    import scipy.stats

    # determine ndiv and binsize
    binsize = len(data) / ndivs
    if binsize < minsamples:
        ndivs = len(data) / minsamples
        binsize = len(data) / ndivs

    data = sorted(data)
    bins = [data[i:i+binsize] for i in xrange(0, len(data), binsize)]
    obs = scipy.array(map(len, bins))
    ind = util.find(lambda x: x[-1] >= start and x[0] <= end, bins)
    obs = util.mget(obs, ind)
    
    x = [bin[0] for bin in bins]
    expected = [len(data) * cdf(x[1], params)]
    expected.extend([len(data) *
                     (cdf(x[i+1], params) - cdf(x[i], params))
                     for i in range(1, len(x)-1)])
    expected.append(len(data) * (1.0 - cdf(x[-1], params)))
    expected = scipy.array(util.mget(expected, ind))
    
    chi2, pval = scipy.stats.chisquare(obs, expected)

    if plot:        
        p = gnuplot.plot(util.mget(x, ind), obs)
        p.plot(util.mget(x, ind), expected)
    
    return chi2, pval

 def _write_directive(self, line, out, delim):
     """Write a directive"""
     
     if line == DIR_VERSION:
         out.write("##version:%s\n" % self.version)
     
     elif line == DIR_TYPES:
         if len(self) > 0:
             entry = self[0]
         else:
             entry = [""] * len(self.headers)
         out.write("##types:" +
                   self._type_lookup.formatTableTypes(
                         util.mget(self.types, self.headers),
                         delim) + "\n")
     elif line == DIR_DEFAULTS:
         out.write("##defaults:" +
                   delim.join(map(str, 
                             util.mget(self.defaults, self.headers))) + "\n")
     
     elif line == DIR_HEADERS:
         out.write("##headers:%d\n" % self.nheaders)
     
     else:
         raise "unknown directive:", line

    def init_distmats(self):
        """Initialize distance matrices
        
           Initialization should by done after trees and alignments
        """
    
        if len(self.distmats) > 0:
            self.matrices = []
            
            # setup matrices            
            for i, distmat in enumerate(self.distmats):
                # convert distmatrix to summon Matrix
                if isinstance(distmat, matrix.Matrix):
                    mat = distmat
                else:
                    mat = matrix.Matrix()
                    mat.from2DList(distmat)            

                # set default colormap
                if mat.colormap == None:
                    mat.colormap = self.matrix_colormap
                
                # determine labels
                if self.dist_labels_from_align and self.align_order != None:
                    # determine row/col labels from alignment if it exists
                    mat.rowlabels = self.align_order
                    mat.collabels = self.align_order
                
                elif self.distlabels != None:
                    mat.rowlabels = self.distlabels[i]
                    mat.collabels = self.distlabels[i]
                    
                else:
                    raise Exception("no labels given for matrix")
                
                # reorder according to any given tree
                if self.order != None:
                    lookup = util.list2lookup(mat.rowlabels)
                    mat.rperm = util.mget(lookup, self.order)
                    mat.cperm = util.mget(lookup, self.order)
                
                mat.setup()

                self.matrices.append(mat)
        
            if self.seqs == None:
                seqs = self.current_align
            else:
                seqs = self.seqs
            
            # create matrix vis
            self.current_matrix = self.matrices[0]
            self.visdist = distmatrixvis.DistMatrixViewer(self.current_matrix, 
                                                          seqs=seqs, 
                                                          bgcolor=(1,1,1))
        else:
            self.visdist = None

 def walk(node):
     if node.isLeaf():
         return smap(node.name)
     else:
         child_hashes = map(walk, node.children)
         ind = util.sortrank(child_hashes)
         child_hashes = util.mget(child_hashes, ind)
         node.children = util.mget(node.children, ind)
         return hash_tree_compose(child_hashes)

def find_orthologs(gtree, stree, recon, counts=True):
    """Find all ortholog pairs within a gene tree"""

    events = label_events(gtree, recon)
    orths = []
    
    for node, event in events.items():
        if event == "spec":
            leavesmat = [x.leaves() for x in node.children]
            sp_counts = [util.hist_dict(util.mget(recon, row))
                         for row in leavesmat]
            
            for i in range(len(leavesmat)):
                for j in range(i+1, len(leavesmat)):
                    for gene1 in leavesmat[i]:
                        for gene2 in leavesmat[j]:
                            if gene1.name > gene2.name:
                                g1, g2 = gene2, gene1
                                a, b = j, i
                            else:
                                g1, g2 = gene1, gene2
                                a, b = i, j
                            
                            if not counts:
                                orths.append((g1.name, g2.name))
                            else:
                                orths.append((g1.name, g2.name,
                                              sp_counts[a][recon[g1]],
                                              sp_counts[b][recon[g2]]))
    
    return orths

 def lookup(self, *keys, **options):
     """Returns a lookup dict based on a column 'key'
        or multiple keys
        
        extra options:
        default=None
        uselast=False    # allow multiple rows, just use last
     """
     
     options.setdefault("default", None)
     options.setdefault("uselast", False)
     lookup = util.Dict(dim=len(keys), default=options["default"])
     uselast = options["uselast"]
     
     for row in self:
         keys2 = util.mget(row, keys)
         ptr = lookup
         for i in xrange(len(keys2) - 1):
             ptr = lookup[keys2[i]]
         if not uselast and keys2[-1] in ptr:
             raise Exception("duplicate key '%s'" % str(keys2[-1]))
         ptr[keys2[-1]] = row
     
     lookup.insert = False
     return lookup

    def test_ml_large(self):
        """Test ML code"""

        # params
        bgfreq = [.258,.267,.266,.209]
        kappa = 1.59

        # data
        tree = treelib.readTree("test/data/verts/19520/19520.ensembl.tree")
        align = fasta.readFasta("test/data/verts/19520/19520.nt.mfa")


        likes = []
        dists = []

        nodes = sorted(tree.nodes.values(), key=lambda x: x.dist)

        l = spidir.calc_seq_likelihood_hky(tree, align, bgfreq, kappa)
        print l
        self.assert_(l != -util.INF)


        l = spidir.find_ml_branch_lengths_hky(
            tree,
            util.mget(align, tree.leafNames()),
            bgfreq, kappa,
            parsinit=False,
            maxiter=1)
        print l
        self.assert_(l != -util.INF)

def fit_distrib(cdf, params_init, data, ndivs=20, minsamples=5,
                start=-util.INF, end=util.INF):

    import scipy
    import scipy.optimize
    import scipy.stats

    # determine ndiv and binsize
    binsize = len(data) / ndivs
    if binsize < minsamples:
        ndivs = len(data) / minsamples
        binsize = len(data) / ndivs

    data = sorted(data)
    bins = [data[i:i+binsize] for i in xrange(0, len(data), binsize)]
    obs = scipy.array(map(len, bins))
    ind = util.find(lambda x: x[-1] >= start and x[0] <= end, bins)
    obs = util.mget(obs, ind)
    
    def optfunc(params):
        x = [bin[0] for bin in bins]
        expected = [len(data) * cdf(x[1], params)]
        expected.extend([len(data) *
                         (cdf(x[i+1], params) - cdf(x[i], params))
                         for i in range(1, len(x)-1)])
        expected.append(len(data) * (1.0 - cdf(x[-1], params)))
        expected = scipy.array(util.mget(expected, ind))
        
        chi2, pval = scipy.stats.chisquare(obs, expected)
        return chi2

    params = scipy.optimize.fmin(optfunc, params_init, disp=False)
    chi2, pval = chi_square_fit(cdf, params, data, ndivs, minsamples)

    return list(params), pval

    def _test_ml_speed(self):
        
        # params
        bgfreq = [.258,.267,.266,.209]
        kappa = 1.59

        # data
        tree = treelib.readTree("test/data/flies.nt/0/0.tree")
        align = fasta.readFasta("test/data/flies.nt/0/0.align")


        likes = []
        dists = []

        nodes = sorted(tree.nodes.values(), key=lambda x: x.dist)

        util.tic("find ML")
        for i in xrange(10):
            l = spidir.find_ml_branch_lengths_hky(
                tree,
                util.mget(align, tree.leafNames()),
                bgfreq, kappa,
                maxiter=10)            
        util.toc()

        dists.append([n.dist for n in nodes])
        likes.append(l)

def read_length_matrix(filename, minlen=.0001, maxlen=1.0,
                       nooutliers=True):
    """Read a length matrix made by spidir-prep"""

    from rasmus import util

    dat = [line.rstrip().split("\t") for line in open(filename)]
    species = dat[0][2:]
    lens = util.map2(float, util.submatrix(dat, range(1, len(dat)),
                                           range(2, len(dat[0]))))
    gene_sizes = map(int, util.cget(dat[1:], 1))
    files = util.cget(dat[1:], 0)

    if nooutliers:
        treelens = map(sum, lens)
        m = mean(treelens)
        ind = util.find(lambda x: x<5*m, treelens)
        files, gene_sizes, lens, treelens = [util.mget(x, ind) for x in
                                             files, gene_sizes, lens, treelens]



    for row in lens:
        for i in xrange(len(row)):
            if row[i] < minlen:
                row[i] = minlen

    
    return species, lens, gene_sizes, files

 def on_reorder_leaves(self):
     leaves = self.current_tree.leaf_names()
     
     # reorder matrix
     for mat in self.matrices:
         lookup = util.list2lookup(mat.rowlabels)
         mat.rperm = util.mget(lookup, leaves)
         mat.cperm = util.mget(lookup, leaves)
         mat.setup()
     if self.visdist:
         self.visdist.redraw()
     
     
     # reorder alignment
     for aln in self.aligns:
         aln.names = leaves
     if self.visalign:
         self.visalign.show()

def parsimony_C(aln, tree):    
    ptree, nodes, nodelookup = makePtree(tree)
    leaves = [x.name for x in nodes if isinstance(x.name, str)]
    seqs = util.mget(aln, leaves)
    
    dists = pyspidir.parsimony(ptree, seqs)
    
    for i in xrange(len(dists)):
        nodes[i].dist = dists[i]

def mlhkydist_C(aln, tree, bgfreq, ratio, maxiter):
    ptree, nodes, nodelookup = makePtree(tree)
    leaves = [x.name for x in nodes if isinstance(x.name, str)]
    seqs = util.mget(aln, leaves)
    
    dists, logl = pyspidir.mlhkydist(ptree, seqs, bgfreq, ratio, maxiter)
    
    for i in xrange(len(dists)):
        nodes[i].dist = dists[i]
    
    return logl

    def draw_matches(self, sp, chrom, start, end, drawn=None):
        vis = []

        if drawn is None:
            drawn = set()
        
        # build list of matches in order of drawing
        
        for gene in iter_chrom(self.db.get_regions(sp, chrom), start, end):
            # need to sort matches by genome order so that mult-genome synteny
            # is drawn top-down

            # get orthologs
            genes2 = [x for x in self.orth_lookup.get(gene.data["ID"], [])
                      if x in self.region_layout]
            if len(genes2) == 0:
                continue
            
            rows = util.groupby(lambda x: self.region_layout[x].y, genes2)
            keys = util.sort(rows.keys(), reverse=True)
            rows = util.mget(rows, keys)

            l = self.region_layout
            
            for i in range(1, len(rows)):
                for botGene in rows[i]:
                    gene1 = self.db.get_region(botGene)
                    for topGene in rows[i-1]:

                        if (botGene, topGene) in drawn:
                            continue

                        drawn.add((botGene, topGene))
                        
                        gene2 = self.db.get_region(topGene)
                        y1 = l[topGene].y 
                        y2 = l[botGene].y + 1
                        x1 = l[topGene].x
                        x2 = l[topGene].x + gene2.length()
                        x3 = l[botGene].x + gene1.length()
                        x4 = l[botGene].x
                        
                        if self.fat_matches:
                            vis.append(quads(
                                    self.colors["matches"],
                                    x1, y1,
                                    x2, y1,
                                    x3, y2,
                                    x4, y2))

                        vis.append(lines(self.colors["matches"],
                                         x1, y1,
                                         x4, y2))
        return group(* vis)

 def optfunc(params):
     x = [bin[0] for bin in bins]
     expected = [len(data) * cdf(x[1], params)]
     expected.extend([len(data) *
                      (cdf(x[i+1], params) - cdf(x[i], params))
                      for i in range(1, len(x)-1)])
     expected.append(len(data) * (1.0 - cdf(x[-1], params)))
     expected = scipy.array(util.mget(expected, ind))
     
     chi2, pval = scipy.stats.chisquare(obs, expected)
     return chi2

    def _test_ml(self):
        """Test ML code"""

        # params
        bgfreq = [.258,.267,.266,.209]
        kappa = 1.59

        # data
        tree = treelib.readTree("test/data/flies.nt/0/0.tree")
        align = fasta.readFasta("test/data/flies.nt/0/0.align")


        likes = []
        dists = []

        nodes = sorted(tree.nodes.values(), key=lambda x: x.dist)

        util.tic("find ML")
        for i in range(40):
            l = spidir.find_ml_branch_lengths_hky(
                    tree,
                    util.mget(align, tree.leafNames()),
                    bgfreq, kappa,
                    parsinit=False,
                    maxiter=1)
            
            dists.append([n.dist for n in nodes])
            likes.append(l)
        util.toc()

        print likes

        prep_dir("test/output/ml/")

        # distances plot
        util.rplot_start("test/output/ml/ml_branches.pdf")
        util.rplot("plot", util.cget(dists, 0),
                   ylim=[0, max(dists[0])], t="l",
                   main="branch length convergence",
                   xlab="iterations",
                   ylab="branch lengths (sub/site)")
        for d in zip(* dists):
            util.rplot("lines", d)
        util.rplot_end(True)

        print util.cget(dists, 4)

        # likelihood plot
        util.rplot_start("test/output/ml/ml_likelihood.pdf")
        util.rplot("plot", likes, t="l",
                   xlab="iterations",
                   ylab="log likelihood",
                   main="likelihood convergence")
        util.rplot_end(True)

def chiSquareFit(xbins, ybins, func, nsamples, nparams, minsamples=5):
    sizes = [xbins[i+1] - xbins[i] for i in xrange(len(xbins)-1)]
    sizes.append(sizes[-1]) # NOTE: assumes bins are of equal size
    
    # only focus on bins that are large enough
    counts = [ybins[i] * sizes[i] * nsamples for i in xrange(len(xbins)-1)]
    
    expected = []
    for i in xrange(len(xbins)-1):
        expected.append((func(xbins[i]) + func(xbins[i+1]))/2.0 * 
                         sizes[i] * nsamples)
        
    # ensure we have enough expected samples in each bin
    ind = util.find(util.gefunc(minsamples), expected)
    counts = util.mget(counts, ind)
    expected = util.mget(expected, ind)
    
    if len(counts) == 0:
        return [0, 1], counts, expected
    else:
        return chiSquare([counts], [expected], nparams), counts, expected

def query_point_regions(point, regions, inc=True):

    ind = util.sortindex(regions, key=lambda r: r[1])
    rind = util.mget(range(len(regions)), ind)
    regions_by_end = util.mget(regions, ind)

    end = util.binsearch([r[0] for r in regions], x)[1]
    start = util.binsearch([r[1] for r in regions_by_end], x)[0]

    if start is None:
        start = 0
    if end is None:
        end = len(regions)

    if inc:
        for i in xrange(start, end):
            if regions[i][0] <= x <= regions[i][1]:
                yield regions[i]
    else:
        for i in xrange(start, end):
            if regions[i][0] < x < regions[i][1]:
                yield regions[i]

    def walk(node):
        if node in leaves:
            colors[node] = phylo.hash_tree(node, gene2species)
        else:
            # recurse
            for child in node.children:
                walk(child)

            childHashes = util.mget(colors, node.children)
            if len(childHashes) > 1 and util.equal(*childHashes):
                nmirrors[0] += 1

            childHashes.sort()
            colors[node] = phylo.hash_tree_compose(childHashes)

def make_pep_colors(prop2color=prop2color):
    pep_colors = util.Dict(default=color(.5, .5, .5))

    AA = 'ARNDCEQGHILKMFPSTWYVU*'
    pep_per_prop = util.hist_dict(util.mget(seqlib.AA_PROPERTY, AA))

    prop_counts = util.Dict(default=0)
    for char in AA:
        prop = seqlib.AA_PROPERTY[char]
        tint = prop_counts[prop] / float(pep_per_prop[prop])
        pep_colors[char] = prop2color(prop, tint * .5)
        prop_counts[prop] += 1
    
    return pep_colors