import numpy as np


def nicenum(x, round=True):
    ''' find a nice number approximately equal to x, where nice means
        a number divisible by 2 or 10.
        x : input floating point number
        round : True = round number, False = take ceiling
    '''
    exponent = np.floor(np.log10(x))
    frac = x/10**exponent
    if round :
        if frac < 1.5:
            ans = 1.
        elif frac < 3.:
            ans = 2.
        elif frac < 7. :
            ans = 5.
        else :
            ans = 10.
    else :
        if frac < 1.:
            ans = 1.
        elif frac <= 2.:
            ans = 2.
        elif frac <= 5. :
            ans = 5.
        else :
            ans = 10.
    return ans*10**exponent

def pretty(data, ntics=5, tight=True):
    '''  Come up with a nice range of values for a graph where the start and end values are
         divisible by 2 or 5, and we can have ntics tics which are "nice" numbers.
         data : input array of float or int
         ntics : number of tics to create
         tight : by default put a buffer of potentially zero space around data.
                 If False, then put a buffer of about 1/2 step past min and max data.
         return array of tic values and number of decimals
    '''
    data_range = data.max() - data.min()
    nice_range = nicenum(data_range, round=False)
    step = nicenum(nice_range/(ntics-1))
    graph_min = np.floor(data.min()/step)*step
    graph_max = np.ceil(data.max()/step)*step
    digits = np.floor(np.log10(step))
    num_frac = max(-digits, 0)
    if tight == False :
        return np.arange(graph_min, graph_max+step, step), int(digits)
    else :
        gmin = graph_min
        if abs(graph_min - data.min()) > 0.75*step:  
            gmin = graph_min + step/4. # pull in boundary a bit
        elif abs(graph_min - data.min()) < 0.25*step:  
            gmin = graph_min - step/4. # extend boundary a bit
        gmax = graph_max
        if abs(graph_max - data.max()) > 0.75*step:  
            gmax = graph_max - step/4. # pull in boundary a bit
        elif abs(graph_max - data.max()) < 0.25*step:  
            gmax = graph_max + step/4. # extend boundary a bit
        return np.array([gmin] + list(np.arange(graph_min+step, graph_max, step)) + [gmax]), int(digits)

if __name__ == '__main__' :

    data = np.array([4.8, 6.5, 12., 14.9, 10., 10.7, 7.9, 21.9, 12.5, 14.5, 9.2])

    print "---------- test 1 ---------------"
    print pretty(data)
    print "---------- test 2 ---------------"
    print pretty(data*100000)
    print "---------- test 3 ---------------"
    print pretty(data/100000)
    print "---------- test 3 ---------------"
    print pretty(data, tight=False)