#!/bin/env python3

import argparse
import random
import operator
import itertools
import subprocess


def tree_gen(depth, children, value_set, curr_depth=0):
    if random.randrange(0, depth) <= curr_depth:
        return random.choice(value_set), []
    else:
        return random.choice(value_set), [tree_gen(depth, children, value_set, curr_depth+1) for _ in range(random.randint(1,children))]


def find_paths(root, op=operator.add, pick_op=max):
    def calc_leafs(tree, aggregated, path='x0', op=op):
        aggregated += tree[0]
        if tree[1]:
            results = []
            for subtree, subpath in zip(tree[1], range(len(tree[1]))):
                results.extend(calc_leafs(subtree, aggregated, f'{path}x{subpath}', op))
            return results
        else:
            return [(aggregated, path)]

    leafs = calc_leafs(root, 0)
    picked = pick_op([x[0] for x in leafs])
    return tuple(filter(lambda x: x[0] == picked, leafs))


def tree2dot(tree, highlight_paths=(), label=''):
    ret = 'digraph isu8cv82 {\n'
    highlight = set()
    for x in highlight_paths:
        highlight.update(itertools.accumulate(tuple(map(''.join, zip(*[iter(x)] * 2)))))

    def conv_tree(parent, subtree, path='x0'):
        node = f'\t{path} [label="{parent}"' + (',color=red,style=bold' if path in highlight else '') + ']\n'
        if subtree:
            for ch_path in [f'{path}x{x}' for x in range(len(subtree))]:
                node += f'\t{path} -> {ch_path}' + ('[color=red,style=bold]' if ch_path in highlight else '') + '\n'

        subtree = ''.join(
            conv_tree(x[0][0], x[0][1], f'{path}x{x[1]}') for x in zip(subtree, range(len(subtree)))) if subtree else ''
        return node + subtree

    ret += conv_tree(tree[0], tree[1])
    ret += f'\tlabel="{label}"\n\tlabelloc="t"\n'
    ret += '}'
    return ret


def tree2asm(tree, start=0):
    # if there are no subtrees it's an easy case
    if not len(tree[1]):
        return f'{tree[0]},0', start + 2

    offset = start + len(tree[1]) + 2
    data = str(tree[0])
    subtree_strs = list()

    for subtree in tree[1]:
        data = f'{data},$+{offset*4}'
        subtree_str, offset = tree2asm(subtree, offset)
        subtree_strs.append(subtree_str)

    data = f'{data},0,{",".join(subtree_strs)}'
    return data, offset


parser = argparse.ArgumentParser(description='Script for generating inputs for the isu8-cv8-2 practice')
parser.add_argument('-s', '--seed', dest='seed', type=str, help='seed for random generator')
parser.add_argument('-d', '--depth', dest='depth', type=int, default=5, help='maximum tree depth, 5 by default')
parser.add_argument('-f', '--fanout', dest='fanout', type=int, default=5, help='maximum number of children for each node, 5 by default')
parser.add_argument('-g', '--generate-svg', dest='svg', type=str, help='generate svg image of the tree (graphwiz dot required)')

args = parser.parse_args()

if args.seed is not None:
    random.seed(args.seed)

tree = tree_gen(args.depth, args.fanout, range(-100, 100))
results, paths = zip(*find_paths(tree))

description = f'Maximum depth: {args.depth}, maximum fanout: {args.fanout}, seed: {"<not specified>" if args.seed is None else args.seed}, correct result: {", ".join(map(str, results))}'
print(description)
print(f'tree dd {tree2asm(tree)[0]}')

if args.svg:
    dot_graph = tree2dot(tree, paths, description)
    dot = subprocess.run(['dot', '-Tsvg'], input=dot_graph.encode(), stdout=subprocess.PIPE)
    with open(args.svg, 'wb') as fp:
        fp.write(dot.stdout)