#!/usr/bin/env python3
#
# collatz2dot - generate GraphViz input for a Collatz spiral
#
# $Id$

import sys
import os
from argparse import ArgumentParser
from math import sin, cos, pi, log, copysign

argparser = ArgumentParser(
  description = 'Generate GraphViz input for a Collatz graph; use -X for details'
)

argparser.add_argument(
  '-X', '--explain',
  action = 'store_true',
  default = False,
  help = 'explain how to use this command'
)

argparser.add_argument(
  '-n', '--maximum',
  type = int,
  default = 16,
  help = 'the maximum integer to include into the output graph'
)

argparser.add_argument(
  '-t', '--torque',
  type = float,
  default = 0,
  help = 'whether to use a spiral layout, and if so, how quickly; try 2 and 3'
)

argparser.add_argument(
  '-I', '--increment-color',
  type = str,
  default = '',
  help = "the color of increment arcs (n -> n+1); use 'invis' for invisible, '' to omit"
)

argparser.add_argument(
  '-H', '--halving-color',
  type = str,
  default = '',
  help = "the color of halving arcs (n -> n/2); use 'invis' for invisible, '' to omit"
)

argparser.add_argument(
  '-S', '--sesquipling-color',
  type = str,
  default = '',
  help = "the color of sexquipling arcs (n -> (3n+1)/2); use 'invis' for invisible, '' to omit"
)

argparser.add_argument(
  '-T', '--tripling-color',
  type = str,
  default = '',
  help = "the color of halving arcs (n -> n/2); use 'invis' for invisible, '' to omit"
)

args = argparser.parse_args()

def warn(*args, **kwargs):
  print(*args, file=sys.stderr, **kwargs)

def die(*args):
  warn('fatal error:', *args)
  exit(1)

if args.explain:
  warn('''
Generate GraphViz input for a Collatz graph.

Optionally, the following types of arcs can be included:

  - increment arcs (n -> n+1)
  - halving arcs (n -> n/2 for even n)
  - sesquipling arcs (n -> (3n+1)/2 for odd n)
  - tripling arcs (n -> 3n + 1 for odd n)

By default, no arcs arc included.
The vertices are colored by degree,
counting only sesquipling and halving arcs.

The -t option positions the vertices in a spiral,
with the argument a measure of torque (spiraling speed):
-t x puts the powers of x on a straight line.
Use neato(1) to render the result.

Without -t, no positions are supplied,
leaving the layout up to the GraphViz program invoked.

Examples:

  collatz2dot -I gray -H blue -T green | dot -Tpdf

    generates a 16-node graph with 3 types of arcs
    and leaves its layout to dot(1)

  collatz2dot -t 3 -I gray -T green -n 1000 | neato -Tpdf

    generates a 1000-node graph with increment arcs
    and tripling arcs in a spiral layout rendered by neato(1)

  collatz2dot -t 2 -I gray -H blue -n 1000 | neato -Tpdf

    generates a 1000-node graph with increment arcs
    and halving arcs that all point to the center

  collatz2dot -t 1.5 -S darkred -n 1000 | neato -Tpdf

    generates a 1000-node graph with only sesquipling arcs,
    curving slightly outwards

'''
  )
  exit(0)

n = args.maximum
torque = args.torque
i_color = args.increment_color
h_color = args.halving_color
s_color = args.sesquipling_color
t_color = args.tripling_color

if abs(torque) == 1:
  die('torque %d means no spiraling; please pick a different value' % (torque))

# generate a graph representing the Collatz function on the first n integers

print('''strict digraph {

  center = true;
  node [ shape = none ];
''')

in_degree = {}
out_degree = {}

def in_d(i):
  return in_degree.get(i, 0)

def out_d(i):
  return out_degree.get(i, 0)

def count(src, dst):
  out_degree[src] = out_d(src) + 1
  in_degree[dst] = in_d(dst) + 1

def degree2node_color(i):
  if in_d(i) > 0: 
    if out_d(i) > 0:
      return 'black'
    else:
      return 'red'
  elif out_d(i) > 0:
    return 'purple'
  else:
    return 'orange'

def print_positioned_node(i):
  alpha_i = 2 * pi * copysign(log(i, abs(torque)), torque)
  print('  %d [ fontcolor = %s, pos = "%f,%f!" ]' % (i, degree2node_color(i), alpha_i * sin(alpha_i), alpha_i * cos(alpha_i)))

def print_node(i):
  print('  %d [ fontcolor = %s ]' % (i, degree2node_color(i)))

def arc_color_attr(c):
  if c == 'invis':
    return 'style = invis'
  else:
    return 'color = %s' % (c)

def print_increment_arc_if_any(i):
  if i_color != '':
    print('  %d -> %d [ %s ]' % (i - 1, i, arc_color_attr(i_color)))

def print_halving_arc_if_any(i):
  if i % 2 == 0:
    count(i, i/2)
    if h_color != '':
      print('  %d -> %d [ %s ]' % (i, i / 2, arc_color_attr(h_color)))

def print_sesquipling_arc_if_any(i):
  if i % 3 == 2:
    count((2*i-1) / 3, i)
    if s_color != '':
      print('  %d -> %d [ %s ]' % ((2*i-1) / 3, i, arc_color_attr(s_color)))

def print_tripling_arc_if_any(i):
  if i % 6 == 4:
    #count((i-1) / 3, i)  # already counted when sesquipling
    if t_color != '':
      print('  %d -> %d [ %s ]' % ((i-1) / 3, i, arc_color_attr(t_color)))

if n == 1:
  print_node(1)
elif n > 1:
  for i in range(2, n+1):
    print_increment_arc_if_any(i)
    print_halving_arc_if_any(i)
    print_sesquipling_arc_if_any(i)
    print_tripling_arc_if_any(i)
  for i in range(1, n+1):
    if torque != 0:
      print_positioned_node(i)
    else:
      print_node(i)

print()
print('}')