#!/usr/bin/env python
"""Pi digits example

Example shows arbitrary precision using mpmath with the
computation of the digits of pi.
"""

from sympy.mpmath import libmpf
from sympy.mpmath import functions as mpf_funs

import math
from time import clock
import sys

def display_fraction(digits, skip=0, colwidth=10, columns=5):
    """Pretty printer for first n digits of a fraction"""
    perline = colwidth * columns
    printed = 0
    for linecount in range((len(digits)-skip) // (colwidth * columns)):
        line = digits[skip+linecount*perline:skip+(linecount+1)*perline]
        for i in range(columns):
            print line[i*colwidth : (i+1)*colwidth],
        print ":", (linecount+1)*perline
        if (linecount+1) % 10 == 0:
            print
        printed += colwidth*columns
    rem = (len(digits)-skip) % (colwidth * columns)
    if rem:
        buf = digits[-rem:]
        s = ""
        for i in range(columns):
            s += buf[:colwidth].ljust(colwidth+1, " ")
            buf = buf[colwidth:]
        print s + ":", printed + colwidth*columns

def calculateit(func, base, n, tofile):
    """Writes first n base-digits of a mpmath function to file"""
    prec = 100
    intpart = libmpf.small_numeral(int(float(func(prec))), base)
    if intpart == 0:
        skip = 0
    else:
        skip = len(intpart)
    print "Step 1 of 2: calculating binary value..."
    prec = int(n*math.log(base,2))+10
    t = clock()
    a = func(prec)
    step1_time = clock() - t
    print "Step 2 of 2: converting to specified base..."
    t = clock()
    d = libmpf.bin_to_radix(a.man, -a.exp, base, n)
    d = libmpf.numeral(d, base, n)
    step2_time = clock() - t
    print "\nWriting output...\n"
    if tofile:
        out_ = sys.stdout
        sys.stdout = tofile
    print "%i base-%i digits of pi:\n" % (n, base)
    print intpart, ".\n"
    display_fraction(d, skip, colwidth=10, columns=5)
    if tofile:
        sys.stdout = out_
    print "\nFinished in %f seconds (%f calc, %f convert)" % \
        ((step1_time + step2_time), step1_time, step2_time)

def interactive():
    """Simple function to interact with user"""
    print "Compute digits of pi with SymPy\n"
    base = input("Which base? (2-36, 10 for decimal) \n> ")
    digits = input("How many digits? (enter a big number, say, 10000)\n> ")
    tofile = raw_input("Output to file? (enter a filename, or just press enter\nto print directly to the screen) \n> ")
    if tofile:
        tofile = open(tofile, "w")
    calculateit(mpf_funs.pi, base, digits, tofile)

def main():
    """A non-interactive runner"""
    base = 16
    digits = 500
    tofile = None
    calculateit(mpf_funs.pi, base, digits, tofile)

if __name__ == "__main__":
    interactive()
