From 4ad6429db68bdae9c7fec87b73ca6abb514ac236 Mon Sep 17 00:00:00 2001
From: Simon Crase <simon@greenweaves.nz>
Date: Thu, 12 Nov 2020 12:40:53 +1300
Subject: [PATCH] #46 Started work on qrtd

---
 README.md    |  2 +-
 phylogeny.py |  5 ++---
 qrtd.py      | 44 ++++++++++++++++++++++++++++++++------------
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 851e91c..0c82fd5 100644
--- a/README.md
+++ b/README.md
@@ -231,7 +231,7 @@ NB: functions generally use zero based indexing; [Rosalind](http://rosalind.info
 | ---- | --------------- | ------------------------------------------------|
 |alph|alph.py|[Alignment-Based Phylogeny](http://rosalind.info/problems/alph/) WIP|
 |chbp|chbp.py phylogeny.py|[Character-Based Phylogeny](http://rosalind.info/problems/chbp/) WIP|
-|cntq|chtq.py |[Counting Quartets](http://rosalind.info/problems/cntq/) WIP|
+|cntq|chtq.py |[Counting Quartets](http://rosalind.info/problems/cntq/)|
 |cset|cset.py|[Fixing an Inconsistent Character Set](http://rosalind.info/problems/cset/) WIP|
 |cstr|cstr.py phylogeny.py|[Creating a Character Table from Genetic Strings](http://rosalind.info/problems/cstr/)|
 |ctbl|ctbl.py phylogeny.py|[Creating a Character Table](http://rosalind.info/problems/ctbl/) |
diff --git a/phylogeny.py b/phylogeny.py
index fd7fb1a..0e47dcd 100644
--- a/phylogeny.py
+++ b/phylogeny.py
@@ -319,7 +319,6 @@ def ds(adj1,adj2):
     
     n       = len(species)
     seiceps = {species[i]:i for i in range(n)}
-    tree1   = replace_leaves(create_adj(parse(newick1,start=n)))
-    tree2   = replace_leaves(create_adj(parse(newick2,start=n)))
 
-    return ds(tree1,tree2)
\ No newline at end of file
+    return ds(replace_leaves(create_adj(parse(newick1,start=n))),
+              replace_leaves(create_adj(parse(newick2,start=n))))
\ No newline at end of file
diff --git a/qrtd.py b/qrtd.py
index 9aef309..d544e82 100644
--- a/qrtd.py
+++ b/qrtd.py
@@ -18,32 +18,52 @@
 import argparse
 import os
 import time
-from helpers import read_strings
+from   helpers import read_strings
+from   phylogeny import create_adj, parse
+
+def qrtd(species,newick1,newick2):
+    
+    def q(n):
+        return (n*(n-1)*(n-2)*(n-3))/(4*3*2*1)  
+    
+    def replace_leaves(adj):
+        return {parent:sorted([seiceps[child] if child in seiceps else child for child in children]) for parent,children in adj.items() }
+ 
+    def dq(T1,T2):
+        return q(count_leaves(T1)) +  q(count_leaves(T2)) -2 * shared_quartets(T1,T2)
+    
+    def count_leaves(adj):
+        return sum([1 for children in adj.values() for child in children if child<n])
+    
+    def shared_quartets(T1,T2):
+        return 0
+    
+    n       = len(species)
+    seiceps = {species[i]:i for i in range(n)}
+    tree1   = replace_leaves(create_adj(parse(newick1,start=n)))
+    tree2   = replace_leaves(create_adj(parse(newick2,start=n)))
+    return dq(tree1,tree2)
 
-  
 if __name__=='__main__':
     start = time.time()
     parser = argparse.ArgumentParser('....')
     parser.add_argument('--sample',   default=False, action='store_true', help='process sample dataset')
-    parser.add_argument('--extra',    default=False, action='store_true', help='process extra dataset')
     parser.add_argument('--rosalind', default=False, action='store_true', help='process Rosalind dataset')
     args = parser.parse_args()
     if args.sample:
-        pass
+        print(qrtd(
+            'A B C D E'.split(),
+            '(A,C,((B,D),E));',
+            '(C,(B,D),(A,E));'            
+        ))
         
-    
-    if args.extra:
-        Input,Expected  = read_strings('data/....txt',init=0)
-        ...
-  
     if args.rosalind:
         Input  = read_strings(f'data/rosalind_{os.path.basename(__file__).split(".")[0]}.txt')
  
-        Result = None
+        Result = qrtd(Input[0].split(), Input[1], Input[2])
         print (Result)
         with open(f'{os.path.basename(__file__).split(".")[0]}.txt','w') as f:
-            for line in Result:
-                f.write(f'{line}\n')
+            f.write(f'{Result}\n')
                 
     elapsed = time.time() - start
     minutes = int(elapsed/60)