import timeit
from astropy.io import ascii
import pandas
import numpy as np
from astropy.table import Table, Column
from cStringIO import StringIO
import matplotlib.pyplot as plt
%matplotlib inline

def make_table(size=10000, n_floats=10, n_ints=0, n_strs=0, float_format=None, str_val=None):
    if str_val is None:
        str_val = "abcde12345"
    cols = []
    for i in xrange(n_floats):
        dat = np.random.uniform(low=1, high=10, size=size)
        cols.append(Column(dat, name='f{}'.format(i)))
    for i in xrange(n_ints):
        dat = np.random.randint(low=-9999999, high=9999999, size=size)
        cols.append(Column(dat, name='i{}'.format(i)))
    for i in xrange(n_strs):
        dat = np.repeat(str_val, size)
        cols.append(Column(dat, name='s{}'.format(i)))
    t = Table(cols)

    if float_format is not None:
        for col in t.columns.values():
            if col.name.startswith('f'):
                col.format = float_format

    return t

t = make_table(5, float_format='%.4f')
print(t)

out = StringIO()
ascii.write(t, out, use_fast_writer=True)
print out.getvalue()

out = StringIO()
pandas_table = pandas.DataFrame(np.array(t))
pandas_table.to_csv(out, float_format='%.4f')
print out.getvalue()

def plot_case(n_floats=10, n_ints=0, n_strs=0, float_format=None, str_val=None, strip=True):
    global table, np_table, pandas_table, flt_format, strip_whitespace
    strip_whitespace=strip
    flt_format = float_format
    n_rows = (100, 200, 500, 1000, 2000, 5000, 10000, 20000)  # include 50000 for publish run
    numbers = (10, 10,   5,     2,    1,    1,     1,     1)
    repeats = (3,   3,    3,    3,    3,    3,     3,     2)
    times_slow = []
    times_fast = []
    times_pandas = []
    for n_row, number, repeat in zip(n_rows, numbers, repeats):
        table = make_table(n_row, n_floats, n_ints, n_strs, float_format)
        np_table = np.array(table)
        pandas_table = pandas.DataFrame(np_table)
        t = timeit.repeat("out = StringIO(); ascii.write(table, out, use_fast_writer=False, strip_whitespace=strip_whitespace)", 
                   setup='from __main__ import ascii, table, StringIO, strip_whitespace', number=number, repeat=repeat)
        times_slow.append(min(t) / number)
        t = timeit.repeat("out = StringIO(); ascii.write(table, out, use_fast_writer=True, strip_whitespace=strip_whitespace)", 
                   setup='from __main__ import ascii, table, StringIO, strip_whitespace', number=number, repeat=repeat)
        times_fast.append(min(t) / number)
        t = timeit.repeat("out = StringIO(); pandas_table.to_csv(out, float_format=flt_format)", 
                   setup='from __main__ import pandas_table, pandas, StringIO, flt_format', number=number, repeat=repeat)
        times_pandas.append(min(t) / number)
    plt.loglog(n_rows, times_slow, '-ob', label='io.ascii Python')
    plt.loglog(n_rows, times_fast, '-or', label='io.ascii Fast-c')
    plt.loglog(n_rows, times_pandas, '-oc', label='Pandas')
    plt.grid()
    plt.legend(loc='best')
    plt.title('n_floats={} n_ints={} n_strs={} float_format={}'.format(n_floats, n_ints, n_strs, float_format))
    plt.xlabel('Number of rows')
    plt.ylabel('Time (sec)')
    print('Fast-C to Python speed ratio: {:.2f} : 1'.format(times_slow[-1] / times_fast[-1]))
    print('Pandas to Fast-C speed ratio: {:.2f} : 1'.format(times_fast[-1] / times_pandas[-1]))

plot_case(n_floats=10, n_ints=0, n_strs=0, float_format=None)

plot_case(n_floats=10, n_ints=10, n_strs=10, float_format=None)

plot_case(n_floats=10, n_ints=10, n_strs=10, float_format='%.4f')

plot_case(n_floats=10, n_ints=0, n_strs=0, float_format='%.4f')

plot_case(n_floats=0, n_ints=0, n_strs=10)

plot_case(n_floats=0, n_ints=0, n_strs=10, str_val="'asdf asdfa'")

plot_case(n_floats=0, n_ints=0, n_strs=10, strip=False)

plot_case(n_floats=0, n_ints=10, n_strs=0)