
import numpy as np

import tempfile

from graphviz import Digraph
import matplotlib.pyplot as plt
from skimage.io import imsave
from skimage.transform import resize

from scatteringtree import scatter_single_node

"""
Display the nodes of a scattering-tree graphically.
"""
def visualize_tree(scattering_tree, signal):
    st = scattering_tree
    return _visualize_tree(signal, st._prop_filters, st._out_filters, st._nonlin, st._pools)

def _visualize_tree(signal, propagation_filters, output_filters, nonlinearities, poolings):
    """
    @param signal_shape  Dimension of signal.
    """
    
    # create temporary directory
    tempdir = tempfile.mkdtemp()
    
    """
    Create a path to a temporary file in a temporary directory.
    """
    def create_path(filename, suffix=None):
        if suffix is not None:
            return '%s/%s.%s'%(tempdir, filename, suffix)
        return '%s/%s'%(tempdir, filename)
    
    def save_image(img, filename):
        img = img.astype(np.float)
        img = img - np.min(img)
        img /= np.max(img)
        img = resize(img, (64, 64))
        
        imsave(filename, img)
    
    graph = Digraph(comment='scattering tree', engine='dot')
    
    output = []
    nodes = [signal]
   
        
    meta = [None]*len(nodes)
    
    parent_id = ['root']
    
    level = 0
    for prop_filters, out_filters, nl, pooling in zip(propagation_filters, output_filters, nonlinearities, poolings):
        next_nodes = []
        next_meta = []
        next_parent_id = []
        
        assert(len(parent_id) == len(nodes))
        assert(len(meta) == len(nodes))
        
        # scatter each node on that level
        for n, m, parent, i in zip(nodes, meta, parent_id, range(0, len(nodes))):
            
            # create node
            node_id = 'node_%d_%d'%(level, i)
            imgfile = create_path(node_id, 'png')
            save_image(n, imgfile)
            graph.node(node_id, label='', shape='square', image=imgfile)
            
            # draw edge
            graph.edge(parent, node_id)
            
            sc = scatter_single_node(n, prop_filters, out_filters, nl, pooling, meta=m)
            out = sc['out']
            
            def gen_node_label(node_id, descr):
                html = """<<TABLE BORDER="0">
                                <TR><TD><IMG SRC="%s"/></TD></TR>
                                %s
                            </TABLE>>"""
                
                descr_html = ''
                for k,v in descr.items():
                    descr_html += """
                    <TR><TD>%s</TD><TD>%s</TD></TR>
                    """%(k,v)
                
                return html%(create_path(out_node, 'png'), descr_html)
            
            for o, j in zip(out, range(len(out))):
                
                out_node = 'node_%d_%d_out_%d'%(level, i, j)
                imgfile = create_path(out_node, 'png')
                save_image(o, imgfile)
                label = gen_node_label(out_node, {'output': '', 'shape': o.shape})
                graph.node(out_node, label=label, color='green', style='dashed', shape='box')
                graph.edge(node_id, out_node, constraint='true')
            
            prop = sc['prop']
            output.extend(out)
            next_nodes.extend(prop)
            next_meta.extend(sc['meta'])
            
            next_parent_id.extend([node_id]*len(prop))
            
        nodes = next_nodes
        meta = next_meta
        parent_id = next_parent_id
        level += 1
    print(graph.source)
    graph.render(create_path('scatteringtree'), view=True)
    return 0
