# This file is a part of Framsticks SDK.  http://www.framsticks.com/
# Copyright (C) 2023  Maciej Komosinski and Szymon Ulatowski.
# See LICENSE.txt for details.

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import sys
import csv

# Visualizing the influence of Framsticks f1/f4 modifier sequences. See also geneprops_test.cpp
# usage:
# ./geneprops_test f 3 -e | python3 frams/_demos/geneprops_graph.py
#
# See also https://colab.research.google.com/drive/1x9FTAUQKFhvxX_IyEcglWRWa4ylZeY-7?usp=sharing

#################################################
# loading

nodes_list = [] # ['modifiers', value]

for line in csv.reader(sys.stdin.readlines(),delimiter='\t'):
    if line[0]=='-':
        line[0]=''
    nodes_list.append([line[0],float(line[1])])

nodes_dict = {} # ['modifiers', value, [children,...]]
    
for n in sorted(nodes_list, key=lambda e: len(e[0])):
    tree_node = [n[0],n[1],[]]
    if n[0]!='':
        nodes_dict[n[0][:-1]][2].append(tree_node)
    nodes_dict[n[0]] = tree_node
    
#################################################
# plotting

fig, ax = plt.subplots(figsize=(4, 4))
ax.set(title="Tree of modifier values")

segs=[]
for n in nodes_dict.values():
    for child in n[2]:
        segs.append( ((len(n[0]),n[1]),(len(n[0])+1,child[1])) )

linecol = LineCollection(segs, linewidths=2, linestyle='solid', color='#ddddff')
ax.add_collection(linecol)

node_x = [len(n[0]) for n in nodes_dict.values()]
node_y = [n[1] for n in nodes_dict.values()]

ax.plot(node_x, node_y, "o", color="w", markerfacecolor="r")

for n in nodes_dict.values():
    ax.annotate(n[0], (len(n[0]),n[1]), textcoords="offset points", horizontalalignment="right", verticalalignment='center', xytext=(-4,0))

ax.xaxis.set_visible(False)
plt.tight_layout()
plt.show()