From 5ce4554256ee15b88204cda3b38551ce7e90c698 Mon Sep 17 00:00:00 2001
From: Katie Mummah <radioactivekate@gmail.com>
Date: Mon, 12 Apr 2021 12:38:00 -0600
Subject: [PATCH] commod dict can find nested commodities (aka streams)

---
 trailmap/commodity_dictionary.py | 74 ++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 4 deletions(-)

diff --git a/trailmap/commodity_dictionary.py b/trailmap/commodity_dictionary.py
index 6280315..27a6a88 100644
--- a/trailmap/commodity_dictionary.py
+++ b/trailmap/commodity_dictionary.py
@@ -27,7 +27,7 @@ def get_commod_names(metadata, uitype, agent):
     '''Return all archetypes and their aliases for a given uitype.
 
     inputs:
-    - metadata
+    - metadata: Cyclus metadata
     - uitype: a string. Example, "incommodity"
     - agent: a string. Example, ":cycamore:Enrichment"
 
@@ -40,14 +40,78 @@ def get_commod_names(metadata, uitype, agent):
 
     for var in agent_data:
         for param in agent_data[var]:
-            if param == "uitype" and uitype in agent_data[var][param]:
-                aliases.append(var)
+            if param == "uitype":
+                if uitype in agent_data[var][param]: #typical for archetypes
+                    aliases.append(var)
+                #if streams are present
+                path = search_var_recursive(agent_data[var][param], uitype)
+                if path is not None:
+                    aliases.append(find_alias(agent_data[var]["alias"], path))
+
+    #drop the "val" (signifing an interleave) and "streams_" aliases 
+    try:
+        while True:
+            aliases.remove("val")
+    except:
+        pass
+
+    try:
+        while True:
+            aliases.remove('streams_')
+    except:
+        pass
 
     return aliases
 
 
+def search_var_recursive(var, uitype):
+    '''Finds the path within streams_ or similar tree that matches desired
+    uitype. Searches recursively
+
+    inputs:
+        - var: a nested list from an archetype's metadata
+        - uitype: a string. Example, "incommodity"
+
+    outputs:
+        - path: the path within var that locates uitype
+    '''
+    for index,item in enumerate(var):
+        if item == uitype:
+            return [index]
+        if isinstance(item, list):
+            path = search_var_recursive(item, uitype)
+            if path:
+                return [index] + path
+
+
+def find_alias(var, path):
+    '''Given path to locate uitype, searches archetype aliases to locate
+    matching alias
+
+    inputs:
+        - var: a nested list of aliases
+        - path: the path to search for the desired alias
+
+    outputs:
+        - alias: a string
+    '''
+    if len(path) == 1:
+        return var[path[0]]
+    else:
+        return find_alias(var[path[0]], path[1:])
+
+
 def build_facility_dictionary(metadata, archetypes):
-    '''Identify commodities for each available archetype'''
+    '''Identify commodities for each available archetype
+    
+    inputs:
+        - metadata: Cyclus metadata
+        - archetypes: a list of archetypes to use
+
+    outputs:
+        - archetype_commods: a dictionary with the Cyclus archetypes available
+        and the names of their incommodities and outcommodities 
+    '''
     archetype_commods = {}
 
     for archetype in archetypes:
@@ -56,6 +120,8 @@ def build_facility_dictionary(metadata, archetypes):
         outcommods = get_commod_names(metadata["annotations"], "outcommodity",
                                       archetype)
 
+
+
         archetype_commods.update({archetype: (incommods,
                                               outcommods)})