#!/usr/bin/env python
# coding: utf-8

# ## Trees from bipartitions files

# In[1]:


from Bio import Phylo
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
plt.rcParams["figure.figsize"] = [12,9]

ipyrad_reftree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ipyrad-reference-empirical"
stackstree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.stacks-empirical"
ddocenttree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ddocent-empirical"
ipyrad_minustree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ipyrad-denovo_minus_reference-empirical"
ipyrad_plustree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ipyrad-denovo_plus_reference-095-empirical"
treefiles = {"ipyrad-reference":ipyrad_reftree, "stacks":stackstree,\
             "ddocent":ddocenttree, "ipyrad-denovo-minus":ipyrad_minustree,\
             "ipyrad-denovo-plus":ipyrad_plustree}

## empirical
emp_pops = {}
emp_sample_names = []
popfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/stacks/popmap.txt"
with open(popfile) as infile:
    lines = [l.strip().split() for l in infile.readlines()]
    emp_sample_names = [x[0] for x in lines]
    pops = set([x[1] for x in lines])
    for pop in pops: emp_pops[pop] = []
    for line in lines:
        p = line[1]
        s = line[0]
        emp_pops[p].append(s)
emp_pop_colors = {k:v for (k,v) in zip(emp_pops, list(matplotlib.colors.cnames))}
#print(emp_pop_colors)
#print(emp_pops)

emp_colors_per_sample = {}
for samp in emp_sample_names:
    for pop in emp_pops:
        if samp in emp_pops[pop]:
            emp_colors_per_sample[samp] = emp_pop_colors[pop]
#print(emp_colors_per_sample)

legend_colors = [matplotlib.patches.Patch(color=v, label=k) for k,v in emp_pop_colors.items()]

## Root the tree on the WBS pop
outgroup = [{'name': taxon_name} for taxon_name in emp_pops["WBS"]]

f, axarr = plt.subplots(5, 1, figsize=(10, 30), dpi=1000)
#axarr = [a for b in axarr for a in b]

for name, ax in zip(treefiles.keys(), axarr):
    print(name, ax)
    tree = Phylo.read(treefiles[name], 'newick')
    tree.root_with_outgroup(*outgroup)

    tree.ladderize()
    ## This could be cool but doesn't work
    Phylo.draw(tree, axes=ax, label_colors=emp_colors_per_sample, do_show=False)
    ax.set_title(name)
    ax.legend(handles=legend_colors, loc="lower left")
#axarr[0].legend(handles=legend_colors)


# ## Trees from bestTree

# In[30]:


from Bio import Phylo
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
plt.rcParams["figure.figsize"] = [12,9]

ipyrad_reftree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bestTree.ipyrad-reference-empirical"
stackstree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bestTree.stacks-empirical"
ddocenttree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bestTree.ddocent-empirical"
ipyrad_minustree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bestTree.ipyrad-denovo_minus_reference-empirical"
ipyrad_plustree = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bestTree.ipyrad-denovo_plus_reference-095-empirical"
treefiles = {"ipyrad-reference":ipyrad_reftree, "stacks":stackstree,\
             "ddocent":ddocenttree, "ipyrad-denovo-minus":ipyrad_minustree,\
             "ipyrad-denovo-plus":ipyrad_plustree}

## empirical
emp_pops = {}
emp_sample_names = []
popfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/stacks/popmap.txt"
with open(popfile) as infile:
    lines = [l.strip().split() for l in infile.readlines()]
    emp_sample_names = [x[0] for x in lines]
    pops = set([x[1] for x in lines])
    for pop in pops: emp_pops[pop] = []
    for line in lines:
        p = line[1]
        s = line[0]
        emp_pops[p].append(s)
emp_pop_colors = {k:v for (k,v) in zip(emp_pops, list(matplotlib.colors.cnames))}
#print(emp_pop_colors)
#print(emp_pops)

emp_colors_per_sample = {}
for samp in emp_sample_names:
    for pop in emp_pops:
        if samp in emp_pops[pop]:
            emp_colors_per_sample[samp] = emp_pop_colors[pop]
#print(emp_colors_per_sample)

## Root the tree on the WBS pop
outgroup = [{'name': taxon_name} for taxon_name in emp_pops["WBS"]]

legend_colors = [matplotlib.patches.Patch(color=v, label=k) for k,v in emp_pop_colors.items()]

f, axarr = plt.subplots(5, 1, figsize=(10, 30), dpi=1000)
#axarr = [a for b in axarr for a in b]

for name, ax in zip(treefiles.keys(), axarr):
    print(name, ax)
    tree = Phylo.read(treefiles[name], 'newick')
    tree.root_with_outgroup(*outgroup)
    tree.ladderize()
    ## This could be cool but doesn't work
    Phylo.draw(tree, axes=ax, label_colors=emp_colors_per_sample, do_show=False)
    ax.set_title(name)
    ax.legend(handles=legend_colors, loc="lower left")


# # Everything below here is crap

# In[20]:


outgroup = [{'name': taxon_name} for taxon_name in emp_pops["WBS"]]
print(outgroup)


# In[20]:


from Bio import Phylo
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
plt.rcParams["figure.figsize"] = [12,9]

f, axarr = plt.subplots(1, 1, figsize=(20,8), dpi=1000)

## empirical
emp_pops = {}
emp_sample_names = []
popfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/stacks/popmap.txt"
with open(popfile) as infile:
    lines = [l.strip().split() for l in infile.readlines()]
    emp_sample_names = [x[0] for x in lines]
    pops = set([x[1] for x in lines])
    for pop in pops: emp_pops[pop] = []
    for line in lines:
        p = line[1]
        s = line[0]
        emp_pops[p].append(s)
emp_pop_colors = {k:v for (k,v) in zip(emp_pops, list(matplotlib.colors.cnames))}
#print(emp_pop_colors)
#print(emp_pops)

emp_colors_per_sample = {}
for samp in emp_sample_names:
    for pop in emp_pops:
        if samp in emp_pops[pop]:
            emp_colors_per_sample[samp] = emp_pop_colors[pop]
print(emp_colors_per_sample)

tfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ipyrad-denovo_minus_reference-empirical"
tree = Phylo.read(tfile, 'newick')
tree.ladderize()
## This could be cool but doesn't work
Phylo.draw(tree, axes=axarr, label_colors=emp_colors_per_sample)
axarr.legend(emp_pop_colors)


# In[21]:


from Bio import Phylo
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
plt.rcParams["figure.figsize"] = [12,9]

f, axarr = plt.subplots(1, 1, figsize=(20,8), dpi=1000)

## empirical
emp_pops = {}
emp_sample_names = []
popfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/stacks/popmap.txt"
with open(popfile) as infile:
    lines = [l.strip().split() for l in infile.readlines()]
    emp_sample_names = [x[0] for x in lines]
    pops = set([x[1] for x in lines])
    for pop in pops: emp_pops[pop] = []
    for line in lines:
        p = line[1]
        s = line[0]
        emp_pops[p].append(s)
emp_pop_colors = {k:v for (k,v) in zip(emp_pops, list(matplotlib.colors.cnames))}
#print(emp_pop_colors)
#print(emp_pops)

emp_colors_per_sample = {}
for samp in emp_sample_names:
    for pop in emp_pops:
        if samp in emp_pops[pop]:
            emp_colors_per_sample[samp] = emp_pop_colors[pop]
print(emp_colors_per_sample)

tfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ddocent-empirical"
tree = Phylo.read(tfile, 'newick')
tree.ladderize()
## This could be cool but doesn't work
Phylo.draw(tree, axes=axarr, label_colors=emp_colors_per_sample)
axarr.legend(emp_pop_colors)


# In[1]:


from Bio import Phylo
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
plt.rcParams["figure.figsize"] = [12,9]

f, axarr = plt.subplots(1, 1, figsize=(20,8), dpi=1000)

## empirical
emp_pops = {}
emp_sample_names = []
popfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/stacks/popmap.txt"
with open(popfile) as infile:
    lines = [l.strip().split() for l in infile.readlines()]
    emp_sample_names = [x[0] for x in lines]
    pops = set([x[1] for x in lines])
    for pop in pops: emp_pops[pop] = []
    for line in lines:
        p = line[1]
        s = line[0]
        emp_pops[p].append(s)
emp_pop_colors = {k:v for (k,v) in zip(emp_pops, list(matplotlib.colors.cnames))}
#print(emp_pop_colors)
#print(emp_pops)

emp_colors_per_sample = {}
for samp in emp_sample_names:
    for pop in emp_pops:
        if samp in emp_pops[pop]:
            emp_colors_per_sample[samp] = emp_pop_colors[pop]
print(emp_colors_per_sample)

tfile = "/home/iovercast/manuscript-analysis/Phocoena_empirical/raxml_outdir/RAxML_bipartitions.ipyrad-reference-empirical"
tree = Phylo.read(tfile, 'newick')
tree.ladderize()
## This could be cool but doesn't work
Phylo.draw(tree, axes=axarr, label_colors=emp_colors_per_sample)
axarr.legend(emp_pop_colors)


# In[ ]: