#!  /usr/bin/perl
#
# Copyright (c) 2004 Matthias Bauer <matthiasb@acm.org>
#
# Permission is hereby granted, free of charge, to any person obtaining a cop
#
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#

# Nodes in our tree
package node;

sub new {
    my $proto = shift;
    my $class = ref($proto) || $proto;
    my $self  = {};

    $self->{0}       = undef;    # left child
    $self->{1}       = undef;    # right child
    $self->{pos}     = [];       # position in img
    $self->{colors}  = [];       # colors of links to children
    $self->{level}   = undef;    # depth
    $self->{tag}     = undef;    # path to the node
    $self->{label}   = undef;    # optional text to print next to
    $self->{caption} = undef;    # dto. to print below

    return bless $self, $class;
}

1;

use GD;
use strict;
use warnings "all";
use Math::Trig;

# GD is sooo sick. Colors are allocated _per image_.
my @dim = ( 800, 640 );
my $b   = new GD::Image(@dim);

my $white = $b->colorAllocate( 255, 255, 255 );
my $black = $b->colorAllocate( 0,   0,   0 );

# For text in the image
my $font = "/usr/local/lib/X11/fonts/mscorefonts/arial.ttf";

# offset of rotated text below a node
my $offset = 3;

# Creates a tree of $depth in an image of ($width, $height)
sub mktree {
    my ( $width, $height, $depth ) = @_;
    my $root = new node;
    $root->{pos}   = [ $width / 2, 0 ];
    $root->{level} = 0;
    $root->{tag}   = "";

    # it's divine ...
    sub btr {
        my ( $n, $l ) = @_;
        return if ( $l >= $depth );
        $n->{0} = new node;
        $n->{1} = new node;
        $n->{colors} = [ $black, $black ];
        $n->{0}->{level} = $n->{1}->{level} = $l;
        $n->{0}->{tag}   = $n->{tag} . "0";
        $n->{1}->{tag}   = $n->{tag} . "1";
        $n->{0}->{pos}   = [
            $n->{pos}->[0] - ( $width * ( 0.5**( $l + 1 ) ) ),
            $n->{pos}->[1] + ( ( $height / $depth ) / $l )
        ];
        $n->{1}->{pos} = [
            $n->{pos}->[0] + ( $width * ( 0.5**( $l + 1 ) ) ),
            $n->{pos}->[1] + ( ( $height / $depth ) / $l )
        ];
        btr( $n->{0}, $l + 1 );
        btr( $n->{1}, $l + 1 );
        return;
    }
    btr( $root, 1 );
    return $root;
}

# Draw it
sub rekdraw {
    my ( $img, $n ) = @_;
    my @box;

    if ( defined $n->{label} ) {
        print STDERR "$n->{label}\n";
        @box = GD::Image->stringFT( $black, $font, 10, 0, @{ $n->{pos} },
            $n->{label} );
        unless (@box) {
            print STDERR "Damn: $@\n";
            exit 1;
        }
        $img->stringFT( $black, $font, 10, 0, @{ $n->{pos} }, $n->{label} );
    }

    if ( defined $n->{caption} ) {
        print STDERR "$n->{caption}\n";

        @box = GD::Image->stringFT( $black, $font, 10, 1.5 * pi, @{ $n->{pos} },
            $n->{caption} );

        unless (@box) {
            print STDERR "Damn: $@\n";
            exit 1;
        }

        my ( $x, $y ) = @{ $n->{pos} };
        $y += $offset;
        die "empty box" if ( $box[4] - $box[0] == 0 );

        # Center the caption
        $x -= ( $box[4] - $box[0] ) / 2;
        $img->stringFT( $black, $font, 10, 1.5 * pi, $x, $y, $n->{caption} );
    }

    return unless $n->{0};

    $img->line(
        $n->{pos}->[0],      $n->{pos}->[1], $n->{0}->{pos}->[0],
        $n->{0}->{pos}->[1], $n->{colors}->[0]
    );

    $img->line(
        $n->{pos}->[0],      $n->{pos}->[1], $n->{1}->{pos}->[0],
        $n->{1}->{pos}->[1], $n->{colors}->[1]
    );

    rekdraw( $img, $n->{0} );
    rekdraw( $img, $n->{1} );
    return;
}

# draws the edges along a given path of zeros and ones in
# the given color.
sub markpath {
    my ( $path, $root, $col ) = @_;
    my $n;
    $n = $root;
    foreach $b ( split //, $path ) {
        $n->{colors}->[$b] = $col;
        $n = $n->{$b};
    }
    return $root;
}

# puts the given labels (@names) to the nodes along the
# given path.
sub labelpath {
    my ( $path, $root, @names ) = @_;
    my $n;
    $n = $root;
    foreach $b ( split //, $path ) {
        $n->{label} = shift @names;
        $n = $n->{$b};
    }
    $n->{label} = shift @names;
    return $root;
}

# puts a caption to a node at the end of the given path.
# the caption is printed vertically below the node.
sub leafcaption {
    my ( $path, $root, $cap ) = @_;
    my $n;
    $n = $root;
    foreach $b ( split //, $path ) {
        $n = $n->{$b} if exists $n->{$b};
    }
    $n->{caption} = $cap;
}

# draws the edges to the children of a give node (by path)
# in the given color.
sub pathpairs {
    my ( $path, $root, $col ) = @_;
    my $n;
    $n = $root;
    foreach $b ( split //, $path ) {
        $n->{colors}->[0] = $col;
        $n->{colors}->[1] = $col;
        $n                = $n->{$b};
    }
    return $root;
}

# Almost Haskelly. It's possible to rewrite almost any
# function here by supplying the right $q to forall.
# iterator function over all nodes of a tree rooted in root
sub forall {
    my ( $root, $q ) = @_;

    sub ftrk {
        my $n = shift;
        &$q($n);
        unless ( defined $n->{0} ) {
            return;
        }
        ftrk( $n->{0}, $q );
        ftrk( $n->{1}, $q );
        return;
    }
    ftrk( $root, $q );
}

# draws all edges between nodes of layer $l and $l+1 in
# the given color.
sub colorlayer {
    my ( $l, $root, $col ) = @_;
    return if $l < 0;

    sub rk {
        my ( $n, $d ) = @_;
        if ( $d == $l ) {
            $n->{colors}->[0] = $col;
            $n->{colors}->[1] = $col;
            return;
        }
        else {
            rk( $n->{0}, $d + 1 );
            rk( $n->{1}, $d + 1 );
        }
        return;
    }
    rk( $root, 0 );
    return $root;
}

# XXX No idea why this recurses endlessly
sub copytree {
    my $src = shift;
    my $dst = new node;
    foreach my $att ( "label", "caption", "pos", "tag", "level" ) {
        $dst->{$att} = $src->{$att};
    }
    my @colors = @{ $src->{colors} };
    $dst->{colors} = \@colors;
    if ( defined $src->{0} ) {
        $dst->{0} = copytree $src->{0};
    }
    if ( defined $src->{1} ) {
        $dst->{1} = copytree $src->{1};
    }
    return $dst;
}

# returns an array of leaves
sub leaves {
    my $root = shift;
    my @l;

    sub lrk {
        my ( $n, $rl ) = @_;
        unless ( $n->{0} ) {
            push @$rl, $n;
            return;
        }
        else {
            lrk( $n->{0}, $rl );
            lrk( $n->{1}, $rl );
        }
        return;
    }
    lrk( $root, \@l );
    return @l;
}

# puts captions on leaves, consisting of $labelbase and
# a number.
sub doclabeltree {
    my ( $root, $labelbase ) = @_;
    my @leaves = leaves($root);
    my $c      = 0;
    foreach my $l (@leaves) {
        $l->{caption} = "$labelbase$c";
        $c++;
    }
    return $root;
}

# XXX Does not work because of GD's aforementioned sickness
sub outtree {
    my ( $root, $fname ) = @_;
    my $img = $b->clone;
    rekdraw( $img, $root );
    open F, ">$fname.png" or die "Could not write to $fname";
    print F $img->png;
    close F;
}

# Tests
my $red = $b->colorAllocate( 255, 0, 0 );
my $foo = mktree( @dim, 7 );
my $outfile = $ARGV[0] or die "please give a file to write to";
pathpairs( "0110110", $foo, $red );
labelpath( "0110110", $foo, ( "", "0", "1", "1", "0", "1", "1", "0" ) );
leafcaption( "011011", $foo, "final" );
colorlayer( 2, $foo, $red );
rekdraw( $b, $foo );
open F, ">$outfile" or die "Could not write to $outfile";
print F $b->png;
close F;