/* **********************************************************************
 * binarytree: An Asymptote module to draw binary trees                 *
 *                                                                      *
 * Copyright(C) 2006                                                    *
 * Tobias Langner tobias[at]langner[dot]nightlabs[dot]de                *
 *                                                                      *
 * Modified by John Bowman                                              *
 *                                                                      *
 * Condensed mode:                                                      *
 * Copyright(C) 2012                                                    *
 * Gerasimos Dimitriadis dimeg [at] intracom [dot] gr                   *
 *                                                                      *
 ************************************************************************
 *                                                                      *
 * This library is free software; you can redistribute it and/or        *
 * modify it under the terms of the GNU Lesser General Public           *
 * License as published by the Free Software Foundation; either         *
 * version 3 of the License, or(at your option) any later version.      *
 *                                                                      *
 * This library is distributed in the hope that it will be useful,      *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of       *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU    *
 * Lesser General Public License for more details.                      *
 *                                                                      *
 * You should have received a copy of the GNU Lesser General Public     *
 * License along with this library; if not, write to the                *
 *     Free Software Foundation, Inc.,                                  *
 *     51 Franklin St, Fifth Floor,                                     *
 *     Boston, MA  02110-1301  USA                                      *
 *                                                                      *
 * Or get it online:                                                    *
 *     http: //www.gnu.org/copyleft/lesser.html                         *
 *                                                                      *
 ***********************************************************************/

// default values
real minDistDefault=0.2cm;
real nodeMarginDefault=0.1cm;

// structure to represent nodes in a binary tree
struct binarytreeNode {
  int key;
  binarytreeNode left;
  binarytreeNode right;
  binarytreeNode parent;
  bool spans_calculated=false;
  int left_span,total_left_span;
  int right_span,total_right_span;
  void update_spans();

  // Get the horizontal span of the tree consisting of the current
  // node plus the whole subtree that is rooted at the right child
  // (condensed mode)
  int getTotalRightSpan() {
    if(spans_calculated == false) {
      update_spans();
    }

    return total_right_span;
  }

  // Get the horizontal span of the tree consisting of the current
  // node plus the whole subtree that is rooted at the left child
  // (condensed mode)
  int getTotalLeftSpan() {
    if(spans_calculated == false) {
      update_spans();
    }
    return total_left_span;
  }

  // Get the horizontal distance between this node and its right child
  // (condensed mode)
  int getRightSpan() {
    if(spans_calculated == false) {
      update_spans();
    }
    return right_span;
  }

  // Get the horizontal distance between this node and its left child
  // (condensed mode)
  int getLeftSpan() {
    if(spans_calculated == false) {
      update_spans();
    }
    return left_span;
  }

  // Update all span figures for this node.
  // condensed mode)
  update_spans=new void() {
                           if(spans_calculated == true)
                             return;

                           left_span=0;
                           total_left_span=0;
                           right_span=0;
                           total_right_span=0;

                           if(left != null) {
                             left_span=left.getTotalRightSpan()+1;
                             total_left_span=left_span+left.getTotalLeftSpan();
                           }

                           if(right != null) {
                             right_span=right.getTotalLeftSpan()+1;
                             total_right_span=right_span+right.getTotalRightSpan();
                           }
                           spans_calculated=true;
  };

  // set the left child of this node
  void setLeft(binarytreeNode left) {
    this.left=left;
    this.left.parent=this;
  }

  // set the right child of this node
  void setRight(binarytreeNode right) {
    this.right=right;
    this.right.parent=this;
  }

  // return a boolean indicating whether this node is the root
  bool isRoot() {
    return parent == null;
  }

  // return the level of the subtree rooted at this node.
  int getLevel() {
    if(isRoot())
      return 1;
    else
      return parent.getLevel()+1;
  }

  // set the children of this binarytreeNode
  void setChildren(binarytreeNode left, binarytreeNode right) {
    setLeft(left);
    setRight(right);
  }

  // create a new binarytreeNode with key <key>
  static binarytreeNode binarytreeNode(int key) {
    binarytreeNode toReturn=new binarytreeNode;
    toReturn.key=key;
    return toReturn;
  }

  // returns the height of the subtree rooted at this node.
  int getHeight() {
    if(left == null && right == null)
      return 1;
    if(left == null)
      return right.getHeight()+1;
    if(right == null)
      return left.getHeight()+1;

    return max(left.getHeight(),right.getHeight())+1;
  }
}

binarytreeNode operator init() {return null;}

// "constructor" for binarytreeNode
binarytreeNode binarytreeNode(int key)=binarytreeNode.binarytreeNode;

// draw the tree rooted at the given <node> at the given position <pos>, with
// <height>=the height of the containing tree,
// <minDist>=the minimal horizontal distance of two nodes at the lowest level,
// <levelDist>=the vertical distance between two levels,
// <nodeDiameter>=the diameter of one node.
object draw(picture pic=currentpicture, binarytreeNode node, pair pos,
            int height, real minDist, real levelDist, real nodeDiameter,
            pen p=currentpen, bool condensed=false) {
  Label label=Label(math((string) node.key),pos);

  binarytreeNode left=node.left;
  binarytreeNode right=node.right;

  // return the distance for two nodes at the given <level> when the
  // containing tree has height <height>
  // and the minimal distance between two nodes is <minDist> .
  real getDistance(int level, int height, real minDist) {
    return(nodeDiameter+minDist)*2^(height-level);
  }

  // return the horiontal distance between node <n> and its left child
  // (condensed mode)
  real getLeftDistance(binarytreeNode n) {
    return(nodeDiameter+minDist) *(real)n.getLeftSpan() * 0.5;
  }

  // return the horiontal distance between node <n> and its right child
  // (condensed mode)
  real getRightDistance(binarytreeNode n) {
    return(nodeDiameter+minDist) *(real)n.getRightSpan() * 0.5;
  }

  real dist=getDistance(node.getLevel(),height,minDist)/2;

  // draw the connection between the two nodes at the given positions
  // by calculating the connection points and drawing the corresponding
  // arrow.
  void deferredDrawNodeConnection(pair parentPos, pair childPos) {
    pic.add(new void(frame f, transform t) {
        pair start,end;
        // calculate connection path
        transform T=shift(nodeDiameter/2*unit(t*childPos-t*parentPos));
        path arr=(T*t*parentPos)--(inverse(T)*t*childPos);
        draw(f,PenMargin(arr,p).g,p,Arrow(5));
      });
    pic.addPoint(parentPos);
    pic.addPoint(childPos);
  }

  if(left != null) {
    pair childPos;
    if(condensed == false) {
      childPos=pos-(0,levelDist)-(dist/2,0);
    }
    else {
      childPos=pos-(0,levelDist)-((real)getLeftDistance(node),0);
    }
    draw(pic,left,childPos,height,minDist,levelDist,nodeDiameter,p,condensed);
    deferredDrawNodeConnection(pos,childPos);
  }

  if(right != null) {
    pair childPos;
    if(condensed == false) {
      childPos=pos-(0,levelDist)+(dist/2,0);
    }
    else {
      childPos=pos-(0,levelDist)+((real)getRightDistance(node),0);
    }
    draw(pic,right,childPos,height,minDist,levelDist,nodeDiameter,p,condensed);
    deferredDrawNodeConnection(pos,childPos);
  }

  picture obj;
  draw(obj,circle((0,0),nodeDiameter/2),p);
  label(obj,label,(0,0),p);

  add(pic,obj,pos);

  return label;
}

struct key {
  int n;
  bool active;
}

key key(int n, bool active=true) {key k; k.n=n; k.active=active; return k;}

key operator cast(int n) {return key(n);}
int operator cast(key k) {return k.n;}
int[] operator cast(key[] k) {
  int[] I;
  for(int i=0; i < k.length; ++i)
    I[i]=k[i].n;
  return I;
}

key nil=key(0,false);

// structure to represent a binary tree.
struct binarytree {
  binarytreeNode root;
  int[] keys;

  // add the given <key> to the tree by searching for its place and
  // inserting it there.
  void addKey(int key) {
    binarytreeNode newNode=binarytreeNode(key);

    if(root == null) {
      root=newNode;
      keys.push(key);
      return;
    }

    binarytreeNode n=root;
    while(n != null) {
      if(key < n.key) {
        if(n.left != null)
          n=n.left;
        else {
          n.setLeft(newNode);
          keys.push(key);
          return;
        }
      } else if(key > n.key) {
        if(n.right != null)
          n=n.right;
        else {
          n.setRight(newNode);
          keys.push(key);
          return;
        }
      }
    }
  }

  // return the height of the tree
  int getHeight() {
    if(root == null)
      return 0;
    else
      return root.getHeight();
  }

  // add all given keys to the tree sequentially
  void addSearchKeys(int[] keys) {
    for(int i=0; i < keys.length; ++i) {
      int key=keys[i];
      // Ignore duplicate keys
      if(find(this.keys == key) == -1)
        addKey(key);
    }
  }

  binarytreeNode build(key[] keys, int[] ind) {
    if(ind[0] >= keys.length) return null;
    key k=keys[ind[0]];
    ++ind[0];
    if(!k.active) return null;
    binarytreeNode bt=binarytreeNode(k);
    binarytreeNode left=build(keys,ind);
    binarytreeNode right=build(keys,ind);
    bt.left=left; bt.right=right;
    if(left != null) left.parent=bt;
    if(right != null) right.parent=bt;
    return bt;
  }

  void addKeys(key[] keys) {
    int[] ind={0};
    root=build(keys,ind);
    this.keys=keys;
  }


  // return all key in the tree
  int[] getKeys() {
    return keys;
  }
}

binarytree searchtree(...int[] keys)
{
  binarytree bt;
  bt.addSearchKeys(keys);
  return bt;
}

binarytree binarytree(...key[] keys)
{
  binarytree bt;
  bt.addKeys(keys);
  return bt;
}

// draw the given binary tree.
void draw(picture pic=currentpicture, binarytree tree,
          real minDist=minDistDefault, real nodeMargin=nodeMarginDefault,
          pen p=currentpen, bool condensed=false)
{
  int[] keys=tree.getKeys();

  // calculate the node diameter so that all keys fit into it
  frame f;
  for(int i=0; i < keys.length; ++i)
    label(f,math(string(keys[i])),p);

  real nodeDiameter=abs(max(f)-min(f))+2*nodeMargin;
  real levelDist=nodeDiameter*1.8;

  draw(pic,tree.root,(0,0),tree.getHeight(),minDist,levelDist,nodeDiameter,p,
       condensed);
}