def calc_joint_prob(arg, seqs, ntimes=20, mu=2.5e-8, rho=1.5e-8, popsizes=1e4,
                    times=None, verbose=False, delete_arg=True):
    """
    Calculate arg_joint_prob
    """
    if times is None:
        times = argweaver.get_time_points(
            ntimes=ntimes, maxtime=80000, delta=.01)
    if isinstance(popsizes, float) or isinstance(popsizes, int):
        popsizes = [popsizes] * len(times)

    if verbose:
        util.tic("calc likelihood")

    trees, names = arg2ctrees(arg, times)
    seqs, nseqs, seqlen = seqs2cseqs(seqs, names)

    p = argweaver_joint_prob(
        trees, times, len(times), popsizes, mu, rho, seqs, nseqs, seqlen)
    if delete_arg:
        delete_local_trees(trees)

    if verbose:
        util.toc()

    return p

def setTreeDistances(conf, tree, distmat, genes):
    if isDebug(DEBUG_MED):
        util.tic("fit branch lengths")
    
    if pyspidir and "parsimony" in conf:
        # estimate branch lengths with parsimony
        parsimony_C(conf["aln"], tree)
        tree.data["error"] = sum(node.dist 
                                 for node in tree.nodes.itervalues())
    
    elif pyspidir and "mlhkydist" in conf:
        # estimate branch lengths with ML
        logl = mlhkydist_C(conf["aln"], tree, conf["bgfreq"], conf["tsvratio"], 
                           3*len(tree.nodes))
        tree.data["distlogl"] = logl
        tree.data["error"] = 0.0
    else:
        # perform LSE
        lse = phylo.least_square_error(tree, distmat, genes)

        # catch unusual case that may occur in greedy search
        if sum(x.dist for x in tree.nodes.values()) == 0:
            for node in tree.nodes.values():
                node.dist = .01

        tree.data["error"] = math.sqrt(scipy.dot(lse.resids, lse.resids)) / \
                                       sum(x.dist for x in tree.nodes.values())

        setBranchError(conf, tree, lse.resids, lse.paths, lse.edges, lse.topmat)
        
    if isDebug(DEBUG_MED):
        util.toc()

    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 trainTree(conf, stree, gene2species):
    args = conf["REST"]
    treefiles = []
    
    for arg in args:
        treefiles.extend(util.shellparser(arg))

    util.tic("reading trees")
    trees = []
    prog = progress.ProgressBar(len(treefiles))
    for treefile in treefiles:
        prog.update()
        trees.append(treelib.read_tree(treefile))
        
        # even out top two branches
        totlen = trees[-1].root.children[0].dist + \
                 trees[-1].root.children[1].dist
        trees[-1].root.children[0].dist = totlen / 2.0
        trees[-1].root.children[1].dist = totlen / 2.0
        
    util.toc()
    
    params = Spidir.learnModel(trees, stree, gene2species, conf["trainstats"],
                               filenames=treefiles)
    
    Spidir.writeParams(conf["param"], params)

def dnadist(seqs, output=None, verbose=True, force = False, args=None):
    if args == None:
        args = "y"
    
    validate_seqs(seqs)
    cwd = create_temp_dir()
    util.tic("dnadist on %d of length %d" % (len(seqs), len(seqs.values()[0])))

    # create input
    labels = write_phylip_align(file("infile", "w"), seqs)
    
    # run phylip
    exec_phylip("dnadist", args, verbose)
    
    util.toc()    
    
    # parse output
    if output != None:
        os.rename("outfile", "../" + output)
        cleanup_temp_dir(cwd)
        return labels
    else:
        name, mat = read_dist_matrix("outfile")    
        cleanup_temp_dir(cwd)
        return labels, mat

 def processFunc():
     # remove old query tempfile if one exists
     if closure["oldtmp"] != None:
         os.remove(closure["oldtmp"])
         elapse = util.toc()
         closure["time"] += elapse
         
         util.log("blasted %d of %d sequences (%.1f%%), elapse %.0f m, left %.0f m" % (
             closure["index"], len(seqs.keys()), 
             100 * float(closure["index"]) / len(seqs.keys()),
             closure["time"] / 60.0, 
             elapse / split * (len(seqs.keys()) - closure["index"]) / 60.0))
         
     util.tic()
     
     # find new subset of query sequences
     i = closure["index"]
     names = seqs.keys()[i:i+split]
     
     # if no more sequences then quit
     if len(names) == 0:
         return False
     
     # start blast
     tmpfile = util.tempfile(".", "blastp", ".fasta")
     seqs.write(tmpfile, names = names)
     pipe = os.popen("blastall -p %s -d %s -i %s -m 8 -e .1 %s" % \
         (prog, databaseFile, tmpfile, options))
     
     # update variables
     closure["oldtmp"] = tmpfile
     closure["index"] = i + split
     
     return pipe

def boot_proml(seqs, iters = 100, seed = 1, jumble=5, output=None, 
               verbose=True, force = False):
    validate_seqs(seqs)
    cwd = create_temp_dir()
    util.tic("bootProml on %d of length %d" % (len(seqs), len(seqs.values()[0])))

    # create input
    labels = write_phylip_align(file("infile", "w"), seqs)
    
    exec_phylip("seqboot", "y\n%d" % seed, verbose)
    
    os.rename("outfile", "infile")
    exec_phylip("proml", "m\nD\n%d\n%d\n%d\ny" % (iters, seed, jumble), verbose)
    
    util.toc()        
    
    # read tree samples
    if output != None:
        os.rename("outtree", "../" + output)
        cleanup_temp_dir(cwd)
        return labels
    else:
        trees = []
        infile = file("outtree")
        for i in xrange(iters):
            tree = treelib.Tree()
            tree.read_newick(infile)
            rename_tree_with_names(tree, labels)
            trees.append(tree)
        infile.close()
        cleanup_temp_dir(cwd)
        return trees

    def addGoTerms(self, gofile):

        if not tableExists(self.cur, "GoTerms"):
            self.makeGoTermsTable()

        util.tic("add go terms")
        goterms = tablelib.read_table(gofile)
        goterms_lookup = goterms.groupby("orf")
        goterms_bygoid = goterms.groupby("goid")

        for goterm in goterms_bygoid:
            term = goterms_bygoid[goterm][0]

            if '"' in term["term"]:
                print term

            self.cur.execute("""INSERT INTO GoTerms VALUES ("%s", "%s")""" %
                             (term["goid"], term["term"]))

        for gene, terms in goterms_lookup.iteritems():
            for term in terms:
                self.cur.execute(
                    """INSERT INTO GeneGoTerms VALUES ("%s", "%s");""" %
                    (gene, term["goid"]))
        util.toc()

    def addEvents(self, eventsfile):

        if not tableExists(self.cur, "Events"):
            self.makeEventsTable()

        util.tic("add events")
        events_tab = tablelib.read_table(eventsfile)
        events_lookup = events_tab.lookup("partid")

        self.cur.execute("SELECT famid FROM Families;")
        famids = [x[0] for x in self.cur]

        for famid in famids:
            if famid not in events_lookup:
                continue
            row = events_lookup[famid]

            for sp in self.stree.nodes:
                sp = str(sp)

                self.cur.execute(
                    """INSERT INTO Events VALUES
                    ("%s", "%s", %d, %d, %d, %d);""" %
                    (famid, sp,
                     row[sp+"-genes"],
                     row[sp+"-dup"],
                     row[sp+"-loss"],
                     row[sp+"-appear"]))
        util.toc()

    def addFamilies(self, eventsfile, discard=[]):

        if not tableExists(self.cur, "Families"):
            self.makeFamiliesTable()

        util.tic("add families")
        events_tab = tablelib.read_table(eventsfile)
        events_lookup = events_tab.lookup("partid")
        familyGeneNames = self.makeFamilyGeneNames()
        discard = set(discard)

        for row in events_tab:
            famid = row["partid"]
            if famid in discard:
                util.logger("discarding '%s'" % famid)
                continue

            tree = treelib.read_tree(self.getTreeFile(famid))
            treelen = sum(x.dist for x in tree)
            seqs = fasta.read_fasta(self.getFastaFile(famid))
            seqlen = stats.median(map(len, seqs.values()))

            self.cur.execute(
                """INSERT INTO Families VALUES
                ("%s", "%s", %f, %f, %f, %d, %d, %d,
                "%s");""" %
                (row["partid"],
                 familyGeneNames.get(row["partid"], ("", ""))[0],
                 row["famrate"], treelen, seqlen * 3,
                 row["dup"], row["loss"], row["genes"],
                 familyGeneNames.get(row["partid"], ("", ""))[1]))
        util.toc()

def sample_thread(arg, seqs, rho=1.5e-8, mu=2.5e-8, popsize=1e4,
                  times=None, ntimes=20, maxtime=200000, verbose=False):

    if times is None:
        times = argweaver.get_time_points(
            ntimes=ntimes, maxtime=maxtime, delta=.01)
    popsizes = [popsize] * len(times)

    if verbose:
        util.tic("sample thread")

    trees, names = arg2ctrees(arg, times)

    seqs2 = [seqs[name] for name in names]

    new_name = [x for x in seqs.keys() if x not in names][0]
    names.append(new_name)
    seqs2.append(seqs[new_name])
    seqlen = len(seqs2[0])

    trees = argweaver_sample_thread(
        trees, times, len(times),
        popsizes, rho, mu,
        (C.c_char_p * len(seqs2))(*seqs2), len(seqs2), seqlen, None)
    arg = ctrees2arg(trees, names, times, verbose=verbose)

    if verbose:
        util.toc()

    return arg

    def walk(node):
        for child in node.children:
            walk(child)

        if not node.is_leaf():
            blastfiles = []
            leaves1 = node.children[0].leaf_names()
            leaves2 = node.children[1].leaf_names()

            # determine sibling blast files
            for leaf1 in leaves1:
                for leaf2 in leaves2:
                    if leaf1 in blastFileLookup and \
                       leaf2 in blastFileLookup[leaf1]:
                        blastfiles.append(blastFileLookup[leaf1][leaf2])

            # determine outgroup blast files (all other files, potentially)
            # go up one level, blastfiles for leaves, and subtract
            # sibling files
            outblastfiles = []
            if node.parent:
                inleaves = leaves1 + leaves2
                outleaves = set(node.parent.leaf_names()) - set(inleaves)

                for leaf1 in inleaves:
                    for leaf2 in outleaves:
                        if leaf1 in blastFileLookup and \
                           leaf2 in blastFileLookup[leaf1]:
                            outblastfiles.append(blastFileLookup[leaf1][leaf2])

            util.tic("merging")
            util.logger("leaves1: ", leaves1)
            util.logger("leaves2: ", leaves2)

            if "merge" in conf and \
               conf["merge"] == "avg":
                node.parts = mergeAvg(conf,
                                      genes,
                                      node.children[0].parts,
                                      node.children[1].parts,
                                      blastfiles,
                                      outblastfiles)
            else:
                node.parts = mergeBuh(conf,
                                      genes,
                                      node.children[0].parts,
                                      node.children[1].parts,
                                      blastfiles)

            if "output" in conf and len(node.parts) > 0:
                util.write_delim(conf["output"] +
                                 str(node.name) +
                                 ".part", node.parts)

            util.logger("number of parts: ", len(node.parts))
            if len(node.parts) > 0:
                util.logger("largest part:", max(map(len, node.parts)))

            util.toc()

def draw_raxml_tree(tr, adef):
    util.tic("Tree to string...")
    treestr = raxml.tree_to_string(tr, adef)
    util.toc()

    util.tic("Drawing tree...")
    T = treelib.parse_newick(treestr)
    T2 = treelib.unroot(T)
    treelib.draw_tree(T2, out=sys.stdout, minlen=5, maxlen=5)
    util.toc()

def pamp(seqs, tree, seqtype="dna", saveOutput="", verbose=False, safe=True):
    
    if safe and seqtype == "dna":
        seqs = alignlib.mapalign(seqs, valfunc=removeStopCodons)
    
    phylip.validate_seqs(seqs)
    cwd = phylip.create_temp_dir()

    util.tic("pamp on %d of length %d" % (len(seqs), len(seqs.values()[0])))

    # create input
    nex = nexus.Nexus("align", "w")
    nex.write_matrix(seqs.keys(), seqs.values(), seqtype, seqs.alignlen())
    nex.close()
    treefile = open("tree", "w")
    treefile.write("%d 1\n" % len(tree.leaves()))
    tree.write(treefile, writeData=lambda x: "")
    treefile.close()
    
    # create control file
    out = file("pamp.ctl", "w")
    print >>out, "seqfile = align"
    print >>out, "treefile = tree"
    print >>out, "outfile = out"    

    if seqtype == "dna":
        print >>out, "seqtype = 0"
    elif seqtype == "pep":
        print >>out, "seqtype = 2"
    else:
        raise Exception("unknown seqtype '%s'" % seqtype)
    print >>out, "ncatG = 8"
    print >>out, "nhomo = 0"
    out.close()
    
    # run pamp
    if verbose:
        os.system("pamp paml.ctl")
    else:
        os.system("pamp paml.ctl > /dev/null")

    res = PamlResults("out")
    aln = res.getPampReconstruction()
    aln.write("recon.mfa")
    tree2 = res.getBranchNames()
    renameTreeAlign(tree2, aln)
    
    if saveOutput != "":
        phylip.save_temp_dir(cwd, saveOutput)
    else:
        phylip.cleanup_temp_dir(cwd)
    
    util.toc()

    return tree2, aln

def resample_arg_region(arg, seqs, region_start, region_end,
                        ntimes=20, rho=1.5e-8, mu=2.5e-8,
                        popsizes=1e4, times=None, carg=False,
                        refine=1, verbose=False):
    """
    Sample ARG for sequences
    """
    if times is None:
        times = argweaver.get_time_points(
            ntimes=ntimes, maxtime=80000, delta=.01)
    if isinstance(popsizes, float) or isinstance(popsizes, int):
        popsizes = [popsizes] * len(times)

    if verbose:
        util.tic("resample arg")

    # convert arg to c++
    if verbose:
        util.tic("convert arg")
    trees, names = arg2ctrees(arg, times)
    if verbose:
        util.toc()

    # get sequences in same order
    # and add all other sequences not in arg yet
    leaves = set(names)
    for name, seq in seqs.items():
        if name not in leaves:
            names.append(name)
    seqs2, nseqs, seqlen = seqs2cseqs(seqs, names)

    # resample arg
    seqlen = len(seqs[names[0]])

    trees = argweaver_resample_arg_region(
        trees, times, len(times),
        popsizes, rho, mu, seqs2, nseqs, seqlen,
        region_start, region_end, refine)

    #trees = argweaver_resample_arg_region(
    #    trees, times, len(times),
    #    popsizes, rho, mu, seqs2, nseqs, seqlen,
    #    region_start, region_end)

    # convert arg back to python
    if carg:
        arg = (trees, names)
    else:
        arg = ctrees2arg(trees, names, times, verbose=verbose)

    if verbose:
        util.toc()

    return arg

    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 sample_all_arg(seqs, ntimes=20, rho=1.5e-8, mu=2.5e-8, popsizes=1e4,
                   refine=1, times=None, verbose=False, carg=False,
                   prob_path_switch=.1):
    """
    Sample ARG for sequences
    """
    if times is None:
        times = argweaver.get_time_points(
            ntimes=ntimes, maxtime=80000, delta=.01)
    if isinstance(popsizes, float) or isinstance(popsizes, int):
        popsizes = [popsizes] * len(times)

    if verbose:
        util.tic("resample arg")

    # convert arg to c++
    if verbose:
        util.tic("convert arg")

    arg = argweaver.make_trunk_arg(
        0, len(seqs.values()[0]), name=seqs.keys()[0])
    trees, names = arg2ctrees(arg, times)
    if verbose:
        util.toc()

    # get sequences in same order
    # and add all other sequences not in arg yet
    seqs2 = [seqs[name] for name in names]
    leaves = set(names)
    for name, seq in seqs.items():
        if name not in leaves:
            names.append(name)
            seqs2.append(seq)

    # resample arg
    seqlen = len(seqs[names[0]])
    trees = argweaver_resample_all_arg(
        trees, times, len(times),
        popsizes, rho, mu,
        (C.c_char_p * len(seqs2))(*seqs2), len(seqs2),
        seqlen, refine, prob_path_switch)

    if carg:
        arg = (trees, names)
    else:
        # convert arg back to python
        arg = ctrees2arg(trees, names, times, verbose=verbose)

    if verbose:
        util.toc()

    return arg

def ctrees2arg(trees, names, times, verbose=False, delete_arg=True):
    """
    Convert a C data structure for the ARG into a python ARG
    """

    if verbose:
        util.tic("convert arg")

    # get local trees info
    nnodes = get_local_trees_nnodes(trees)
    ntrees = get_local_trees_ntrees(trees)

    # allocate data structures for treeset
    ptrees = []
    ages = []
    sprs = []
    blocklens = [0] * ntrees
    for i in range(ntrees):
        ptrees.append([0] * nnodes)
        ages.append([0] * nnodes)
        sprs.append([0, 0, 0, 0])

    # populate data structures
    get_local_trees_ptrees(trees, ptrees, ages, sprs, blocklens)

    # fully convert to python
    for i in range(ntrees):
        ptrees[i] = ptrees[i][:nnodes]
        ages[i] = ages[i][:nnodes]
        sprs[i] = sprs[i][:4]

    # convert treeset to arg data structure
    blocks = []
    start = 0
    for blocklen in blocklens:
        end = start + blocklen
        blocks.append((start, end))
        start = end

    assert len(names) == ((nnodes + 1) / 2)

    arg = treeset2arg(ptrees, ages, sprs, blocks, names, times)

    if delete_arg:
        delete_local_trees(trees)

    if verbose:
        util.toc()

    return arg

    def addGenes(self, species, gff_files, region_filter=lambda x: x):
        """populate genes table"""

        # clear Genes Table
        if not tableExists(self.cur, "Genes"):
            self.makeGenesTable()

        dups = set()

        util.tic("add genes")
        for sp, gff_file in zip(species, gff_files):
            for region in gff.read_gff(gff_file, regionFilter=region_filter):
                gene = region.data["ID"]
                #gene = row["name"]

                if gene in self.fams.genelookup:
                    famid = self.fams.getFamid(gene)
                    if len(self.fams.getGenes(famid)) < 2:
                        famid = "NONE"
                else:
                    famid = "NONE"

                if gene in dups:
                    continue
                dups.add(gene)

                assert region.start <= region.end

                if gene in self.gene2name:
                    common = self.gene2name[gene]["name"]
                    desc = self.gene2name[gene]["description"]
                else:
                    common = ""
                    desc = ""

                cmd = ("""INSERT INTO Genes VALUES
                         ("%s", "%s", "%s", "%s", %d, %d, %d, "%s", "%s");""" %
                       (gene,
                        common,
                        self.gene2species(gene),
                        region.seqname,
                        region.start,
                        region.end,
                        region.strand,
                        desc.replace('"', ''),
                        famid))

                self.cur.execute(cmd)
        util.toc()

def resample_mcmc_arg(arg, seqs, ntimes=20,
                      rho=1.5e-8, mu=2.5e-8, popsizes=1e4,
                      refine=1, times=None, verbose=False, carg=False,
                      window=200000, niters2=5):
    """
    Sample ARG for sequences
    """
    if times is None:
        times = argweaver.get_time_points(
            ntimes=ntimes, maxtime=80000, delta=.01)
    if isinstance(popsizes, float) or isinstance(popsizes, int):
        popsizes = [popsizes] * len(times)

    if verbose:
        util.tic("resample arg")

    # convert arg to c++
    if verbose:
        util.tic("convert arg")
    trees, names = arg2ctrees(arg, times)
    if verbose:
        util.toc()

    # get sequences in same order
    # and add all other sequences not in arg yet
    leaves = set(names)
    names = list(names)
    for name in seqs:
        if name not in leaves:
            names.append(name)
    seqs2, nseqs, seqlen = seqs2cseqs(seqs, names)

    # resample arg
    trees = argweaver_resample_mcmc_arg(
        trees, times, len(times),
        popsizes, rho, mu,
        seqs2, nseqs, seqlen, refine, niters2, window)

    if carg:
        arg = (trees, names)
    else:
        # convert arg back to python
        arg = ctrees2arg(trees, names, times, verbose=verbose)

    if verbose:
        util.toc()

    return arg