package main

import (
	"fmt"
	"os"
	"runtime"
	"strconv"
)

type node struct {
	left  *node
	right *node
}

type result struct {
	depth int
	sum   int
}

var ch chan result
var results []int

func (n *node) Checksum() int {
	if n.left == nil {
		return 1
	}
	return n.left.Checksum() + n.right.Checksum() + 1
}

func NewNode(depth int) *node {
	n := new(node)

	if depth > 0 {
		n.left = NewNode(depth - 1)
		n.right = NewNode(depth - 1)
	}

	return n
}

func ManyTrees(iterations, depth int) {
	sum := 0

	var n *node
	for i := 0; i < iterations; i++ {
		n = NewNode(depth)
		sum += n.Checksum()
	}

	ch <- result{
		depth: depth,
		sum:   sum,
	}
}

func main() {
	minDepth := 4
	maxDepth := minDepth + 2

	if len(os.Args) == 2 {
		val, err := strconv.Atoi(os.Args[1])
		if err == nil {
			maxDepth = val
		}
	}

	ch = make(chan result, runtime.NumCPU())
	defer close(ch)
	results = make([]int, maxDepth+1)

	// Stretch tree
	stretch := NewNode(maxDepth + 1)
	fmt.Printf("stretch tree of depth %d\t check: %d\n", maxDepth+1, stretch.Checksum())

	// Long lived tree for later use
	long := NewNode(maxDepth)

	// Lots of trees in parallel
	c := 0
	for i := minDepth; i <= maxDepth; i += 2 {
		c++
		go ManyTrees(1<<uint(maxDepth-i+minDepth), i)
	}
	for i := 0; i < c; i++ {
		r := <-ch
		results[r.depth] = r.sum
	}

	for i := minDepth; i <= maxDepth; i += 2 {
		fmt.Printf("%d\t trees of depth %d\t check: %d\n", 1<<uint(maxDepth-i+minDepth), i, results[i])
	}

	// Long lived tree stats
	fmt.Printf("long lived tree of depth %d\t check: %d\n", maxDepth, long.Checksum())
}