from matplotlib.pyplot import fill
import numpy as np

def parse_arg(tree1, tree2, aln1, aln2, MCC_file, fill_overhangs=True):
    from Bio import Phylo, AlignIO, Seq
    from Bio.Align import MultipleSeqAlignment
    from treetime.seq_utils import seq2array

    t1 = Phylo.read(tree1, 'newick')
    t2 = Phylo.read(tree2, 'newick')

    MCCs = []
    with open(MCC_file) as fh:
        for line in fh:
            if line.strip():
                MCCs.append(line.strip().split(','))

    a1 = {s.id:s for s in AlignIO.read(aln1, 'fasta')}
    a2 = {s.id:s for s in AlignIO.read(aln2, 'fasta')}
    all_leaves = set.union(set(a1.keys()), set(a2.keys()))
    for aln in [a1,a2]:
        for s,seq in aln.items():
            seqstr = "".join(seq2array(seq, fill_overhangs=fill_overhangs))
            seq.seq = Seq.Seq(seqstr)

    aln_combined = []
    for leaf in all_leaves:
        seq = a1[leaf] + a2[leaf]
        seq.id = leaf
        aln_combined.append(seq)

    l1 = len(a1[leaf])
    l2 = len(a2[leaf])
    combined_mask = np.ones(l1 + l2)
    mask1 = np.zeros(l1 + l2)
    mask2 = np.zeros(l1 + l2)
    mask1[:l1] = 1
    mask2[l1:] = 1

    return {"MCCs": MCCs, "trees":[t1,t2], "alignment":MultipleSeqAlignment(aln_combined),
            "masks":[mask1,mask2], "combined_mask":combined_mask}


def setup_arg(T, aln, total_mask, segment_mask, dates, MCCs, gtr='JC69',
              verbose=0, fill_overhangs=True):
    from treetime import TreeTime
    from collections import defaultdict

    tt = TreeTime(dates=dates, tree=T,
                  aln=aln, gtr=gtr, verbose=verbose,
                  fill_overhangs=fill_overhangs, compress=False)


    tt.reroot("least-squares", force_positive=True)

    leaf_to_MCC = {}
    for mi,mcc in enumerate(MCCs):
        for leaf in mcc:
            leaf_to_MCC[leaf] = mi

    for leaf in tt.tree.get_terminals():
        leaf.child_mccs = set([leaf_to_MCC[leaf.name]])
        leaf.mcc = leaf_to_MCC[leaf.name]

    for n in tt.tree.get_nonterminals(order='postorder'):
        n.child_mccs = set.union(*[c.child_mccs for c in n])

    mcc_intersection = set.intersection(*[c.child_mccs for c in tt.tree.root])
    if len(mcc_intersection):
        tt.tree.root.mcc = list(mcc_intersection)[0]
    else:
        tt.tree.root.mcc = None

    for n in tt.tree.get_nonterminals(order='preorder'):
        if n==tt.tree.root:
            continue
        else:
            if n.up.mcc in n.child_mccs:
                n.mcc = n.up.mcc
            elif len(n.child_mccs)==1:
                n.mcc = list(n.child_mccs)[0]
            else:
                n.mcc = None

    for n in tt.tree.find_clades():
        if (n.mcc is not None) and n.up and n.up.mcc==n.mcc:
            n.mask=total_mask
        else:
            n.mask = segment_mask

    return tt
@@ -1412,35 +1412,48 @@
 l_overhangs=True
+, reroot=True
 ):%0A
-
     from treetim
@@ -1670,16 +1670,35 @@
 alse)%0A%0A%0A
+    if reroot:%0A    
     tt.r
@@ -1741,16 +1741,16 @@
 e=True)%0A
-
 %0A    lea
@@ -2049,32 +2049,185 @@
 r='postorder'):%0A
+        common_mccs = set.intersection(*%5Bc.child_mccs for c in n%5D)%0A        if len(common_mccs):%0A            n.child_mccs = common_mccs%0A        else:%0A    
         n.child_

fix mcc computation and add flag not to root